ResNet18迁移学习指南：预训练模型+你的数据，1小时出效果-开发者社区

ResNet18迁移学习指南：预训练模型+你的数据，1小时出效果

1. 为什么电商运营需要ResNet18迁移学习

作为电商运营人员，你可能经常遇到这样的困扰：每天要处理成千上万的商品图片，手动分类耗时耗力。传统方法要么准确率低，要么需要大量标注数据。这时候，ResNet18迁移学习就是你的救星。

迁移学习就像让AI"站在巨人的肩膀上"学习。想象一下，ResNet18已经在ImageNet数据集（包含1400万张图片）上训练过，相当于一个见过世面的专家。我们只需要让它针对你的商品图"补补课"，就能快速获得专业级的分类能力。

相比从零训练，迁移学习有三大优势： - 训练速度快：通常1小时内就能出效果 - 数据需求少：几百张标注图片就能用 - 准确率高：借助预训练模型的强大特征提取能力

2. 准备工作：5分钟搞定环境

2.1 基础环境配置

首先确保你的电脑或服务器满足以下条件： - 操作系统：Linux或Windows（推荐Linux） - GPU：NVIDIA显卡（显存≥4GB） - 驱动：已安装CUDA和cuDNN

如果你使用CSDN算力平台，可以直接选择预装PyTorch的镜像，省去环境配置的麻烦。

2.2 安装必要库

打开终端，执行以下命令安装依赖：

pip install torch torchvision pillow numpy matplotlib

2.3 准备数据集

把你的商品图片按类别整理成如下结构：

your_dataset/ ├── train/ │ ├── 类别1/ │ │ ├── img1.jpg │ │ └── img2.jpg │ └── 类别2/ │ ├── img1.jpg │ └── img2.jpg └── val/ ├── 类别1/ └── 类别2/

建议每个类别至少准备50-100张训练图片，20-30张验证图片。图片尺寸不需要统一，ResNet18会自动处理。

3. 实战：1小时训练你的分类模型

3.1 加载预训练模型

创建train.py文件，添加以下代码：

import torch import torchvision.models as models from torchvision import transforms # 加载预训练ResNet18 model = models.resnet18(pretrained=True) # 修改最后一层全连接层 num_classes = 2 # 改为你的类别数 model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

3.2 数据预处理

在train.py中继续添加：

# 数据增强和归一化 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]) ]), }

3.3 加载数据集

继续添加：

from torchvision.datasets import ImageFolder import torch.utils.data as data # 加载数据集 image_datasets = { 'train': ImageFolder('your_dataset/train', data_transforms['train']), 'val': ImageFolder('your_dataset/val', data_transforms['val']) } # 创建数据加载器 dataloaders = { 'train': data.DataLoader(image_datasets['train'], batch_size=32, shuffle=True), 'val': data.DataLoader(image_datasets['val'], batch_size=32, shuffle=False) }

3.4 训练模型

添加训练代码：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = model.to(device) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9) # 训练20个epoch for epoch in range(20): model.train() for inputs, labels in dataloaders['train']: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() # 验证 model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in dataloaders['val']: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f'Epoch {epoch+1}, Val Acc: {100 * correct / total:.2f}%')

运行训练脚本：

python train.py

4. 模型优化与使用技巧

4.1 关键参数调整

学习率(lr)：新手建议0.001-0.0001
批量大小(batch_size)：根据显存调整，32是常用值
训练轮数(epochs)：20-50轮通常足够

4.2 常见问题解决

过拟合：数据太少时可能出现
解决方法：增加数据增强，使用Dropout层
准确率低：
检查数据标注是否正确
尝试微调更多层（不只是最后一层）
显存不足：
减小batch_size
使用更小的图片尺寸

4.3 保存和使用模型

训练完成后保存模型：

torch.save(model.state_dict(), 'resnet18_product.pth')

加载模型进行预测：

model.load_state_dict(torch.load('resnet18_product.pth')) model.eval() # 单张图片预测 def predict(image_path): img = Image.open(image_path) img = data_transforms['val'](img).unsqueeze(0).to(device) output = model(img) _, pred = torch.max(output, 1) return image_datasets['train'].classes[pred.item()]