PyTorch实战6：好莱坞明星识别--VGG16_pytorch vgg16 参数设置

本次实战主要学习内容：

• 调用官方的VGG-16网络框架
• 保存训练过程中的最佳模型权重

一、调用官方的VGG-16模型

使用 PyTorch 的 torchvision 库，加载预训练的 VGG16 模型，并对模型进行微调以用于图像分类任务。

第一行代码检查是否有可用的 GPU 设备，如果有则将设备设置为 “cuda”，否则为 “cpu”。

from torchvision.models import vgg16

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using {} device".format(device))
1
2
3
4

接下来，通过调用 vgg16() 函数并传入参数 pretrained=True 从 torchvision 库中加载预训练的 VGG16 模型，并将其移动到指定的设备上。

model = vgg16(pretrained=True).to(device)
1

为了冻结模型的所有参数，即使在微调过程中也不对它们进行更新，使用一个简单的 for 循环遍历所有参数并将 requires_grad 属性设置为 False。

for param in model.parameters():
    param.requires_grad = False
1
2

最后，修改了 VGG16 模型的最后一层全连接层，以输出目标类别个数。在本例中，将输出个数设置为 len(classeNames)，其中 classeNames 是定义了数据集中类别名称的 Python 列表。

model.classifier._modules['6'] = nn.Linear(4096,len(classeNames)) 
model.to(device)  
model
1
2
3

值得注意的是：在以上代码中，直接修改了 VGG16 模型的最后一层，这意味着将无法使用预训练的权重。如果你想要使用预训练的权重，可以在创建新的全连接层时调用 nn.Linear() 函数并传入参数 in_features 和 out_features，以保持原有的权重不变。

完整代码如下：

from torchvision.models import vgg16

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using {} device".format(device))
    
# 加载预训练模型，并且对模型进行微调
model = vgg16(pretrained = True).to(device) # 加载预训练的vgg16模型

for param in model.parameters():
    param.requires_grad = False # 冻结模型的参数，这样子在训练的时候只训练最后一层的参数

# 修改classifier模块的第6层（即：(6): Linear(in_features=4096, out_features=2, bias=True)）
# 注意查看我们下方打印出来的模型
model.classifier._modules['6'] = nn.Linear(4096,len(classeNames)) # 修改vgg16模型中最后一层全连接层，输出目标类别个数
model.to(device)  
model
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

二、设置动态学习率

• 本例调用官方动态学习率接口
• 动态学习率具体详解可参考PyTorch实战5：运动鞋识别之动态学习率

使用PyTorch框架中的学习率调整技术，具体来说是使用了LambdaLR调度器。

首先，定义了一个名为lambda1的匿名函数，其输入为epoch，表示当前训练的轮数（epoch），返回值为0.92的epoch//4次方。这个函数的作用是计算每一轮训练的学习率，采用指数衰减的方式，即每经过4个epoch，学习率就会降低为原来的0.92倍。

接下来，定义了一个SGD优化器，它会对模型参数进行优化，其中lr参数指定了初始学习率。

然后，通过调用torch.optim.lr_scheduler.LambdaLR方法创建了一个LambdaLR调度器，并将优化器和上面定义的lambda1函数作为参数传入。这个调度器的作用是根据指定的学习率变化规律，动态地调整优化器中的学习率。在每次更新模型参数之前，都会调用LambdaLR调度器的step()方法，从而更新当前的学习率。

代码如下：

lambda1 = lambda epoch: 0.92 ** (epoch // 4)
optimizer = torch.optim.SGD(model.parameters(), lr=learn_rate)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1) #选定调整方法
1
2
3

三、保存训练模型最佳权重

在模型训练过程中，如果当前测试准确率（epoch_test_acc）比之前的最好准确率（best_acc）更高，就保存当前的模型（使用深拷贝方法），并更新最好准确率。具体的注释如下：

# 判断当前测试准确率是否超过了之前记录的最佳准确率
if epoch_test_acc > best_acc:
    # 如果是，则更新最佳准确率和最佳模型（进行深拷贝）
    best_acc   = epoch_test_acc
    best_model = copy.deepcopy(model)
1
2
3
4
5

需要注意的是，这段代码只会在测试集上进行准确率的比较和模型的保存，并不会对训练集上的准确率进行记录或统计。同时，由于可能存在过拟合等问题，最终的“最佳模型”不一定能够在实际场景中取得最好的表现。

最后将最佳模型保存到文件中

PATH = './model_save/best_model.pth'  # 保存的参数文件名
torch.save(model.state_dict(), PATH)
1
2

• 此处代码详解可移步至PyTorch实战4：猴痘病识别了解并学习
• model.train()、model.eval()的用法在PyTorch实战1：实现mnist手写数字识别中有详细解释

完整代码如下：

import copy

# 定义交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()

# 设置训练轮数
epochs = 40

# 定义空列表，用于保存训练和测试的准确率和损失值
train_loss = []
train_acc = []
test_loss = []
test_acc = []

# 设置一个最佳模型的初始准确率，作为最佳模型的判别指标
best_acc = 0    

# 开始进行模型训练
for epoch in range(epochs):
    
    # 更新学习率（使用自定义学习率时使用）
    # adjust_learning_rate(optimizer, epoch, learn_rate)
    
    # 将模型设置为训练模式
    model.train()
    
    # 进行一次训练，并获取训练准确率和训练损失值
    epoch_train_acc, epoch_train_loss = train(train_dl, model, loss_fn, optimizer)
    
    # 更新学习率（调用官方动态学习率接口时使用）
    scheduler.step() 
    
    # 将模型设置为评估模式
    model.eval()
    
    # 进行一次测试，并获取测试准确率和测试损失值
    epoch_test_acc, epoch_test_loss = test(test_dl, model, loss_fn)
    
    # 如果本轮测试准确率比之前的最佳准确率更高，则更新最佳准确率和所对应的模型
    if epoch_test_acc > best_acc:
        best_acc   = epoch_test_acc
        best_model = copy.deepcopy(model)   # 复制当前模型参数
    
    # 将本轮训练和测试的准确率和损失值保存到列表中
    train_acc.append(epoch_train_acc)
    train_loss.append(epoch_train_loss)
    test_acc.append(epoch_test_acc)
    test_loss.append(epoch_test_loss)
    
    # 获取当前的学习率
    lr = optimizer.state_dict()['param_groups'][0]['lr']
    
    # 输出当前轮次的训练和测试指标
    template = ('Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%, Test_loss:{:.3f}, Lr:{:.2E}')
    print(template.format(epoch+1, epoch_train_acc*100, epoch_train_loss, 
                          epoch_test_acc*100, epoch_test_loss, lr))

# 保存最佳模型到文件中
PATH = './model_save/best_model.pth'  # 保存的参数文件名
torch.save(model.state_dict(), PATH)

print('Done')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62