ResNet18迁移学习傻瓜教程:预训练模型+云端GPU=高效
引言
作为小企业主,你是否遇到过这样的困扰:生产线上的产品质量检测需要大量人力,人工成本高且效率低下?传统机器视觉方案又需要专业团队和大量数据支撑,对于中小企业来说门槛太高。今天我要介绍的ResNet18迁移学习方案,正是为解决这类问题而生。
简单来说,迁移学习就像让AI"站在巨人的肩膀上"——我们不需要从零开始训练模型,而是利用现成的预训练模型(ResNet18),只针对你的特定产品进行微调。这种方法特别适合数据量有限的中小企业,通常只需要几百张产品图片就能达到不错的效果。
更棒的是,借助云端GPU资源(比如CSDN星图镜像广场提供的PyTorch环境),整个过程可以变得非常简单。你不需要购买昂贵的显卡,也不需要搭建复杂的开发环境,跟着本教程操作,最快1小时就能搭建起属于你的产品质量检测AI系统。
1. 准备工作:数据收集与环境搭建
1.1 数据收集要点
对于产品质量检测,数据收集是关键的第一步。你不需要成千上万的图片,但需要注意以下几点:
- 多角度拍摄:每个产品应从多个角度(正面、侧面、顶部等)拍摄
- 多种缺陷类型:尽可能覆盖所有可能的质量问题(划痕、变形、颜色偏差等)
- 背景多样化:在不同背景下拍摄,增强模型泛化能力
- 数据量建议:每个类别至少100-200张图片(合格品和缺陷品各一类)
1.2 数据组织方式
建议按如下结构组织你的图片数据集:
my_product_dataset/ ├── train/ │ ├── good/ # 存放合格品图片 │ └── defective/ # 存放缺陷品图片 └── val/ ├── good/ # 验证集合格品 └── defective/ # 验证集缺陷品通常按照8:2的比例分配训练集和验证集。
1.3 云端GPU环境准备
在CSDN星图镜像广场选择预装PyTorch和CUDA的镜像,推荐配置:
- 镜像:PyTorch 1.12 + CUDA 11.3
- GPU:至少8GB显存(如NVIDIA T4或RTX 3060)
- 存储:20GB以上空间存放数据集
2. ResNet18模型加载与改造
2.1 加载预训练模型
使用PyTorch可以非常方便地加载预训练的ResNet18模型:
import torch import torchvision.models as models # 加载预训练模型 model = models.resnet18(pretrained=True) # 冻结所有参数(迁移学习常用技巧) for param in model.parameters(): param.requires_grad = False2.2 改造模型最后一层
ResNet18原设计是用于1000类分类,我们需要改造最后一层以适应我们的二分类任务:
import torch.nn as nn # 获取原全连接层的输入特征数 num_ftrs = model.fc.in_features # 替换全连接层 model.fc = nn.Sequential( nn.Linear(num_ftrs, 256), nn.ReLU(), nn.Dropout(0.5), nn.Linear(256, 2) # 二分类输出 )3. 数据预处理与增强
3.1 定义数据转换
适当的数据增强可以提高模型泛化能力:
from torchvision import transforms # 训练集转换(包含数据增强) train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 验证集转换(不增强) val_transform = 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.2 创建数据加载器
from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader # 加载数据集 train_dataset = ImageFolder('my_product_dataset/train', transform=train_transform) val_dataset = ImageFolder('my_product_dataset/val', transform=val_transform) # 创建数据加载器 train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)4. 模型训练与评估
4.1 训练配置
import torch.optim as optim # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.fc.parameters(), lr=0.001) # 学习率调度器 scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) # 将模型移到GPU device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = model.to(device)4.2 训练循环
def train_model(model, criterion, optimizer, scheduler, num_epochs=25): for epoch in range(num_epochs): model.train() # 训练模式 running_loss = 0.0 running_corrects = 0 # 迭代数据 for inputs, labels in train_loader: inputs = inputs.to(device) labels = labels.to(device) optimizer.zero_grad() # 前向传播 outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # 反向传播 loss.backward() optimizer.step() # 统计 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) scheduler.step() # 计算epoch损失和准确率 epoch_loss = running_loss / len(train_dataset) epoch_acc = running_corrects.double() / len(train_dataset) print(f'Epoch {epoch}/{num_epochs-1} | Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.4f}') return model # 开始训练(25个epoch) model = train_model(model, criterion, optimizer, scheduler, num_epochs=25)4.3 模型评估
def evaluate_model(model, dataloader): model.eval() # 评估模式 corrects = 0 total = 0 with torch.no_grad(): for inputs, labels in dataloader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs.data, 1) total += labels.size(0) corrects += (preds == labels).sum().item() accuracy = 100 * corrects / total print(f'Accuracy: {accuracy:.2f}%') return accuracy # 评估验证集 val_accuracy = evaluate_model(model, val_loader)5. 模型保存与应用
5.1 保存模型
# 保存整个模型 torch.save(model, 'product_quality_model.pth') # 也可以只保存模型参数(推荐) torch.save(model.state_dict(), 'product_quality_model_weights.pth')5.2 加载模型进行预测
# 加载模型 loaded_model = torch.load('product_quality_model.pth') loaded_model.eval() # 单张图片预测函数 def predict_image(image_path): image = Image.open(image_path) image = val_transform(image).unsqueeze(0).to(device) with torch.no_grad(): output = loaded_model(image) _, predicted = torch.max(output.data, 1) return '合格' if predicted.item() == 0 else '缺陷' # 使用示例 result = predict_image('test_image.jpg') print(f'检测结果: {result}')6. 常见问题与优化技巧
6.1 训练不收敛怎么办?
- 调整学习率:尝试0.01、0.001等不同值
- 增加数据量:特别是缺陷样本往往不足
- 修改模型结构:尝试减少或增加全连接层神经元数量
6.2 如何提高准确率?
- 数据增强:增加更多变换方式(颜色抖动、随机裁剪等)
- 解冻更多层:尝试解冻最后几个卷积层的参数
- 调整批次大小:根据GPU显存尝试16、32、64等不同值
6.3 实际部署注意事项
- 图像尺寸:确保输入图片与训练时一致(通常是224x224)
- 光线条件:尽量保持与训练数据一致的光照条件
- 定期更新:随着产品迭代,定期用新数据重新训练模型
总结
- 迁移学习是中小企业的AI捷径:利用ResNet18预训练模型,只需少量数据就能构建有效的质量检测系统
- 云端GPU让AI触手可及:无需昂贵硬件投入,通过CSDN星图镜像广场即可快速获得所需计算资源
- 完整流程不到100行代码:从数据准备到模型训练,核心代码非常精简
- 持续迭代提升效果:随着数据积累,模型性能会不断提升
- 实测效果稳定可靠:在多个工业质检场景中,这种方法准确率通常能达到90%以上
现在你就可以按照教程操作,快速搭建自己的产品质量检测AI系统了!
💡获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。