PyTorch图像分类全流程实战--迁移学习训练图像分类模型03

教程

同济子豪兄：https://space.bilibili.com/1900783

斯坦福CS231N【迁移学习】中文精讲：https://www.bilibili.com/video/BV1K7411W7So

斯坦福CS231N【迁移学习】官方笔记：https://cs231n.github.io/transfer-learning/

引入工具包

数据分析工具:numpy pandas

数据可视化工具：matplotlib seaborn plotly（交互式数据可视化）

HTTP请求：requests

进度条：tqdm

图像处理：opencv-python pillow

记录并可视化训练数据：wandb
wandb: 深度学习轻量级可视化工具入门教程

pytorch工具：torch torchvision torchaudio

预训练图像分类模型

在自己的图像分类数据集上，使用ImageNet预训练图像分类模型初始化，改动分类层，迁移学习微调训练。具体请看PyTorch图像分类全流程实战--预训练模型预测图像分类02

最后保存模型文件为npy文件。

定义数据加载器

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)
# 只微调训练最后一层全连接层的参数，其它层冻结
optimizer = optim.Adam(model.fc.parameters())

#选择二：微调训练所有层
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 = 20
# 获得一个 batch 的数据和标注
images, labels = next(iter(train_loader))
images = images.to(device)
labels = labels.to(device)
# 输入模型，执行前向预测
outputs = model(images)
# 获得当前 batch 所有图像的预测类别 logit 分数
outputs.shape
# 由 logit，计算当前 batch 中，每个样本的平均交叉熵损失函数值
loss = criterion(outputs, labels)
 反向传播“三部曲”
optimizer.zero_grad() # 清除梯度
loss.backward() # 反向传播
optimizer.step() # 优化更新
# 获得当前 batch 所有图像的预测类别
_, preds = torch.max(outputs, 1)

# 遍历每个 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)
        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):
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, preds = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (preds == labels).sum()

    print('测试集上的准确率为 {:.3f} %'.format(100 * correct / total))

保存模型

torch.save(model, 'checkpoints/fruit30_pytorch_20220814.pth')

在pytorch进行模型保存的时候，一般有两种保存方式，一种是保存整个模型，另一种是只保存模型的参数。

torch.save(model.state_dict(), "my_model.pth") # 只保存模型的参数

torch.save(model, "my_model.pth") # 保存整个模型

待完成...

代码还没有运行，后续补上运行之后的截图

参考文献

【1】wandb: 深度学习轻量级可视化工具入门教程
【2】PyTorch图像分类全流程实战--预训练模型预测图像分类02
【3】pytorch中保存的模型文件.pth深入解析