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 os
import numpy as np
from tqdm import tqdm
import 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 DataLoader
BATCH_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 = 0
total = 0
for 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 = 0
total = 0
for 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))
结果如下：
对于微调全连接层的选择一，测试集准确率为 72.078%
而所有权重随机的选择三测试集准确率为 43.228%
总体而言，迁移学习能够利用已有的知识和经验，加速模型的训练过程，提高模型的性能。
