Pytorch loading uses pre-training model and fine tune model fine-tuning (freezing part of the layer) actual combat

In the past two days, I have been learning pytorch's loading pre-training model and fine tune. For the convenience of future viewing, I will feature it as a blog.

1. pytorch pre-training model and modification

Pytorch comes with several commonly used deep learning network pre-training models. The torchvision.models package contains alexnet, densenet, inception, resnet, squeezenet, vgg and other common network structures, and provides pre-training models, which can be read by calling Network structure and pre-trained model (model parameters). Often in order to speed up the learning progress, the pre-trained parameters in the pretrain model are directly loaded in the early stage of training.

There are two ways to load the model: a, load the network structure and pre-training parameters; b, only load the network structure, not load the pre-training parameters. How to choose the two forms depends on your needs.

# pretrained=True 加载网络结构和预训练参数，
# pretrained=False 时代表只加载网络结构，不加载预训练参数，即不需要用预训练模型的参数来初始化
# pretrained 参数默认是False,为了代码清晰，最好还是加上参数赋值
net = models.vgg16(pretrained=True)

There are two ways to modify the loaded pre-training model:
a. The last layer of the network classification layer fc is to divide 1000 types. For your own data set, if there are only 2 types, only modify the output of the last fully connected layer as 2.
b. To modify the hierarchical structure in the network by adding or subtracting convolution, at this time, only the method of parameter coverage can be used, that is, first define a similar network by yourself, and then extract the parameters in the pre-training into your own network.
How to implement the two forms Let's look at the code:

# 方法一：增减卷积 要修改网络中的层次结构，这时只能用参数覆盖的方法，即自己先定义一个类似的网络，再将预训练中的参数提取到自己的网络中来
class Dgo_Cat_Net1(nn.Module):
    def __init__(self, num_classes=2):
        super(Dgo_Cat_Net1, self).__init__()

        # pretrained=True 加载网络结构和预训练参数，
        # pretrained=False 时代表只加载网络结构，不加载预训练参数，即不需要用预训练模型的参数来初始化
        # pretrained 参数默认是False,为了代码清晰，最好还是加上参数赋值
        net = models.vgg16(pretrained=True)
        net.classifier = nn.Sequential()    # 将分类层（fc）置空
        self.features = net
        self.classifier = nn.Sequential(    # 定义一个卷积网络结构
            nn.Linear(512*7*7, 512),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(512, 128),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x


# 方法二：网络最后一层分类层fc是对1000种类型进行划分，对于自己的数据集，如果只有2类
class Dgo_Cat_Net2(nn.Module):
    def __init__(self, num_classes=2):
        super(Dgo_Cat_Net2, self).__init__()

        # pretrained=True 加载网络结构和预训练参数，False 时代表只加载网络结构，不加载预训练参数，即不需要用预训练模型的参数来初始化
        # pretrained 参数默认是False,为了代码清晰，最好还是加上参数赋值
        self.model = models.resnet50(pretrained=True)                                   # 调用模型
        fc_features = self.model.fc.in_features                                         # 提取 fc 层中固定的参数 in_features
        self.model.fc = nn.Linear(in_features=fc_features, out_features=num_classes)    # 修改 fc 层中 out_features 参数，修改分类为9

    def forward(self, x):
        x = self.model(x)
        return x

2. pytorch fine tune fine-tuning (freeze part of the layer)

Another way to use a pretrained model is to partially train it. The specific method is to keep the weights of some initial layers of the model unchanged, freeze them, and retrain the subsequent layers to obtain new weights. In this process, multiple attempts can be made, so that the best match between frozen layers and retrain layers can be found based on the results.
How to use the pre-trained model is determined by the size of the data set and the similarity of the data between the old and new data sets (the pre-trained data set and the data set to be solved).

model = Dgo_Cat_Net1(num_classes=2)
for i, param in enumerate(model.parameters()):
    print(i)
    if i < 27:      # 前面一些参数冻结
        param.requires_grad = False

All codes of the project:
main.py

import torch
from torch import nn
import torchvision.models as models
from dataload import new_test_loader, new_train_loader
import torch.optim as optim
import matplotlib.pyplot as plt
import time
import math


LR = 0.0005    # 设置学习率
EPOCH_NUM = 10  # 训练轮次

def time_since(since):
    s = time.time() - since
    m = math.floor(s/60)
    s -= m*60
    return '%dm %ds' % (m, s)

# 方法一：增减卷积 要修改网络中的层次结构，这时只能用参数覆盖的方法，即自己先定义一个类似的网络，再将预训练中的参数提取到自己的网络中来
class Dgo_Cat_Net1(nn.Module):
    def __init__(self, num_classes=2):
        super(Dgo_Cat_Net1, self).__init__()

        # pretrained=True 加载网络结构和预训练参数，
        # pretrained=False 时代表只加载网络结构，不加载预训练参数，即不需要用预训练模型的参数来初始化
        # pretrained 参数默认是False,为了代码清晰，最好还是加上参数赋值
        net = models.vgg16(pretrained=True)
        net.classifier = nn.Sequential()    # 将分类层（fc）置空
        self.features = net
        self.classifier = nn.Sequential(    # 定义一个卷积网络结构
            nn.Linear(512*7*7, 512),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(512, 128),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x


# 方法二：网络最后一层分类层fc是对1000种类型进行划分，对于自己的数据集，如果只有2类
class Dgo_Cat_Net2(nn.Module):
    def __init__(self, num_classes=2):
        super(Dgo_Cat_Net2, self).__init__()

        # pretrained=True 加载网络结构和预训练参数，False 时代表只加载网络结构，不加载预训练参数，即不需要用预训练模型的参数来初始化
        # pretrained 参数默认是False,为了代码清晰，最好还是加上参数赋值
        self.model = models.resnet50(pretrained=True)                                   # 调用模型
        fc_features = self.model.fc.in_features                                         # 提取 fc 层中固定的参数 in_features
        self.model.fc = nn.Linear(in_features=fc_features, out_features=num_classes)    # 修改 fc 层中 out_features 参数，修改分类为9

    def forward(self, x):
        x = self.model(x)
        return x

model = Dgo_Cat_Net1(num_classes=2)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.to(device)

train_data = new_train_loader
test_data = new_test_loader

# 总共有 44 层网络，37层预加载VGG模型里面的，还有 7 层外面自己加的，我们把前面一些层预加载的模型冻结住，后面的一些层更新
para_optim = []
# for i, single_layer in  enumerate(model.modules()):
#     # print(i, single_layer)
#     if i > 36:      # 前面37层冻结
#         for param in single_layer.parameters():
#             para_optim.append(param)
#     else:           # 后面7层不冻结正常更新
#         for param in single_layer.parameters():
#             param.requires_grad = False
# print(f'para_optim len = {len(para_optim)}')

for i, param in enumerate(model.parameters()):
    print(i)
    if i < 27:      # 前面一些参数冻结
        param.requires_grad = False



criterion = torch.nn.CrossEntropyLoss()
# optimizer = optim.Adam(model.parameters(), lr=LR)
# optimizer = optim.Adam(para_optim, lr=LR)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR)

def train(epoch, loss_list):
    # print(f'moduls len:{len(model.modules())}  children len:{len(model.children())}')
    n, m = 0, 0
    for k in model.modules():
        n += 1
        # print(k)
    print(f'n={n}')
    for k in model.children():
        # print(k)
        m += 1
    print(f'm={m}')
    running_loss = 0.0
    for batch_idx, data in enumerate(train_data, 0):
        inputs, target = data[0], data[1]
        inputs, target = inputs.to(device), target.to(device)
        optimizer.zero_grad()

        outputs = model(inputs)

        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        loss_list.append(loss.item())
        running_loss += loss.item()
        if batch_idx % 100 == 99:
            print(f'[{time_since(start)}] Epoch {epoch}', end='')
            print('[%d, %5d] loss:%.3f' % (epoch + 1, batch_idx + 1, running_loss / 100))
            running_loss = 0.0

    return loss_list

def test():
    correct = 0
    total = 0
    with torch.no_grad():
        for _, data in enumerate(new_test_loader, 0):
            inputs, target = data[0], data[1]
            inputs, target = inputs.to(device), target.to(device)
            outputs = model(inputs)
            _, prediction = torch.max(outputs.data, dim=1)

            total += target.size(0)
            correct += (prediction == target).sum().item()
        print('Accuracy on test set: (%d/%d)%d %%' % (correct, total, 100 * correct / total))
        with open("test.txt", "a") as f:
            f.write('Accuracy on test set: (%d/%d)%d %% \n' % (correct, total, 100 * correct / total))


if __name__ == '__main__':
    start = time.time()

    with open("test.txt", "a") as f:
        f.write('Start write!!! \n')

    loss_list = []
    for epoch in range(EPOCH_NUM):
        train(epoch, loss_list)
        test()
    torch.save(model.state_dict(), 'Model.pth')

    x_ori = []
    for i in range(len(loss_list)):
        x_ori.append(i)
    plt.title("Graph")
    plt.plot(x_ori, loss_list)
    plt.ylabel("Y")
    plt.xlabel("X")
    plt.show()

dataload.py (cat and dog dataset loading)

import torch
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import os
from PIL import Image

# 初始化根目录
train_path = 'D:\\DeapLearn Project\\Vgg16_findtune_classificate_CatDog\\CatDogData\\train\\'
test_path = 'D:\\DeapLearn Project\\Vgg16_findtune_classificate_CatDog\\CatDogData\\test\\'

# 定义读取文件的格式
# 数据集
class MyDataSet(Dataset):
    def __init__(self, data_path:str, transform=None):
        super(MyDataSet, self).__init__()
        self.data_path = data_path
        if transform is None:
            self.transform = transforms.Compose(
                [
                    transforms.Resize(size=(224, 224)),
                    transforms.ToTensor(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                ]
            )
        else:
            self.transform = transforms
        self.path_list = os.listdir(data_path)

    def __getitem__(self, idx:int):
        img_path = self.path_list[idx]
        if img_path.split('.')[0] == 'dog':
            label = 1
        else:
            label = 0
        label = torch.as_tensor(label, dtype=torch.int64)
        img_path = os.path.join(self.data_path, img_path)
        img = Image.open(img_path)
        img = self.transform(img)
        return img, label

    def __len__(self)->int:
        return len(self.path_list)

train_ds = MyDataSet(train_path)

full_ds = train_ds
train_size = int(0.8*len(full_ds))
test_size = len(full_ds) - train_size
new_train_ds, test_ds = torch.utils.data.random_split(full_ds, [train_size, test_size])

# 数据加载
new_train_loader = DataLoader(new_train_ds, batch_size=32, shuffle=True, pin_memory=True, num_workers=0)
new_test_loader = DataLoader(test_ds, batch_size=32, shuffle=False, pin_memory=True, num_workers=0)

In the end, the model ran for 10 rounds, and the maximum accuracy rate was 94%. We can see that using the pre-training model that comes with pytorch, adding some layers added by ourselves, and then freezing the parameters in some of the previous pre-loaded models, only After updating the parameters in the latter part, the effect of the final model is still very good.

Reference article:
pytorch pre-training
project GitHub: Vgg16_findtune_classificate_CatDog
Contact email: [email protected]