Pytorch03: 使用预训练模型训练 Kaggle: Dogs vs Cats

Pytorch使用预训练模型训练 Kaggle: Dogs vs Cats

1、Pytorch加载数据集

​ pytorch 的torchvision中继承了常用的数据集，可以直接在torchvision.datasets中调用，当我们想要加载自己的数据集来训练模型时，需要借助torchvision.datasets.ImageFolder这个类，在当前的Pytorch版本中， ImageFolder继承自DatasetFolder官方源码

​ 我们可以直接使用ImageFolder进行数据的读取，但是这要求我们的数据集按照以下方式被组织：

'''
A generic data loader where the samples are arranged in this way: ::
        root/class_x/xxx.ext
        root/class_x/xxy.ext
        root/class_x/xxz.ext
        root/class_y/123.ext
        root/class_y/nsdf3.ext
        root/class_y/asd932_.ext
'''

即保证在文件根目录root下，每一个类别的图像被置于一个以标签为名的文件夹中。

import torch
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torchvision.models as models
import torch.nn.functional as F
import torch.optim as optim
import os
import time
from datetime import timedelta

#指定训练数据与测试数据的路径
traindir = r'F:\DeepLearning\datasets\dogvscat\training_set'
testdir = r'F:\DeepLearning\datasets\dogvscat\test_set'
batch_size = 32
num_worker = 2

#定义加载数据函数load_data() 返回train_loder和test_loder
def load_data(traindir, testdir, batch_size = batch_size, num_worker = num_worker, shuffle = True):
    print('starting load data...')
    train_loder = torch.utils.data.DataLoader(
    datasets.ImageFolder(traindir, transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor()
    ])),
    batch_size= batch_size,
    shuffle = shuffle,
    num_workers = num_worker)
    
    test_loder = torch.utils.data.DataLoader(
    datasets.ImageFolder(testdir, transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor()
    ])),
    batch_size= batch_size,
    shuffle = shuffle,
    num_workers = num_worker)
    
    print('ending load data...')
    
    return train_loder, test_loder

2、加载预训练模型

​ 在进行神经网络模型的训练时，我们往往不会重头开始训练一个网络，而是借助一些在大型数据集上已经训练好的模型完成特征的提取，再增加一些自定义的层完成深度学习任务（如：分类）。在训练的过程中需要冻结预训练模型的各层，只训练自定义的层。

​ 在这个Dogs vs Cats的分类任务中，我主要使用了预训练的VGG16模型，可以直接在torchvision.models中使用vgg16, 并在最后一层FC后增加了一层线性层nn.Linear(input_features=1000, output_features=2)，自定义模型Mymodel:

class Mymodel(nn.Module):
    def __init__(self):
        super(Mymodel, self).__init__()
        self.vgg = models.vgg16()
        self.last_layer = torch.nn.Linear(in_features=1000, out_features=2)
    def forward(self, x):
        x = self.vgg(x)
        x = F.relu(x)
        x = self.last_layer(x)
        
        return x

随后build_model()函数完成预训练模型的权重加载与冻结：

def build_model(path):
    #加载vgg预训练参数
    vgg_state_dict = torch.load(path)
    model = Mymodel()
    model.vgg.load_state_dict(vgg_state_dict)
    
    #冻结部分参数
    for name, p in model.vgg.named_parameters():
        p.requires_grad = False
        
    return model

在定义模型之后可以查看以下模型结构：

print(model)

Mymodel(
  (vgg): VGG(
    (features): Sequential(
      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (1): ReLU(inplace=True)
      (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (3): ReLU(inplace=True)
      (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (6): ReLU(inplace=True)
      (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (8): ReLU(inplace=True)
      (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (11): ReLU(inplace=True)
      (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (13): ReLU(inplace=True)
      (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (15): ReLU(inplace=True)
      (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (18): ReLU(inplace=True)
      (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (20): ReLU(inplace=True)
      (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (22): ReLU(inplace=True)
      (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (25): ReLU(inplace=True)
      (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (27): ReLU(inplace=True)
      (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (29): ReLU(inplace=True)
      (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
    (classifier): Sequential(
      (0): Linear(in_features=25088, out_features=4096, bias=True)
      (1): ReLU(inplace=True)
      (2): Dropout(p=0.5, inplace=False)
      (3): Linear(in_features=4096, out_features=4096, bias=True)
      (4): ReLU(inplace=True)
      (5): Dropout(p=0.5, inplace=False)
      (6): Linear(in_features=4096, out_features=1000, bias=True)
    )
  )
  (last_layer): Linear(in_features=1000, out_features=2, bias=True)
)

定义训练函数train(), 参数包括：

• data_loder: 数据集
• model: 需要训练的模型
• lossFunc： 损失函数（eg:torch.nn.CrossEntropyLoss())
• optimizer: 优化器（eg: torch.optim.Adam)

def train(data_loder, model, lossFunc, optimizer, num_epoch = 1):
    print('Start training...')
    for epoch in range(num_epoch):
        print('Starting epoch %d / %d' %(epoch+1, num_epoch))
        #模型训练
        for i,(images, labels) in enumerate(data_loder):
            predict = model(images)
            optimizer.zero_grad()
            loss = lossFunc(predict, labels)
            loss.backward()
            optimizer.step()
            print('epoch %d, iteration %d, loss %s' % (epoch, i, loss.item()))
    print('End training...')

3、训练模型

#定义模型， 损失， 优化规则
model = build_model(path = r'F:\DeepLearning\models\vgg16-397923af.pth')
lossFunc = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

#加载数据
train_loder, test_loder = load_data(traindir, testdir)

#开始训练
train(train_loder, model, lossFunc, optimizer, num_epoch=2 )

4、测试模型

先可视化部分测试数据：

import matplotlib.pyplot as plt
batch_images = batch_images.permute(0,2,3,1)  #对于多维tensor使用permute()
f = plt.figure(figsize=(12,6))
for i in range(4):
    f.add_subplot(1,4, i+1)
    plt.imshow(batch_images[i])

#在测试集上测试准确性
#在测试集上测试性能
total = 0
correct = 0

iter_data = iter(test_loder)

with torch.no_grad():
    for i in range(20):
        data = iter_data.next()
        images, labels = data
        out = model(images)
        _, predict = torch.max(out, 1)
        total += labels.size(0)
        correct += (labels == predict).sum().item()
        
print('Accuracy: %s' % (100*correct / total))

Accuracy: 94.21875

由于算力的问题，无法训练多个epoch， 并进行模型的fine-tuning，准确率有点小低。。。

5、模型的保存与恢复

​ Pytorch中有两种保存模型的方式，一种是以state_dict的形式保存所有的模型参数，模型恢复时可以调用load_state_dict()进行恢复：

#保存模型参数
torch.save(model.state_dict(), path)
#恢复模型
model = Mymodel()
state_dict = torch.load(path)
model.load_state_dict(state_dict)

只保存模型参数的方式执行起来比较快。

​ 另一种保存模型的方法是保存模型结构+参数：

#保存完整模型
torch.save(model, path)
#恢复完整模型
model = torch.load(path)

这个例子中我采用了第一种方式：

#模型保存参数
save_dir = r'F:\DeepLearning\models'
torch.save(model.state_dict(), os.path.join(save_dir, 'mymodel.pkl'))

6、Reference

1. 数据集地址
2. 完整代码Github