pytorch迁移学习训练图像分类

一、环境配置
二、迁移学习关键代码
三、完整代码
四、结果对比

代码和图片等资源均来源于哔哩哔哩up主：同济子豪兄
讲解视频：Pytorch迁移学习训练自己的图像分类模型

一、环境配置

1，安装所需的包

pip install numpy pandas matplotlib seaborn plotly requests tqdm opencv-python pillow wandb -i https://pypi.tuna.tsinghua.edu.cn/simple

2，安装Pytorch

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

3，创建目录

import os
# 存放训练得到的模型权重
os.mkdir('checkpoint')

4，下载数据集压缩包（下载之后需要解压数据集）

wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/fruit30/fruit30_split.zip

二、迁移学习关键代码

以下是迁移学习的三种选择，根据训练的需求选择不同的迁移方法：

选择一：只微调训练模型最后一层（全连接分类层）

model = models.resnet18(pretrained=True) # 载入预训练模型
# 修改全连接层，使得全连接层的输出与 当前数据集类别数n_class 对应
model.fc = nn.Linear(model.fc.in_features, n_class)
# 只微调训练最后一层全连接层的参数，其它层冻结
optimizer = optim.Adam(model.fc.parameters())

选择二：微调训练所有层。

适用于训练数据集与预训练模型相差大时，可以选择微调训练所有层，此时只使用预训练模型的部分权重和特征，例如原始模型为imageNet，而训练数据为医疗相关

model = models.resnet18(pretrained=True) # 载入预训练模型
model.fc = nn.Linear(model.fc.in_features, n_class)
optimizer = optim.Adam(model.parameters())

选择三：随机初始化模型全部权重，从头训练所有层

model = models.resnet18(pretrained=False) # 只载入模型结构，不载入预训练权重参数
model.fc = nn.Linear(model.fc.in_features, n_class)
optimizer = optim.Adam(model.parameters())

三、完整代码

import time
import osimport numpy as np
from tqdm import tqdmimport torch
import torchvision
import torch.nn as nn# 忽略出现的红色提示
import warnings
warnings.filterwarnings("ignore")# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device', device)from torchvision import transforms# 训练集图像预处理：缩放裁剪、图像增强、转 Tensor、归一化
train_transform = transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])# 测试集图像预处理-RCTN：缩放、裁剪、转 Tensor、归一化
test_transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])# 数据集文件夹路径
dataset_dir = 'fruit30_split'
train_path = os.path.join(dataset_dir, 'train')	# 测试集路径
test_path = os.path.join(dataset_dir, 'val')	# 测试集路径from torchvision import datasets# 载入训练集
train_dataset = datasets.ImageFolder(train_path, train_transform)# 载入测试集
test_dataset = datasets.ImageFolder(test_path, test_transform)# 各类别名称
class_names = train_dataset.classes
n_class = len(class_names)# 定义数据加载器DataLoader
from torch.utils.data import DataLoaderBATCH_SIZE = 32# 训练集的数据加载器
train_loader = DataLoader(train_dataset,batch_size=BATCH_SIZE,shuffle=True,num_workers=4)# 测试集的数据加载器
test_loader = DataLoader(test_dataset,batch_size=BATCH_SIZE,shuffle=False,num_workers=4)from torchvision import models
import torch.optim as optim# 选择一：只微调训练模型最后一层（全连接分类层）
model = models.resnet18(pretrained=True) # 载入预训练模型
# 修改全连接层，使得全连接层的输出与当前数据集类别数对应
# 新建的层默认 requires_grad=True，指定张量需要梯度计算
model.fc = nn.Linear(model.fc.in_features, n_class)
model.fc	# 查看全连接层
# 只微调训练最后一层全连接层的参数，其它层冻结
optimizer = optim.Adam(model.fc.parameters())    # optim 是 PyTorch 的一个优化器模块，用于实现各种梯度下降算法的优化方法# 选择二：微调训练所有层
# 训练数据集与预训练模型相差大时，可以选择微调训练所有层，只使用预训练模型的部分权重和特征，例如原始模型为imageNet，训练数据为医疗相关
# model = models.resnet18(pretrained=True) # 载入预训练模型
# model.fc = nn.Linear(model.fc.in_features, n_class)
# optimizer = optim.Adam(model.parameters())# 选择三：随机初始化模型全部权重，从头训练所有层
# model = models.resnet18(pretrained=False) # 只载入模型结构，不载入预训练权重参数
# model.fc = nn.Linear(model.fc.in_features, n_class)
# optimizer = optim.Adam(model.parameters())# 训练配置
model = model.to(device)# 交叉熵损失函数
criterion = nn.CrossEntropyLoss()# 训练轮次 Epoch
EPOCHS = 30# 遍历每个 EPOCH
for epoch in tqdm(range(EPOCHS)):model.train()for images, labels in train_loader:  # 获取训练集的一个 batch，包含数据和标注images = images.to(device)labels = labels.to(device)outputs = model(images)           # 前向预测，获得当前 batch 的预测结果loss = criterion(outputs, labels) # 比较预测结果和标注，计算当前 batch 的交叉熵损失函数optimizer.zero_grad()loss.backward()                   # 损失函数对神经网络权重反向传播求梯度optimizer.step()                  # 优化更新神经网络权重# 测试集上初步测试
model.eval()
with torch.no_grad():correct = 0total = 0for images, labels in tqdm(test_loader): # 获取测试集的一个 batch，包含数据和标注images = images.to(device)labels = labels.to(device)outputs = model(images)              # 前向预测，获得当前 batch 的预测置信度_, preds = torch.max(outputs, 1)     # 获得最大置信度对应的类别，作为预测结果total += labels.size(0)correct += (preds == labels).sum()   # 预测正确样本个数print('测试集上的准确率为 {:.3f} %'.format(100 * correct / total))# 保存模型
torch.save(model, 'checkpoint/fruit30_pytorch_A1.pth') # 选择一：微调全连接层
# torch.save(model, 'checkpoint/fruit30_pytorch_A2.pth') # 选择二：微调所有层
# torch.save(model, 'checkpoint/fruit30_pytorch_A3.pth') # 选择三：随机权重

四、结果对比

调用不同迁移学习得到的模型对比测试集准确率

# 测试集导入和图像预处理等代码和上述完整代码中一致，此处省略……# 调用自己训练的模型
model = torch.load('checkpoint/fruit30_pytorch_A1.pth')# 测试集上进行测试
model.eval()
with torch.no_grad():correct = 0total = 0for images, labels in tqdm(test_loader): # 获取测试集的一个 batch，包含数据和标注images = images.to(device)labels = labels.to(device)outputs = model(images)              # 前向预测，获得当前 batch 的预测置信度_, preds = torch.max(outputs, 1)     # 获得最大置信度对应的类别，作为预测结果total += labels.size(0)correct += (preds == labels).sum()   # 预测正确样本个数print('测试集上的准确率为 {:.3f} %'.format(100 * correct / total))