【预训练-微调】迁移学习项目实战自定义数据集宝可梦精灵-学习笔记

本学习笔记来源于B站：深度学习—迁移学习项目实战自定义数据集宝可梦精灵。

在本预训练-微调代码中，重点要学习的内容包括：加载官方提供的经典网络架构resnet18和已经训练好的模型，对最后一层全连接层进行修改，改为适合自己任务的网络架构。在此基础上，训练最后一层全连接层，并保存最优模型。

与上一篇 预训练-微调的学习笔记 不同，本视频教程中提供了另一种分割经典网络架构与预训练模型的方法，在代码中主要改动了两个位置。

改动位置一：

from torchvision.models import resnet18      # 加载官网上训练好的 resnet18

改动位置二：

    trained_model = resnet18(pretrained=True)
    # 通过 pretrained=True 指定该模型使用在大规模数据集上（如 ImageNet）预训练的权重。
    model = nn.Sequential(*list(trained_model.children())[:-1],           # [b, 512, 1, 1]
                          # trained_model.children() 返回一个包含该模型所有子模块的迭代器。使用 list(...) 将其转换为列表。
                          # [:-1] 表示切片操作，将列表中的最后一个子模块（通常是分类层）去掉，因此取出除了最后一层以外的所有层（即 1-17层）。
                          nn.Flatten(),         # 添加了一个 Flatten 层     # [b, 512, 1, 1] => [b, 512]
                          nn.Linear(512, 5)     # 添加了一个全连接层（线性层）
                          ).to(device)

主程序完整代码如下：

import torch
from torch import optim, nn
import visdom
from torch.utils.data import DataLoader
from pokemon import Pokemon

############################################ 改动位置一
from torchvision.models import resnet18      # 加载官网上训练好的 resnet18
############################################


batchsz = 32
lr = 1e-3
epochs = 10

device = torch.device('cuda')
torch.m