Pytorch学习笔记（I）——预训练模型（一）：加载与使用_如何判断预训练模型加载成功了

为完成自己的科研任务，当前我需要基于VGG16做fine-tuning。于是写下这一节笔记。
我使用的是torch1.0，因此本博客主要基于这篇博客——pytorch finetuning 自己的图片进行行训练做调整

一、加载预训练模型

import torch
import torchvision
# prepare model
mode1_vgg16 = torchvision.models.vgg16(pretrained=True)
1
2
3
4

pretrained 设置为 True，会自动下载模型所对应权重，并加载到模型中

也可以自己下载权重，然后 load 到 模型中，源码中有权重的地址：https://github.com/aaron-xichen/pytorch-playground。

二、设置网络参数

假设我们的类别数为2（即需要修改最后一层全连接层的神经元数量）
先来看看vgg16最后几层

与Resnet（分类器只有一个名为fc的全连接层）不同，vgg16的分类器（classifier）由多层组成，而我们要修改的仅为最后一层的输出1000
因此在，第一步的基础上添加以下代码

#读取输入特征的维度（因为这一维度不需要修改）
num_fc = mode1_vgg16.classifier[6].in_features
#修改最后一层的输出维度，即分类数
mode1_vgg16.classifier[6] = torch.nn.Linear(num_fc,2)
1
2
3
4

可以看到最后一层的输出被修改了。

三、固定权值参数

在先前的基础上添加以下代码

#对于模型的每个权重，使其不进行反向传播，即固定参数
for param in mode1_vgg16.parameters():
    param.requires_grad = False
#但是参数全部固定了，也没法进行学习，所以我们不固定最后一层，即全连接层
for param in mode1_vgg16.classifier[6].parameters():
    param.requires_grad = True
1
2
3
4
5
6

注：可以通过以下代码查看权值数值（可加可不加）

for child in mode1_vgg16.children():
    print(child)
    for param in child.parameters():
        print(param)
1
2
3
4

未固定的权重会有requires_grad=True，而被固定的权值为requires_grad=False默认不显示

四、训练模型

先在代码开头补上接下来需要使用的模块

import numpy as np
from torchvision import transforms,utils
from torch.utils.data import DataLoader
from torch.autograd import Variable
1
2
3
4

1、加载图片

1） 这种读取图片的方式用的是torch自带的 ImageFolder，读取的文件夹必须在一个大的子文件下，按类别归好类。就像我现在要区分2个类别

2） 随后用代码导入数据

#train data
train_data=torchvision.datasets.ImageFolder('/disk2/lockonlxf/pin/trainData',transform=transforms.Compose(
                                                                        [
                                                                            transforms.Resize(256),
                                                                            transforms.CenterCrop(224),
                                                                            transforms.ToTensor()
                                                                        ]))
train_loader = DataLoader(train_data,batch_size=20,shuffle=True)

#test data
test_data = torchvision.datasets.ImageFolder('/disk2/lockonlxf/pin/testData',transform=transforms.Compose(
                                                                        [
                                                                            transforms.Resize(256),
                                                                            transforms.CenterCrop(224),
                                                                            transforms.ToTensor()
                                                                        ]))
test_loader = DataLoader(test_data,batch_size=20,shuffle=True)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

2、训练

1） 定义自己的优化器

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(mode1_vgg16.parameters(),lr=0.001)
1
2

2） 训练模型

EPOCH = 1
for epoch in range(EPOCH):
    train_loss = 0.
    train_acc = 0.
    for step, data in enumerate(train_loader):
        batch_x, batch_y = data
        batch_x, batch_y = Variable(batch_x), Variable(batch_y)
        batch_x, batch_y = batch_x.cuda(), batch_y.cuda()#GPU
        # batch_y not one hot
        # out is the probability of eatch class
        # such as one sample[-1.1009  0.1411  0.0320],need to calculate the max index
        # out shape is batch_size * class
        out = mode1_vgg16(batch_x)
        loss = criterion(out, batch_y)
        train_loss += loss.item()
        # pred is the expect class
        # batch_y is the true label
        pred = torch.max(out, 1)[1]
        train_correct = (pred == batch_y).sum()
        train_acc += train_correct.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step % 100 == 0:
            print('Epoch: ', epoch, 'Step', step,
                  'Train_loss: ', train_loss / ((step + 1) * 20), 'Train acc: ', train_acc / ((step + 1) * 20))

    print('Epoch: ', epoch, 'Train_loss: ', train_loss / len(train_data), 'Train acc: ', train_acc / len(train_data))
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

五、测试

mode1_vgg16.eval()
eval_loss=0
eval_acc=0
for step ,data in enumerate(test_loader):
    batch_x,batch_y=data
    batch_x,batch_y=Variable(batch_x),Variable(batch_y)
    batch_x, batch_y = batch_x.cuda(), batch_y.cuda()
    out = mode1_vgg16(batch_x)
    loss = criterion(out, batch_y)
    eval_loss += loss.item()
    # pred is the expect class
    # batch_y is the true label
    pred = torch.max(out, 1)[1]
    test_correct = (pred == batch_y).sum()
    eval_acc += test_correct.item()
print( 'Test_loss: ', eval_loss / len(test_data), 'Test acc: ', eval_acc / len(test_data))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16