微调预训练模型训练自己的模型

查看一个batch的图像和标注

# DataLoader 是 python生成器，每次调用返回一个 batch 的数据
images, labels = next(iter(train_loader))

可视化

将tensor转为array后即可使用plt调用

images = images.numpy()
plt.hist(images[5].flatten(), bins=50)
plt.show()

载入模型微调

可以在载入后微调训练所有层（适用于源域相似的情景）

model = models.resnet18(pretrained=True) # 载入预训练模型

model.fc = nn.Linear(model.fc.in_features, n_class)

optimizer = optim.Adam(model.parameters())

也可以只训练最后一层（fc layer)，即改变传入优化器的参数。

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 = model.to(device)

# 交叉熵损失函数
criterion = nn.CrossEntropyLoss() 
# 训练轮次 Epoch
EPOCHS = 20
# 学习率降低策略(每隔5轮，降为原先一半)
lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

训练过程

将batch逐一放入模型并进行反向传播

# 获得一个 batch 的数据和标注
images, labels = next(iter(train_loader))
images = images.to(device)
labels = labels.to(device)
# 输入模型，执行前向预测
outputs = model(images)
# 由 logit，计算当前 batch 中，每个样本的平均交叉熵损失函数值
loss = criterion(outputs, labels)
# 反向传播“三部曲”
optimizer.zero_grad() # 清除梯度
loss.backward() # 反向传播
optimizer.step() # 优化更新

完整训练

# 遍历每个 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()

保存模型

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