ViT 图像分类训练实战自定义分类数据集与模型导出这篇教程根据我复现 Vision Transformer 图像分类流程时整理重点演示分类数据准备、模型定义、训练循环、单图验证和模型保存。本文整理自我的学习和项目复现过程尽量按实操顺序保留 notebook 的关键步骤同时把数据集获取方式调整为适合中文教程发布的写法。本文会重点跑通以下流程安装 Transformers 依赖从数据集后台获取分类数据使用 ImageFolder 加载数据定义 ViT 分类模型并训练保存并重新加载训练好的模型如果你正在系统学习目标检测、实例分割、OCR、多目标跟踪或视觉大模型建议收藏本文配套 notebook、示例图片和运行环境说明后续会继续整理。如果环境配置卡住可以在评论区说明具体报错。 文章目录ViT 图像分类训练实战自定义分类数据集与模型导出⚙️ 安装依赖 从数据集后台获取分类数据 加载分类数据集 定义 ViT 分类模型 设置训练参数️ 训练模型 单图验证 保存与导出模型 加载模型推理 小结 同系列教程汇总⚙️ 安装依赖先安装 Hugging Face Transformers后续会基于预训练 ViT 做分类微调。!pip install-q githttps://github.com/huggingface/transformers 从数据集后台获取分类数据将分类数据导出为文件夹结构后放到 Colab 对应目录。# 从数据集后台下载文件夹分类格式数据集并解压到 /content。DATASET_DIR/content# 目录下应包含 train、valid 等子目录TRAIN_DIR/content/trainVALID_DIR/content/valid 加载分类数据集使用 torchvision 的 ImageFolder 读取 train 和 valid 目录。importtorchvisionfromtorchvision.transformsimportToTensor train_dstorchvision.datasets.ImageFolder(/content/train/,transformToTensor())valid_dstorchvision.datasets.ImageFolder(/content/valid/,transformToTensor())test_dstorchvision.datasets.ImageFolder(/content/test/,transformToTensor()) 定义 ViT 分类模型在预训练 ViT 主干后接一个线性分类头用于适配自定义类别数。fromtransformersimportViTModelfromtransformers.modeling_outputsimportSequenceClassifierOutputimporttorch.nnasnnimporttorch.nn.functionalasFclassViTForImageClassification(nn.Module):def__init__(self,num_labels3):super(ViTForImageClassification,self).__init__()self.vitViTModel.from_pretrained(google/vit-base-patch16-224-in21k)self.dropoutnn.Dropout(0.1)self.classifiernn.Linear(self.vit.config.hidden_size,num_labels)self.num_labelsnum_labelsdefforward(self,pixel_values,labels):outputsself.vit(pixel_valuespixel_values)outputself.dropout(outputs.last_hidden_state[:,0])logitsself.classifier(output)lossNoneiflabelsisnotNone:loss_fctnn.CrossEntropyLoss()lossloss_fct(logits.view(-1,self.num_labels),labels.view(-1))iflossisnotNone:returnlogits,loss.item()else:returnlogits,None 设置训练参数配置 epoch、batch size、学习率、特征提取器、优化器和损失函数。EPOCHS3BATCH_SIZE10LEARNING_RATE2e-5fromtransformersimportViTFeatureExtractorimporttorch.nnasnnimporttorch# Define ModelmodelViTForImageClassification(len(train_ds.classes))# Feature Extractorfeature_extractorViTFeatureExtractor.from_pretrained(google/vit-base-patch16-224-in21k)# Adam Optimizeroptimizertorch.optim.Adam(model.parameters(),lrLEARNING_RATE)# Cross Entropy Lossloss_funcnn.CrossEntropyLoss()# Use GPU if availabledevicetorch.device(cudaiftorch.cuda.is_available()elsecpu)iftorch.cuda.is_available():model.cuda()️ 训练模型使用常规 PyTorch 训练循环完成 ViT 分类模型微调。importtorch.utils.dataasdatafromtorch.autogradimportVariableimportnumpyasnpprint(Number of train samples: ,len(train_ds))print(Number of test samples: ,len(test_ds))print(Detected Classes are: ,train_ds.class_to_idx)train_loaderdata.DataLoader(train_ds,batch_sizeBATCH_SIZE,shuffleTrue,num_workers4)test_loaderdata.DataLoader(test_ds,batch_sizeBATCH_SIZE,shuffleTrue,num_workers4)# Train the modelforepochinrange(EPOCHS):forstep,(x,y)inenumerate(train_loader):# Change input array into list with each batch being one elementxnp.split(np.squeeze(np.array(x)),BATCH_SIZE)# Remove unecessary dimensionforindex,arrayinenumerate(x):x[index]np.squeeze(array)# Apply feature extractor, stack back into 1 tensor and then convert to tensorxtorch.tensor(np.stack(feature_extractor(x)[pixel_values],axis0))# Send to GPU if availablex,yx.to(device),y.to(device)b_xVariable(x)# batch x (image)b_yVariable(y)# batch y (target)# Feed through modeloutput,lossmodel(b_x,None)# Calculate lossiflossisNone:lossloss_func(output,b_y)optimizer.zero_grad()loss.backward()optimizer.step()ifstep%500:# Get the next batch for testing purposestestnext(iter(test_loader))test_xtest[0]# Reshape and get feature matrices as neededtest_xnp.split(np.squeeze(np.array(test_x)),BATCH_SIZE)forindex,arrayinenumerate(test_x):test_x[index]np.squeeze(array)test_xtorch.tensor(np.stack(feature_extractor(test_x)[pixel_values],axis0))# Send to appropirate computing devicetest_xtest_x.to(device)test_ytest[1].to(device)# Get output ( respective class) and compare to targettest_output,lossmodel(test_x,test_y)test_outputtest_output.argmax(1)# Calculate Accuracyaccuracy(test_outputtest_y).sum().item()/BATCH_SIZEprint(Epoch: ,epoch,| train loss: %.4f%loss,| test accuracy: %.2f%accuracy) 单图验证从验证集中取一张图查看模型预测类别。importmatplotlib.pyplotaspltimportnumpyasnp EVAL_BATCH1eval_loaderdata.DataLoader(valid_ds,batch_sizeEVAL_BATCH,shuffleTrue,num_workers4)# Disable gradwithtorch.no_grad():inputs,targetnext(iter(eval_loader))# Reshape and get feature matrices as neededprint(inputs.shape)inputsinputs[0].permute(1,2,0)# Save original InputoriginalInputinputsforindex,arrayinenumerate(inputs):inputs[index]np.squeeze(array)inputstorch.tensor(np.stack(feature_extractor(inputs)[pixel_values],axis0))# Send to appropriate computing deviceinputsinputs.to(device)targettarget.to(device)# Generate predictionprediction,lossmodel(inputs,target)# Predicted class value using argmaxpredicted_classnp.argmax(prediction.cpu())value_predictedlist(valid_ds.class_to_idx.keys())[list(valid_ds.class_to_idx.values()).index(predicted_class)]value_targetlist(valid_ds.class_to_idx.keys())[list(valid_ds.class_to_idx.values()).index(target)]# Show resultplt.imshow(originalInput)plt.xlim(224,0)plt.ylim(224,0)plt.title(fPrediction:{value_predicted}- Actual target:{value_target})plt.show() 保存与导出模型训练完成后保存模型文件便于后续推理或迁移到其他环境。torch.save(model,/content/model.pt)fromgoogle.colabimportdrive drive.mount(/content/gdrive)%cp/content/model.pt/content/gdrive/My\ Drive 加载模型推理重新加载本地模型并切换到 eval 模式。MODEL_PATH/content/model.ptmodeltorch.load(MODEL_PATH)model.eval() 小结这篇教程完整整理了ViT 图像分类训练的核心复现流程。实际操作时建议先确认 GPU、依赖版本、数据集路径和模型权重路径再逐段运行 notebook。后续我会继续按源项目顺序整理同系列中的目标检测、实例分割、OCR、多目标跟踪和视觉大模型教程。 同系列教程汇总Google Gemini 3.5 Flash 零样本目标检测教程从提示词到可视化结果GLM-OCR 文档识别实战教程从验证码、公式到车牌 OCRRF-DETR ByteTrack 多目标跟踪实战教程从命令行到 Python 视频轨迹可视化SAM 3 图像分割实战教程文本、框和点提示的多种分割方式-ViT 图像分类训练实战自定义分类数据集与模型导出-本文
ViT 图像分类训练实战：自定义分类数据集与模型导出
发布时间：2026/7/21 23:10:15
ViT 图像分类训练实战自定义分类数据集与模型导出这篇教程根据我复现 Vision Transformer 图像分类流程时整理重点演示分类数据准备、模型定义、训练循环、单图验证和模型保存。本文整理自我的学习和项目复现过程尽量按实操顺序保留 notebook 的关键步骤同时把数据集获取方式调整为适合中文教程发布的写法。本文会重点跑通以下流程安装 Transformers 依赖从数据集后台获取分类数据使用 ImageFolder 加载数据定义 ViT 分类模型并训练保存并重新加载训练好的模型如果你正在系统学习目标检测、实例分割、OCR、多目标跟踪或视觉大模型建议收藏本文配套 notebook、示例图片和运行环境说明后续会继续整理。如果环境配置卡住可以在评论区说明具体报错。 文章目录ViT 图像分类训练实战自定义分类数据集与模型导出⚙️ 安装依赖 从数据集后台获取分类数据 加载分类数据集 定义 ViT 分类模型 设置训练参数️ 训练模型 单图验证 保存与导出模型 加载模型推理 小结 同系列教程汇总⚙️ 安装依赖先安装 Hugging Face Transformers后续会基于预训练 ViT 做分类微调。!pip install-q githttps://github.com/huggingface/transformers 从数据集后台获取分类数据将分类数据导出为文件夹结构后放到 Colab 对应目录。# 从数据集后台下载文件夹分类格式数据集并解压到 /content。DATASET_DIR/content# 目录下应包含 train、valid 等子目录TRAIN_DIR/content/trainVALID_DIR/content/valid 加载分类数据集使用 torchvision 的 ImageFolder 读取 train 和 valid 目录。importtorchvisionfromtorchvision.transformsimportToTensor train_dstorchvision.datasets.ImageFolder(/content/train/,transformToTensor())valid_dstorchvision.datasets.ImageFolder(/content/valid/,transformToTensor())test_dstorchvision.datasets.ImageFolder(/content/test/,transformToTensor()) 定义 ViT 分类模型在预训练 ViT 主干后接一个线性分类头用于适配自定义类别数。fromtransformersimportViTModelfromtransformers.modeling_outputsimportSequenceClassifierOutputimporttorch.nnasnnimporttorch.nn.functionalasFclassViTForImageClassification(nn.Module):def__init__(self,num_labels3):super(ViTForImageClassification,self).__init__()self.vitViTModel.from_pretrained(google/vit-base-patch16-224-in21k)self.dropoutnn.Dropout(0.1)self.classifiernn.Linear(self.vit.config.hidden_size,num_labels)self.num_labelsnum_labelsdefforward(self,pixel_values,labels):outputsself.vit(pixel_valuespixel_values)outputself.dropout(outputs.last_hidden_state[:,0])logitsself.classifier(output)lossNoneiflabelsisnotNone:loss_fctnn.CrossEntropyLoss()lossloss_fct(logits.view(-1,self.num_labels),labels.view(-1))iflossisnotNone:returnlogits,loss.item()else:returnlogits,None 设置训练参数配置 epoch、batch size、学习率、特征提取器、优化器和损失函数。EPOCHS3BATCH_SIZE10LEARNING_RATE2e-5fromtransformersimportViTFeatureExtractorimporttorch.nnasnnimporttorch# Define ModelmodelViTForImageClassification(len(train_ds.classes))# Feature Extractorfeature_extractorViTFeatureExtractor.from_pretrained(google/vit-base-patch16-224-in21k)# Adam Optimizeroptimizertorch.optim.Adam(model.parameters(),lrLEARNING_RATE)# Cross Entropy Lossloss_funcnn.CrossEntropyLoss()# Use GPU if availabledevicetorch.device(cudaiftorch.cuda.is_available()elsecpu)iftorch.cuda.is_available():model.cuda()️ 训练模型使用常规 PyTorch 训练循环完成 ViT 分类模型微调。importtorch.utils.dataasdatafromtorch.autogradimportVariableimportnumpyasnpprint(Number of train samples: ,len(train_ds))print(Number of test samples: ,len(test_ds))print(Detected Classes are: ,train_ds.class_to_idx)train_loaderdata.DataLoader(train_ds,batch_sizeBATCH_SIZE,shuffleTrue,num_workers4)test_loaderdata.DataLoader(test_ds,batch_sizeBATCH_SIZE,shuffleTrue,num_workers4)# Train the modelforepochinrange(EPOCHS):forstep,(x,y)inenumerate(train_loader):# Change input array into list with each batch being one elementxnp.split(np.squeeze(np.array(x)),BATCH_SIZE)# Remove unecessary dimensionforindex,arrayinenumerate(x):x[index]np.squeeze(array)# Apply feature extractor, stack back into 1 tensor and then convert to tensorxtorch.tensor(np.stack(feature_extractor(x)[pixel_values],axis0))# Send to GPU if availablex,yx.to(device),y.to(device)b_xVariable(x)# batch x (image)b_yVariable(y)# batch y (target)# Feed through modeloutput,lossmodel(b_x,None)# Calculate lossiflossisNone:lossloss_func(output,b_y)optimizer.zero_grad()loss.backward()optimizer.step()ifstep%500:# Get the next batch for testing purposestestnext(iter(test_loader))test_xtest[0]# Reshape and get feature matrices as neededtest_xnp.split(np.squeeze(np.array(test_x)),BATCH_SIZE)forindex,arrayinenumerate(test_x):test_x[index]np.squeeze(array)test_xtorch.tensor(np.stack(feature_extractor(test_x)[pixel_values],axis0))# Send to appropirate computing devicetest_xtest_x.to(device)test_ytest[1].to(device)# Get output ( respective class) and compare to targettest_output,lossmodel(test_x,test_y)test_outputtest_output.argmax(1)# Calculate Accuracyaccuracy(test_outputtest_y).sum().item()/BATCH_SIZEprint(Epoch: ,epoch,| train loss: %.4f%loss,| test accuracy: %.2f%accuracy) 单图验证从验证集中取一张图查看模型预测类别。importmatplotlib.pyplotaspltimportnumpyasnp EVAL_BATCH1eval_loaderdata.DataLoader(valid_ds,batch_sizeEVAL_BATCH,shuffleTrue,num_workers4)# Disable gradwithtorch.no_grad():inputs,targetnext(iter(eval_loader))# Reshape and get feature matrices as neededprint(inputs.shape)inputsinputs[0].permute(1,2,0)# Save original InputoriginalInputinputsforindex,arrayinenumerate(inputs):inputs[index]np.squeeze(array)inputstorch.tensor(np.stack(feature_extractor(inputs)[pixel_values],axis0))# Send to appropriate computing deviceinputsinputs.to(device)targettarget.to(device)# Generate predictionprediction,lossmodel(inputs,target)# Predicted class value using argmaxpredicted_classnp.argmax(prediction.cpu())value_predictedlist(valid_ds.class_to_idx.keys())[list(valid_ds.class_to_idx.values()).index(predicted_class)]value_targetlist(valid_ds.class_to_idx.keys())[list(valid_ds.class_to_idx.values()).index(target)]# Show resultplt.imshow(originalInput)plt.xlim(224,0)plt.ylim(224,0)plt.title(fPrediction:{value_predicted}- Actual target:{value_target})plt.show() 保存与导出模型训练完成后保存模型文件便于后续推理或迁移到其他环境。torch.save(model,/content/model.pt)fromgoogle.colabimportdrive drive.mount(/content/gdrive)%cp/content/model.pt/content/gdrive/My\ Drive 加载模型推理重新加载本地模型并切换到 eval 模式。MODEL_PATH/content/model.ptmodeltorch.load(MODEL_PATH)model.eval() 小结这篇教程完整整理了ViT 图像分类训练的核心复现流程。实际操作时建议先确认 GPU、依赖版本、数据集路径和模型权重路径再逐段运行 notebook。后续我会继续按源项目顺序整理同系列中的目标检测、实例分割、OCR、多目标跟踪和视觉大模型教程。 同系列教程汇总Google Gemini 3.5 Flash 零样本目标检测教程从提示词到可视化结果GLM-OCR 文档识别实战教程从验证码、公式到车牌 OCRRF-DETR ByteTrack 多目标跟踪实战教程从命令行到 Python 视频轨迹可视化SAM 3 图像分割实战教程文本、框和点提示的多种分割方式-ViT 图像分类训练实战自定义分类数据集与模型导出-本文
