【AI实战】Pytorch版使用预训练卷积神经网络快速训练自己的分类模型

常见的卷积神经网络

• alexnet
• vgg
• resnet
• inception
• densenet
• googlenet

预训练模型

ResNet 有很多变种，包括 ResNet 18、ResNet 34、ResNet 50、ResNet 101、ResNet 152。

这里以 ResNet 18 为例子详细说明Pytorch版使用预训练模型快速训练自己的分类模型的过程。

使用预训练的 ResNet 18 快速训练自己的分类模型核心步骤

• 1、首先加载训练好的模型参数：

resnet18 = models.resnet18()

• 2、修改全连接层的输出

num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs, 4)#这里我们的模型分为 4 类

• 3、加载模型参数

checkpoint = torch.load(m_path)
resnet18.load_state_dict(checkpoint['model_state_dict'])

• 4、resnet18预训练模型下载

  wget “https://download.pytorch.org/models/resnet18-5c106cde.pth”

  或者：

  浏览器下载：https://download.pytorch.org/models/resnet18-5c106cde.pth

使用预训练的 ResNet 18 快速训练自己的分类模型详细过程

• 1、依赖包
  先导入依赖包：

import os
import torch.utils.data as data
import torch
import torch.optim as optim
import torch.nn as nn
from torch.optim import lr_scheduler
from torchvision import datasets, models, transforms
from PIL import Image
import numpy as np
import cv2
import sys

• 2、加载预训练的模型

    model_ft = models.resnet18(pretrained=True)
    num_fits = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层 NUMCLASS=4

• 3、设置 device

    device = torch.device('cuda:0') # 默认使用 GPU
    model_ft = model_ft.to(device)
    model_ft.cuda()

• 4、定义优化器

    criterion = nn.CrossEntropyLoss()
    optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)
    

• 5、训练模型

num_epochs = 100
model_ft,arr_acc = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs)

• 6、保存模型

    torch.save(model_ft.state_dict(), './model/my_model.pth')

• 7、定义自己的数据类

def default_loader(path):
    with open(path, 'rb') as f:
        with Image.open(f) as img:
            return img.convert('RGB')

class CustomImageLoader(data.Dataset): # 定义自己的数据类
    ##自定义类型数据输入
    def __init__(self, img_path, txt_path, dataset = '', data_transforms=None, loader = default_loader):
        im_list = []
        im_labels = []
        with open(txt_path, 'r') as files:
            for line in files:
                #/x/y/a.jpg 1
                #/x/y/b.jpg 2
                items = line.split()
                im_list.append(items[0])
                im_labels.append(int(items[1]))
        self.imgs = im_list
        self.labels = im_labels
        self.data_tranforms = data_transforms
        self.loader = loader
        self.dataset = dataset
 
    def __len__(self):
        return len(self.imgs)
 
    def __getitem__(self, item):
        img_name = self.imgs[item]
        label = self.labels[item]
        img = self.loader(img_name)
 
        if self.data_tranforms is not None:
            try:
                img = self.data_tranforms[self.dataset](img)
            except:
                print("Cannot transform image: {}".format(img_name))
        return img, label

数据转换器：根据自己的需要设置，比如下面的 64、50 是根据我的输入图像尺寸决定的。
resnet18中的默认值：
transforms.Resize(256),
transforms.CenterCrop(224),
我的data_tranforms：

data_tranforms={
    'Train':transforms.Compose([
        transforms.RandomResizedCrop(50), # 随机裁剪为不同的大小和宽高比,缩放所为制定的大小
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
    ]),
    'Test':transforms.Compose([
        transforms.Resize(64), # 变换大小
        transforms.CenterCrop(50), # 中心裁剪
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
    ])
}

• 8、定义训练函数

image_datasets = {x : CustomImageLoader('/', # 默认目录为根目录，配搭文件中使用全路径
                                        txt_path=('./data/{0}.txt'.format(x)), # 标签文件
                                        data_transforms=data_tranforms,
                                        dataset=x) for x in ['Train', 'Test']
                  }
 
dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],
                                                 batch_size=batch_size,
                                                 shuffle=True) for x in ['Train', 'Test']}
 
dataset_sizes = {x: len(image_datasets[x]) for x in ['Train', 'Test']} # 数据大小

def train_model(model, crtiation, optimizer,schedular, num_epochs=NUM_EPOCH):
    
    begin_time = time.time()
    best_weights = copy.deepcopy(model.state_dict())#copy the weights from the model
    best_acc = 0.0
    arr_acc = [] # 用于作图
 
    for epoch in range(num_epochs):
        print("-*-" * 20)
        item_acc = []
        for phase in ['Train', 'Test']:
            if phase=='Train':
                schedular.step()
                model.train()
            else:
                model.eval()
            running_loss = 0.0
            running_acc = 0.0
 
            for images, labels in dataloders[phase]:
                images.to(device)
                labels.to(device)
                optimizer.zero_grad()
 
                with torch.set_grad_enabled(phase=='Train'):
                    opt = model(images.cuda())
                    # opt = model(images)
                    _,pred = torch.max(opt,1)
                    labels = labels.cuda()
                    loss = crtiation(opt, labels)
                    if phase=='Train':
                        loss.backward()
                        optimizer.step()
 
                running_loss += loss.item()*images.size(0)
                running_acc += torch.sum(pred==labels)
            epoch_loss = running_loss/dataset_sizes[phase]
            epoch_acc = running_acc.double()/dataset_sizes[phase]
            print('epoch={}, Phase={}, Loss={:.4f}, ACC:{:.4f}'.format(epoch, phase, 
                                                                       epoch_loss, epoch_acc))
            item_acc.append(epoch_acc)
 
            if phase == 'Test' and epoch_acc>best_acc:
                # Upgrade the weights
                best_acc=epoch_acc
                best_weights = copy.deepcopy(model.state_dict())
        arr_acc.append(item_acc)
        
    time_elapes = time.time() - begin_time
    print('Training Complete in{:.0f}m {:0f}s'.format(
        time_elapes // 60, time_elapes % 60
    ))
    print('Best Val ACC: {:}'.format(best_acc))
 
    model.load_state_dict(best_weights) # 保存最好的参数
    return model,arr_acc

完整代码

train.py：

import os
import torch.utils.data as data
import torch
import torch.optim as optim
import torch.nn as nn
from torch.optim import lr_scheduler
from torchvision import datasets, models, transforms
from PIL import Image
import time
import copy
#import pandas as pd
#import matplotlib.pyplot as plt
import numpy as np
#%matplotlib inline

NUM_EPOCH = 100 # 默认迭代次数
batch_size = 64
device = torch.device('cuda:0') # 默认使用 GPU
NUMCLASS = 4 # 类别数

def default_loader(path):
    with open(path, 'rb') as f:
        with Image.open(f) as img:
            return img.convert('RGB')

class CustomImageLoader(data.Dataset): # 定义自己的数据类
    ##自定义类型数据输入
    def __init__(self, img_path, txt_path, dataset = '', data_transforms=None, loader = default_loader):
        im_list = []
        im_labels = []
        with open(txt_path, 'r') as files:
            for line in files:
                #/x/y/a.jpg 1
                #/x/y/b.jpg 2
                items = line.split()
                im_list.append(items[0])
                im_labels.append(int(items[1]))
        self.imgs = im_list
        self.labels = im_labels
        self.data_tranforms = data_transforms
        self.loader = loader
        self.dataset = dataset
 
    def __len__(self):
        return len(self.imgs)
 
    def __getitem__(self, item):
        img_name = self.imgs[item]
        label = self.labels[item]
        img = self.loader(img_name)
 
        if self.data_tranforms is not None:
            try:
                img = self.data_tranforms[self.dataset](img)
            except:
                print("Cannot transform image: {}".format(img_name))
        return img, label
    
data_tranforms={
    'Train':transforms.Compose([
        transforms.RandomResizedCrop(50), # 随机裁剪为不同的大小和宽高比,缩放所为制定的大小
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
    ]),
    'Test':transforms.Compose([
        transforms.Resize(64), # 变换大小
        transforms.CenterCrop(50), # 中心裁剪
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
    ])
}
 
image_datasets = {x : CustomImageLoader('/', # 默认目录为根目录，配搭文件中使用全路径
                                        txt_path=('./data/{0}.txt'.format(x)), # 标签文件
                                        data_transforms=data_tranforms,
                                        dataset=x) for x in ['Train', 'Test']
                  }
 
dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],
                                                 batch_size=batch_size,
                                                 shuffle=True) for x in ['Train', 'Test']}
 
dataset_sizes = {x: len(image_datasets[x]) for x in ['Train', 'Test']} # 数据大小

def train_model(model, crtiation, optimizer,schedular, num_epochs=NUM_EPOCH):
    
    begin_time = time.time()
    best_weights = copy.deepcopy(model.state_dict())#copy the weights from the model
    best_acc = 0.0
    arr_acc = [] # 用于作图
 
    for epoch in range(num_epochs):
        print("-*-" * 20)
        item_acc = []
        for phase in ['Train', 'Test']:
            if phase=='Train':
                schedular.step()
                model.train()
            else:
                model.eval()
            running_loss = 0.0
            running_acc = 0.0
 
            for images, labels in dataloders[phase]:
                images.to(device)
                labels.to(device)
                optimizer.zero_grad()
 
                with torch.set_grad_enabled(phase=='Train'):
                    opt = model(images.cuda())
                    # opt = model(images)
                    _,pred = torch.max(opt,1)
                    labels = labels.cuda()
                    loss = crtiation(opt, labels)
                    if phase=='Train':
                        loss.backward()
                        optimizer.step()
 
                running_loss += loss.item()*images.size(0)
                running_acc += torch.sum(pred==labels)
            epoch_loss = running_loss/dataset_sizes[phase]
            epoch_acc = running_acc.double()/dataset_sizes[phase]
            print('epoch={}, Phase={}, Loss={:.4f}, ACC:{:.4f}'.format(epoch, phase, 
                                                                       epoch_loss, epoch_acc))
            item_acc.append(epoch_acc)
 
            if phase == 'Test' and epoch_acc>best_acc:
                # Upgrade the weights
                best_acc=epoch_acc
                best_weights = copy.deepcopy(model.state_dict())
        arr_acc.append(item_acc)
        
    time_elapes = time.time() - begin_time
    print('Training Complete in{:.0f}m {:0f}s'.format(
        time_elapes // 60, time_elapes % 60
    ))
    print('Best Val ACC: {:}'.format(best_acc))
 
    model.load_state_dict(best_weights) # 保存最好的参数
    return model,arr_acc

if __name__ == '__main__':
    
    model_ft = models.resnet18(pretrained=True)
    num_fits = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层
    model_ft = model_ft.to(device)
    model_ft.cuda()
    criterion = nn.CrossEntropyLoss()
    optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)
    model_ft,arr_acc = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, NUM_EPOCH)
    ## 保存模型 
    torch.save(model_ft.state_dict(), './model/my_model.pth')