基于torchvision的CV迁移学习

前面我们用过了cifar10，这里因为我们模型的体量更大，他能够理解更加复杂的数据集，所以这里我们就使用更加复杂的数据集叫做cifar100，顾名思义就是它是一个100分类的图像数据集，分类数据更多，复杂度更多。

定义数据集

import torchvision
import torch


#定义数据集
class Dataset(torch.utils.data.Dataset):

    def __init__(self, train):

        #在线加载数据集
        #更多数据集:https://pytorch.org/vision/stable/datasets.html
        self.data = torchvision.datasets.CIFAR100(root='data',
                                                  train=train,
                                                  download=True)

        #更多数据增强:https://pytorch.org/vision/stable/transforms.html
        self.compose = torchvision.transforms.Compose([

            #原本是32*32的,缩放到300*300,这是为了适应预训练模型的习惯,便于它抽取图像特征
            torchvision.transforms.Resize(300),

            #随机左右翻转,这是一种图像增强,很显然,左右翻转不影响图像的分类结果
            torchvision.transforms.RandomHorizontalFlip(p=0.5),

            #图像转矩阵数据,值域是0-1之间
            torchvision.transforms.ToTensor(),

            #让图像的3个通道的数据分别服从3个正态分布,这3分数据是从一个大的数据集上统计得出的
            #投影也是为了适应预训练模型的习惯
            torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                             std=[0.229, 0.224, 0.225]),
        ])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        #取数据
        x, y = self.data[i]

        #应用compose,图像转数据
        x = self.compose(x)

        return x, y


dataset = Dataset(train=True)

x, y = dataset[0]

print(len(dataset), x.shape, y)

这里我们使用了torchvision的加载数据集的方式，它能够在线的加载数据集，那么更多的数据集可以通过上面注释中的连接获得。root='data'，是指将下载的数据集保存到本地磁盘的路径，也就是数据缓存的位置，下面这个参数train它的取值是一个布尔值，指的是要下载的数据集的训练的部分还是测试的部分，这里用到是一个变量因为两部分训练集我们都需要。

compose这个变量是torchvision提供的另外一个功能，就是图像的数据增强，具体的方法也可以通过注释当中的链接查到，下面演示的是常用的几个，第一个是resize也就是图像的缩放，原本是32x32，这里统一缩放到300x300，这是为了适应预训练模型的习惯，便于预训练模型抽取图像的特征，因为我们的预训练模型它训练的时候都是使用300x300的图像来训练的。第二个应用的图像增强就是随机的左右翻转，翻转的概率设置为0.5，但对于cifar100这个数据集来说，左右翻转是不影响图像的分类结果，加上这个数据增强是让我们的数据集更加的丰富。使用ToTensor这个工具类将图像转为矩阵，值域在0到1之间。最后我们对数据进行一个normalize，也就是让我们的图像数据的三个通道分别的服从三个正态分布，三个正态分布它的均值和标准差都写在上面了。

然后就是len和getitem，getitem这个函数每次取一批数据后，然后对这个图像应用我们的compose，这样对我们的图像增强，然后把图像转换为了数据。

定义loader

#每次从loader获取一批数据时回回调,可以在这里做一些数据整理的工作
#这里写的只是个例子,事实上这个回调函数什么也没干..
def collate_fn(data):
    #取数据
    x = [i[0] for i in data]
    y = [i[1] for i in data]

    #比如可以手动转换数据格式
    x = torch.stack(x)
    y = torch.LongTensor(y)

    return x, y


#数据加载器
loader = torch.utils.data.DataLoader(dataset=dataset,
                                     batch_size=8,
                                     shuffle=True,
                                     drop_last=True,
                                     collate_fn=collate_fn)

x, y = next(iter(loader))

print(len(loader), x.shape, y)

loader的代码没有什么可讲的，要提到的只有collate_fn，这个函数是每次从loader取一批数据的时候都会回调的，所以可以在这一个函数里面做一些数据整理的工作。

(6250, torch.Size([8, 3, 300, 300]), tensor([50, 54, 98, 51, 77, 96, 72, 81]))

很显然，x就是8张图像，y就是8个整数，取值是在0到100之间。

迁移学习

一般模型的第一部都是将数据读进去，然后一层层的抽取特征，最后把数据抽取成一个向量后，放到一个全连接的神经网络当中取进行分类，那么对于一个训练好的神经网络模型来说，其中的很多层其实是可以复用的。比方这里的一个模型它是一个回归的结果，然后我又不想要回归了怎么办。很简单，我将最后一层剪掉，然后重新接上三层新的，然后在这三层当中想做分类还是回归就有我自己决定了。也就是说前面的这些层我是不对它进行训练的，或者说这些层基本上训练好了，即使我对它进行重新的训练，难度也会想对的较小。

这种就是迁移学习，它的核心思想就是复用以前训练好的模型，它其中的一些层的参数，尤其是浅层的，因为这些层它是负责图像数据的特征抽取的，在我新的模型当中是可以复用的，因为数据特征抽取这个工作我一样是要做的。

定义模型

按照之前所说的，我们需要一个预训练好的模型，使用torchvision来完成这项工作，在这里面它提供了很多的预训练模型，更多的选择可以去链接中查找，这样通过torchvision加载后， 重新组装模型，而我们也只需要这里面的feature部分，后面给它接一个全连接的输出层，那么我要做的是分类还是回归就可以自己决定了。

class Model(torch.nn.Module):

    def __init__(self):
        super().__init__()

        #加载预训练模型
        #更多模型:https://pytorch.org/vision/stable/models.html#table-of-all-available-classification-weights
        pretrained = torchvision.models.efficientnet_v2_s(
            weights=torchvision.models.EfficientNet_V2_S_Weights.IMAGENET1K_V1)

        #重新组装模型,只要特征抽取部分
        pretrained = torch.nn.Sequential(
            pretrained.features,
            pretrained.avgpool,
            torch.nn.Flatten(start_dim=1),
        )

        #锁定参数,不训练
        for param in pretrained.parameters():
            param.requires_grad_(False)

        pretrained.eval()
        self.pretrained = pretrained

        #线性输出层,这部分是要重新训练的
        self.fc = torch.nn.Sequential(
            torch.nn.Linear(1280, 256),
            torch.nn.ReLU(),
            torch.nn.Linear(256, 256),
            torch.nn.ReLU(),
            torch.nn.Linear(256, 100),
        )

    def forward(self, x):
        #调用预训练模型抽取参数,因为预训练模型是不训练的,所以这里不需要计算梯度
        with torch.no_grad():
            #[8, 3, 300, 300] -> [8, 1280]
            x = self.pretrained(x)

        #计算线性输出
        #[8, 1280] -> [8, 100]
        return self.fc(x)


model = Model()

x = torch.randn(8, 3, 300, 300)

print(model.pretrained(x).shape, model(x).shape)

模型训练

#训练
def train():
    #注意这里的参数列表,只包括要训练的参数即可
    optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)
    loss_fun = torch.nn.CrossEntropyLoss()
    model.fc.train()

    #定义计算设备,优先使用gpu
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model.to(device)

    print('device=', device)

    for i, (x, y) in enumerate(loader):
        #如果使用gpu,数据要搬运到显存里
        x = x.to(device)
        y = y.to(device)

        out = model(x)
        loss = loss_fun(out, y)

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if i % 500 == 0:
            acc = (out.argmax(dim=1) == y).sum().item() / len(y)
            print(i, loss.item(), acc)

    #保存模型,只保存训练的部分即可
    torch.save(model.fc.to('cpu'), 'model/8.model')

定义设备是我们普遍会写的，如果有GPU那么就使用GPU进行运算

测试

@torch.no_grad()
def test():

    #加载保存的模型
    model.fc = torch.load('model/8.model')
    model.fc.eval()

    #加载测试数据集,共10000条数据
    loader_test = torch.utils.data.DataLoader(dataset=Dataset(train=False),
                                              batch_size=8,
                                              shuffle=True,
                                              drop_last=True)

    correct = 0
    total = 0
    for i in range(100):
        x, y = next(iter(loader_test))

        #这里因为数据量不大,使用cpu计算就可以了
        out = model(x).argmax(dim=1)

        correct += (out == y).sum().item()
        total += len(y)

    print(correct / total)

这里加载的是测试数据集，注意这里的train是False，最后得出的正确率为70%，还是比较的高的。