使用 pytorch训练自己的图片分类模型

如何自己训练一个图片分类模型，如果一切从头开始，对于一般公司或个人基本是难以实现的。其实，我们可以利用一个现有的图片分类模型，加上新的分类，这种方式叫做迁移学习，就是把现有的模式知识，转移到新的模型。Pytorch 官网提供已经训练好的模型，可以在此基础上训练自己的模型。我们用的模型是 VGG 分类模型，首先，先运行一个已经训练好的模型可做 1000 个分类。

安装依赖

# 去官网根据系统进行下载
pip3 install torch torchvision torchaudio
pip3 install tqdm

现有模型进行图片识别

可以去百度上下载一个狗或者鸟的图片，运行下面的程序进行识别。

# 导入软件包
import numpy as np
import json
from PIL import Image

import torch
import torchvision
from torchvision import models, transforms

#生成VGG-16模型的实例
use_pretrained = True  # 使用已经训练好的参数
net = models.vgg16(pretrained=use_pretrained)
net.eval()  # 设置为推测模式

# 对输入图片进行预处理的类
class BaseTransform():
    """
    调整图片的尺寸，并对颜色进行规范化。

    Attributes
    ----------
    resize : int
       指定调整尺寸后图片的大小
    mean : (R, G, B)
       各个颜色通道的平均值
    std : (R, G, B)
       各个颜色通道的标准偏差
    """

    def __init__(self, resize, mean, std):
        self.base_transform = transforms.Compose([
            transforms.Resize(resize),  #将较短边的长度作为resize的大小
            transforms.CenterCrop(resize),  #从图片中央截取resize × resize大小的区域
            transforms.ToTensor(),  #转换为Torch张量
            transforms.Normalize(mean, std)  #颜色信息的正规化
        ])

    def __call__(self, img):
        return self.base_transform(img)

# 根据输出结果对标签进行预测的后处理类
class ILSVRCPredictor():
    """
    根据ILSVRC数据，从模型的输出结果计算出分类标签

    Attributes
    ----------
    class_index : dictionary
           将类的index与标签名关联起来的字典型变量
    """

    def __init__(self, class_index):
        self.class_index = class_index

    def predict_max(self, out):
        """
        获得概率最大的ILSVRC分类标签名

        Parameters
        ----------
        out : torch.Size([1, 1000])
            从Net中输出结果

        Returns
        -------
        predicted_label_name : str
            预测概率最高的分类标签的名称
        """
        maxid = np.argmax(out.detach().numpy())
        predicted_label_name = self.class_index[str(maxid)][1]

        return predicted_label_name
# 载入ILSVRC的标签信息，并生成字典型变量
ILSVRC_class_index = json.load(open('./data/imagenet_class_index.json', 'r'))

# 生成ILSVRCPredictor的实例
predictor = ILSVRCPredictor(ILSVRC_class_index)

# 读取输入的图像
image_file_path = './data/jww2.webp'
img = Image.open(image_file_path)  # [ 高度 ][ 宽度 ][ 颜色RGB]

# 完成预处理后，添加批次尺寸的维度
resize = 224
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
transform = BaseTransform(resize, mean, std)  #创建预处理类
img_transformed = transform(img)  # torch.Size([3, 224, 224])
inputs = img_transformed.unsqueeze_(0)  # torch.Size([1, 3, 224, 224])

# 输入数据到模型中，并将模型的输出转换为标签
out = net(inputs)  # torch.Size([1, 1000])
result = predictor.predict_max(out)

# 输出预测结果
print("输入图像的预测结果：", result)

我识别的是一只吉娃娃的图片，结果正确，Chihuahua。

现有的模型已经可以正常工作了，下面就是添加新的分类了，这里使用了蚂蚁和蜜蜂。把 1000 个分类改为了 2个分类。
net.classifier[6] = nn.Linear(in_features=4096, out_features=2)

import glob
import os.path as osp
import random
import numpy as np
import json
from PIL import Image
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision
from torchvision import models, transforms

torch.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)




class ImageTransform():
    """
    图像的预处理类。训练时和推测时采用不同的处理方式
    对图像的大小进行调整，并将颜色信息标准化
    训练时采用 RandomResizedCrop 和 RandomHorizontalFlip 进行数据增强处理


    Attributes
    ----------
    resize : int
       指定调整后图像的尺寸
    mean : (R, G, B)
        各个颜色通道的平均值
    std : (R, G, B)
        各个颜色通道的标准偏差
    """

    def __init__(self, resize, mean, std):
        self.data_transform = {
            'train': transforms.Compose([
                transforms.RandomResizedCrop(
                    resize, scale=(0.5, 1.0)), #数据增强处理
                transforms.RandomHorizontalFlip(),  #数据增强处理
                transforms.ToTensor(),  # 转换为张量
                transforms.Normalize(mean, std)  # 归一化
            ]),
            'val': transforms.Compose([
                transforms.Resize(resize),  #调整大小
                transforms.CenterCrop(resize),  #从图像中央截取resize×resize大小的区域
                transforms.ToTensor(), #转换为张量
                transforms.Normalize(mean, std)  #归一化
            ])
        }

    def __call__(self, img, phase='train'):
        """
        Parameters
        ----------
        phase : 'train' or 'val'
            指定预处理所使用的模式
        """
        return self.data_transform[phase](img)

#  创建用于保存蚂蚁和蜜蜂的图片的文件路径的列表变量


def make_datapath_list(phase="train"):
    """
    创建用于保存数据路径的列表

    Parameters
    ----------
    phase : 'train' or 'val'
        指定是训练数据还是验证数据

    Returns
    -------
    path_list : list
       保存了数据路径的列表
    """

    rootpath = "./data/hymenoptera_data/"
    target_path = osp.join(rootpath+phase+'/**/*.jpg')
    print(target_path)

    path_list = []  #  保存到这里

    #  使用 glob 取得包括示例目录的文件路径
    for path in glob.glob(target_path):
        path_list.append(path)

    return path_list


class HymenopteraDataset(data.Dataset):
    """
    蚂蚁和蜜蜂图片的Dataset类，继承自PyTorch的Dataset类

    Attributes
    ----------
    file_list : 列表
        列表中保存了图片路径
    transform : object
        预处理类的实例
    phase : 'train' or 'test'
        指定是学习还是验证
    """

    def __init__(self, file_list, transform=None, phase='train'):
        self.file_list = file_list  # 文件路径列表
        self.transform = transform  # 预处理类的实例
        self.phase = phase  # 指定是train 还是val

    def __len__(self):
        '''返回图片张数'''
        return len(self.file_list)

    def __getitem__(self, index):
        '''
        获取预处理完毕的图片的张量数据和标签
        '''

        #载入第index张图片
        img_path = self.file_list[index]
        img = Image.open(img_path) #[高度][宽度][颜色RGB]

        #对图片进行预处理
        img_transformed = self.transform(
            img, self.phase)  # torch.Size([3, 224, 224])

        #从文件名中抽取图片的标签
        if self.phase == "train":
            label = img_path[30:34]
        elif self.phase == "val":
            label = img_path[28:32]

      #将标签转换为数字
        if label == "ants":
            label = 0
        elif label == "bees":
            label = 1

        return img_transformed, label

#  执行
train_list = make_datapath_list(phase="train")
val_list = make_datapath_list(phase="val")

#执行
size = 224
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
train_dataset = HymenopteraDataset(
    file_list=train_list, transform=ImageTransform(size, mean, std), phase='train')

val_dataset = HymenopteraDataset(
    file_list=val_list, transform=ImageTransform(size, mean, std), phase='val')

#指定小批次尺寸
batch_size = 32

#创建DataLoader
train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True)

val_dataloader = torch.utils.data.DataLoader(
    val_dataset, batch_size=batch_size, shuffle=False)

#集中到字典变量中
dataloaders_dict = {"train": train_dataloader, "val": val_dataloader}

#确认执行结果
batch_iterator = iter(dataloaders_dict["train"])  #转换成迭代器
inputs, labels = next(
    batch_iterator) #取出第一个元素

# 载入已经学习完毕的VGG−16模型
#创建VGG−16模型的实例
use_pretrained = True #指定使用已经训练好的参数
net = models.vgg16(pretrained=use_pretrained)

#指定使用已经训练好的参数
net.classifier[6] = nn.Linear(in_features=4096, out_features=2)

#设定为训练模式
net.train()

print('网络设置完毕 ：载入已经学习完毕的权重，并设置为训练模式')

# #设置损失函数
criterion = nn.CrossEntropyLoss()

params_to_update = []

#需要学习的参数名称
update_param_names = ["classifier.6.weight", "classifier.6.bias"]

#除了需要学习的那些参数外，其他参数设置为不进行梯度计算，禁止更新
for name, param in net.named_parameters():
    if name in update_param_names:
        param.requires_grad = True
        params_to_update.append(param)
        print(name)
    else:
        param.requires_grad = False

optimizer = optim.SGD(params=params_to_update, lr=0.001, momentum=0.9)


def train_model(net, dataloaders_dict, criterion, optimizer, num_epochs):

    #epoch循环
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-------------')

        # 每个epoch中的学习和验证循环
        for phase in ['train', 'val']:
            if phase == 'train':
                net.train()  #将模式设置为训练模式
            else:
                net.eval()   #将模式设置为验证模式

            epoch_loss = 0.0  #epoch的合计损失
            epoch_corrects = 0 #epoch的正确答案数量

            #为了确认训练前的验证能力，省略epoch=0时的训练
            if (epoch == 0) and (phase == 'train'):
                continue

            #载入数据并切取出小批次的循环
            for inputs, labels in tqdm(dataloaders_dict[phase]):

                #初始化optimizer
                optimizer.zero_grad()

                #计算正向传播（forward）
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = net(inputs)
                    loss = criterion(outputs, labels) #计算损失
                    _, preds = torch.max(outputs, 1)  #预测标签
                    
  
                    ##训练时的反向传播
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                    #计算迭代的结果
                    # 计算迭代的结果
                    epoch_loss += loss.item() * inputs.size(0)  
                    # 更新正确答案数量的总和
                    epoch_corrects += torch.sum(preds == labels.data)

            #显示每个epoch的loss和正解率
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double(
            ) / len(dataloaders_dict[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

num_epochs=2
train_model(net, dataloaders_dict, criterion, optimizer, num_epochs=num_epochs)

在这里插入图片描述
通过运行结果可以看到，首次没有训练直接在原始模型进行测试，正确率 33%，第二轮，经过 8 次迭代学习，正确率提高到 72%，这里比较奇怪的是验证集的正确率更高。原因是训练集做了数据增广，有些图片是变形的，所以识别起来更加困难。