PyTorch 示例框架如何使用分类模型ResNet-152/ EfficientNet 深度学习架构 使用EfficientNet模型对262种水果进行分类 水果分类数据集的训练及应用
文章目录
- 1. 模型选择
- 使用以下几种主流分类模型:
- ✅ **EfficientNet(B4-B7)**
- ✅ **ResNet-152 / ResNeXt-101**
- ✅ **Vision Transformer (ViT) 或 ConvNeXt**
- ✅ **使用预训练模型 + 微调(Fine-tuning)**
- 步骤 1: 安装必要的库
- 步骤 2: 数据加载与预处理
- 步骤 3: 模型构建
- 步骤 4: 训练模型
- 数据增强(Data Augmentation)
- 3. 模型构建(PyTorch 示例)
- 4. 训练流程设计
- 超参数设置:
- 训练代码框架(PyTorch)
- 5. 性能评估与指标
- 6. 模型导出与部署
- 导出 ONNX 模型:
- 部署选项:
- 7. 进阶技巧
- 📌 标签平滑(Label Smoothing)
- 📌 半监督学习(伪标签)
- 📌 使用 Mixup / CutMix 增强训练
- 8. 推荐训练硬件配置
- 总结
- 1. 安装必要的库
- 2. 数据准备
- 3. 数据加载器
- 4. 模型定义
- 5. 训练过程
- 6. 模型评估
以下文字及代码仅供参考学习。
262类水果分类数据集
包含 262 种不同水果的 225,640 张图像的数据集
涵盖常见与稀有水果,像苹果、香蕉等常见水果,以及阿比乌、巴西莓、针叶樱桃等稀有水果均有涉及。这些水果来自世界各地,反映了水果的多样性,能满足研究不同产地、特性水果的需求
包含 225,640 张图像,数量较多。且每个水果类别图像数量分布相对均匀,平均值为 861,中位数为 1007,标准差为 276,为训练水果分类模型提供了充足的数据,有助于模型学习不同水果的特征
具体类别如下:
abiu, acai, acerola, ackee, alligator apple, ambarella, apple, apricot, araza, avocado, bael, banana, barbadine, barberry, bayberry, beach plum, bearberry, bell pepper, betel nut, bignay, bilimbi, bitter gourd, black berry, black cherry, black currant, black mullberry, black sapote, blueberry, bolwarra, bottle gourd, brazil nut, bread fruit, buddha s hand, buffaloberry, burdekin plum, burmese grape, caimito, camu camu, canistel, cantaloupe, cape gooseberry, carambola, cardon, cashew, cedar bay cherry, cempedak, ceylon gooseberry, che, chenet, cherimoya, cherry, chico, chokeberry, clementine, cloudberry, cluster fig, cocoa bean, coconut, coffee, common buckthorn, corn kernel, cornelian cherry, crab apple, cranberry, crowberry, cupuacu, custard apple, damson, date, desert fig, desert lime, dewberry, dragonfruit, durian, eggplant, elderberry, elephant apple, emblic, entawak, etrog, feijoa, fibrous satinash, fig, finger lime, galia melon, gandaria, genipap, goji, gooseberry, goumi, grape, grapefruit, greengage, grenadilla, guanabana, guarana, guava, guavaberry, hackberry, hard kiwi, hawthorn, hog plum, honeyberry, honeysuckle, horned melon, illawarra plum, indian almond, indian strawberry, ita palm, jaboticaba, jackfruit, jalapeno, jamaica cherry, jambul, japanese raisin, jasmine, jatoba, jocote, jostaberry, jujube, juniper berry, kaffir lime, kahikatea, kakadu plum, keppel, kiwi, kumquat, kundong, kutjera, lablab, langsat, lapsi, lemon, lemon aspen, leucaena, lillipilli, lime, lingonberry, loganberry, longan, loquat, lucuma, lulo, lychee, mabolo, macadamia, malay apple, mamey apple, mandarine, mango, mangosteen, manila tamarind, marang, mayhaw, maypop, medlar, melinjo, melon pear, midyim, miracle fruit, mock strawberry, monkfruit, monstera deliciosa, morinda, mountain papaya, mountain soursop, mundu, muskmelon, myrtle, nance, nannyberry, naranjilla, native cherry, native gooseberry, nectarine, neem, nungu, nutmeg, oil palm, old world sycomore, olive, orange, oregon grape, otaheite apple, papaya, passion fruit, pawpaw, pea, peanut, pear, pequi, persimmon, pigeon plum, pigface, pili nut, pineapple, pineberry, pitomba, plumcot, podocarpus, pomegranate, pomelo, prikly pear, pulasan, pumpkin, pupunha, purple apple berry, quandong, quince, rambutan, rangpur, raspberry, red mulberry, redcurrant, riberry, ridged gourd, rimu, rose hip, rose myrtle, rose-leaf bramble, saguaro, salak, salal, salmonberry, sandpaper fig, santol, sapodilla, saskatoon, sea buckthorn, sea grape, snowberry, soncoya, strawberry, strawberry guava, sugar apple, surinam cherry, sycamore fig, tamarillo, tangelo, tanjong, taxus baccata, tayberry, texas persimmon, thimbleberry, tomato, toyon, ugli fruit, vanilla, velvet tamarind, watermelon, wax gourd, white aspen, white currant, white mulberry, white sapote, wineberry, wongi, yali pear, yellow plum, yuzu, zigzag vine, zucchini
非常庞大且具有挑战性的分类任务。为了有效训练这样的模型,需要选择适合大规模多类别分类的深度学习架构,并结合良好的数据预处理和优化策略。
以下是针对该水果分类任务的完整解决方案,包括模型选择、训练流程、代码实现以及性能优化建议。
仅供参考学习使用。
1. 模型选择
使用以下几种主流分类模型:
✅EfficientNet(B4-B7)
- 优点:参数量适中,精度高,适用于中大规模分类任务。
- 可扩展性强,支持迁移学习。
✅ResNet-152 / ResNeXt-101
- 优点:经典的残差结构,泛化能力强。
- 对于大类别数表现稳定。
✅Vision Transformer (ViT) 或 ConvNeXt
- 更现代的架构,在超多类别分类任务中表现出色。
- ViT 在 ImageNet-21K 上预训练后迁移到你的 262 类效果更佳。
✅使用预训练模型 + 微调(Fine-tuning)
- 使用在 ImageNet 上预训练的模型作为 backbone。
- 替换最后的全连接层为输出 262 个类别的分类器。
步骤 1: 安装必要的库
确保安装了所有需要的Python库:
pipinstalltorch torchvision timm scikit-learn matplotlib tqdm步骤 2: 数据加载与预处理
由于数据集已经按照类别分文件夹存放,直接使用ImageFolder来加载数据。同时进行一些基本的数据增强操作以提高模型泛化能力。
fromtorchvisionimporttransforms,datasetsfromtorch.utils.dataimportDataLoaderimportos# 设置路径data_dir='path_to_your_dataset'# 替换为你的数据集路径# 图像变换定义data_transforms={'train':transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]),'val':transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])]),}image_datasets={x:datasets.ImageFolder(os.path.join(data_dir,x),data_transforms[x])forxin['train','val']}dataloaders={x:DataLoader(image_datasets[x],batch_size=32,shuffle=True,num_workers=4)forxin['train','val']}dataset_sizes={x:len(image_datasets[x])forxin['train','val']}class_names=image_datasets['train'].classes步骤 3: 模型构建
这里我们使用EfficientNet模型,并调整最后一层以适应262类别的分类任务。
importtimmimporttorch.nnasnn model=timm.create_model('efficientnet_b4',pretrained=True)num_ftrs=model.classifier.in_features model.classifier=nn.Linear(num_ftrs,262)# 修改输出类别数device=torch.device("cuda:0"iftorch.cuda.is_available()else"cpu")model=model.to(device)步骤 4: 训练模型
定义损失函数、优化器以及学习率调度策略,然后开始训练过程。
importtorch.optimasoptimfromtqdmimporttqdm criterion=nn.CrossEntropyLoss()optimizer=optim.Adam(model.parameters(),lr=0.0001)scheduler=optim.lr_scheduler.StepLR(optimizer,step_size=7,gamma=0.1)num_epochs=25forepochinrange(num_epochs):print(f'Epoch{epoch}/{num_epochs-1}')print('-'*10)forphasein['train','val']:ifphase=='train':model.train()else:model.eval()running_loss=0.0running_corrects=0forinputs,labelsintqdm(dataloaders[phase]):inputs=inputs.to(device)labels=labels.to(device)optimizer.zero_grad()withtorch.set_grad_enabled(phase=='train'):outputs=model(inputs)_,preds=torch.max(outputs,1)loss=criterion(outputs,labels)ifphase=='train':loss.backward()optimizer.step()running_loss+=loss.item()*inputs.size(0)running_corrects+=torch.sum(preds==labels.data)ifphase=='train':scheduler.step()epoch_loss=running_loss/dataset_sizes[phase]epoch_acc=running_corrects.double()/dataset_sizes[phase]print(f'{phase}Loss:{epoch_loss:.4f}Acc:{epoch_acc:.4f}')这个脚本提供了一个完整的框架,用于使用EfficientNet模型对262种水果进行分类。请根据自己的需求调整参数或添加额外的功能(例如早停机制等)。记得替换path_to_your_dataset为你实际的数据集路径。
数据增强(Data Augmentation)
fromtorchvisionimporttransforms transform=transforms.Compose([transforms.Resize((224,224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225]),transforms.RandomHorizontalFlip(),transforms.RandomRotation(10),transforms.ColorJitter(brightness=0.2,contrast=0.2,saturation=0.2)])3. 模型构建(PyTorch 示例)
importtorchimporttorchvision.modelsasmodelsimporttorch.nnasnndefbuild_model(num_classes=262):model=models.resnet152(pretrained=True)# 替换最后一层全连接层model.fc=nn.Linear(model.fc.in_features,num_classes)returnmodel4. 训练流程设计
超参数设置:
| 参数 | 值 |
|---|---|
| Batch Size | 64 - 128 |
| Epochs | 50 - 100 |
| Optimizer | Adam / SGD with momentum |
| Learning Rate | 1e-4 (Adam), 0.01 (SGD with cosine LR) |
| Loss Function | CrossEntropyLoss |
| LR Scheduler | CosineAnnealingLR / ReduceLROnPlateau |
训练代码框架(PyTorch)
importtorchfromtorch.utils.dataimportDataLoaderfromtorchvisionimportdatasets,transformsfromtqdmimporttqdm# 构建数据集train_dataset=datasets.ImageFolder(root='dataset/train',transform=transform)val_dataset=datasets.ImageFolder(root='dataset/val',transform=transform)train_loader=DataLoader(train_dataset,batch_size=64,shuffle=True)val_loader=DataLoader(val_dataset,batch_size=64)# 构建模型device='cuda'iftorch.cuda.is_available()else'cpu'model=build_model(num_classes=262).to(device)optimizer=torch.optim.Adam(model.parameters(),lr=1e-4)criterion=nn.CrossEntropyLoss()# 训练循环forepochinrange(50):model.train()total_loss=0forimages,labelsintqdm(train_loader):images,labels=images.to(device),labels.to(device)outputs=model(images)loss=criterion(outputs,labels)optimizer.zero_grad()loss.backward()optimizer.step()total_loss+=loss.item()print(f"Epoch{epoch+1}, Loss:{total_loss:.4f}")# 验证model.eval()correct=0total=0withtorch.no_grad():forimages,labelsinval_loader:images,labels=images.to(device),labels.to(device)outputs=model(images)_,preds=torch.max(outputs,1)total+=labels.size(0)correct+=(preds==labels).sum().item()print(f"Validation Accuracy:{correct/total*100:.2f}%")5. 性能评估与指标
- Top-1 Accuracy:单次预测准确率
- Top-5 Accuracy:前五预测中是否包含正确类别
- 混淆矩阵(Confusion Matrix):分析哪些类别容易混淆
- F1 Score、Precision、Recall:适用于不平衡类别分布情况下的评估
6. 模型导出与部署
导出 ONNX 模型:
dummy_input=torch.randn(1,3,224,224).to(device)torch.onnx.export(model,dummy_input,"fruit_classifier.onnx",export_params=True,opset_version=13,do_constant_folding=True,input_names=['input'],output_names=['output'])部署选项:
- ONNX Runtime:轻量级推理引擎,跨平台。
- TensorRT:适用于 NVIDIA GPU 的高性能推理。
- TorchScript / TorchServe:原生 PyTorch 部署方式。
- OpenVINO:适用于 Intel CPU/GPU 的加速推理。
7. 进阶技巧
📌 标签平滑(Label Smoothing)
criterion=LabelSmoothingCrossEntropy(smoothing=0.1)📌 半监督学习(伪标签)
- 在测试集上生成伪标签,加入训练。
📌 使用 Mixup / CutMix 增强训练
fromtimm.data.mixupimportMixup mixup_fn=Mixup(mixup_alpha=0.2,cutmix_alpha=1.0,num_classes=262)images,labels=mixup_fn(images,labels)8. 推荐训练硬件配置
| 组件 | 推荐配置 |
|---|---|
| GPU | NVIDIA RTX 3090 / A100 / V100 |
| RAM | ≥ 64GB |
| 存储 | ≥ 500GB SSD(用于缓存数据) |
总结
| 步骤 | 工具/方法 |
|---|---|
| 模型选择 | EfficientNet / ResNet / Vision Transformer |
| 数据增强 | Resize + Normalize + Flip + ColorJitter |
| 损失函数 | CrossEntropyLoss + LabelSmoothing |
| 优化器 | Adam / SGD + CosineAnnealingLR |
| 部署格式 | ONNX / TorchScript |
| 评估指标 | Top-1/Top-5 Acc、F1、混淆矩阵 |
如果你希望我提供完整的训练脚本或 GUI 界面代码(如 PyQt),也可以继续告诉我,我可以为你打包整个项目模板。
使用EfficientNet模型来训练一个262类水果分类的数据集。这个指南将包括数据加载、模型构建、训练过程以及评估模型性能。
1. 安装必要的库
首先,确保安装了所有需要的Python库:
pipinstalltorch torchvision timm scikit-learn matplotlib tqdmtimm库用于加载预训练的EfficientNet模型。
2. 数据准备
假设你的数据集已经按照以下结构组织好了:
dataset/ ├── train/ │ ├── apple/ │ ├── banana/ │ └── ... ├── val/ │ ├── apple/ │ ├── banana/ │ └── ... └── test/ ├── apple/ ├── banana/ └── ...3. 数据加载器
接下来,编写代码以加载和增强图像数据:
importtorchfromtorch.utils.dataimportDataLoaderfromtorchvisionimportdatasets,transformsfromtorchvision.datasetsimportImageFolder# 定义图像转换transform=transforms.Compose([transforms.Resize((224,224)),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])# 加载数据集train_dataset=ImageFolder(root='dataset/train',transform=transform)val_dataset=ImageFolder(root='dataset/val',transform=transform)train_loader=DataLoader(train_dataset,batch_size=32,shuffle=True)val_loader=DataLoader(val_dataset,batch_size=32)4. 模型定义
使用timm库加载预训练的EfficientNet模型,并替换最后的全连接层以适应262类输出:
importtimmimporttorch.nnasnndefcreate_efficientnet_model(num_classes):model=timm.create_model('efficientnet_b4',pretrained=True)num_ftrs=model.classifier.in_features model.classifier=nn.Linear(num_ftrs,num_classes)returnmodel model=create_efficientnet_model(262)device=torch.device("cuda"iftorch.cuda.is_available()else"cpu")model=model.to(device)5. 训练过程
定义损失函数、优化器和学习率调度器,然后开始训练:
importtorch.optimasoptimfromtqdmimporttqdm criterion=nn.CrossEntropyLoss()optimizer=optim.Adam(model.parameters(),lr=0.0001)scheduler=optim.lr_scheduler.StepLR(optimizer,step_size=7,gamma=0.1)num_epochs=25forepochinrange(num_epochs):print(f'Epoch{epoch}/{num_epochs-1}')print('-'*10)forphasein['train','val']:ifphase=='train':model.train()dataloader=train_loaderelse:model.eval()dataloader=val_loader running_loss=0.0running_corrects=0forinputs,labelsintqdm(dataloader):inputs=inputs.to(device)labels=labels.to(device)optimizer.zero_grad()withtorch.set_grad_enabled(phase=='train'):outputs=model(inputs)_,preds=torch.max(outputs,1)loss=criterion(outputs,labels)ifphase=='train':loss.backward()optimizer.step()running_loss+=loss.item()*inputs.size(0)running_corrects+=torch.sum(preds==labels.data)ifphase=='train':scheduler.step()epoch_loss=running_loss/len(dataloader.dataset)epoch_acc=running_corrects.double()/len(dataloader.dataset)print(f'{phase}Loss:{epoch_loss:.4f}Acc:{epoch_acc:.4f}')6. 模型评估
在测试集上评估模型的性能:
test_dataset=ImageFolder(root='dataset/test',transform=transform)test_loader=DataLoader(test_dataset,batch_size=32)model.eval()running_corrects=0forinputs,labelsintqdm(test_loader):inputs=inputs.to(device)labels=labels.to(device)withtorch.no_grad():outputs=model(inputs)_,preds=torch.max(outputs,1)running_corrects+=torch.sum(preds==labels.data)accuracy=running_corrects.double()/len(test_loader.dataset)print(f'Test Accuracy:{accuracy:.4f}')使用EfficientNet模型对262种水果进行分类。调整一些参数或添加额外的功能(例如数据增强、早停机制等)。