当前位置：首页 > news >正文

ResNet18迁移学习傻瓜教程：预训练模型+云端GPU=高效

news 2026/8/13 19:23:24

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 = False

2.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}')