卷积神经网络|迁移学习-猫狗分类完整代码实现

还记得这篇文章吗？迁移学习|代码实现

在这篇文章中，我们知道了在构建模型时，可以借助一些非常有名的模型，这些模型在ImageNet数据集上早已经得到了检验。

同时torchvision模块也提供了预训练好的模型。我们只需稍作修改，便可运用到自己的实际任务中！

我们仍然按照这个步骤开始我们的模型的训练

• 准备一个可迭代的数据集

• 定义一个神经网络

• 将数据集输入到神经网络进行处理

• 计算损失

• 通过梯度下降算法更新参数

import torch import torchvisionimport torchvision.transforms as transformsimport torch.nn as nnimport torch.optim as optimimport matplotlib.pyplot as pltfrom torchvision import models

数据集准备

cifar10_train = torchvision.datasets.CIFAR10(    root = 'cifar10/',    train = True,    download = True)cifar10_test=torchvision.datasets.CIFAR10(    root = 'cifar10/',    train = False,    download = True)
transform = transforms.Compose([        transforms.ToTensor(),        transforms.Resize((224,224))    ])

cifar2_train=[(transform(img),[3,5].index(label)) for img,label in cifar10_train if label in [3,5]]
cifar2_test=[(transform(img),[3,5].index(label)) for img,label in cifar10_test if label in [3,5]]
train_loader = torch.utils.data.DataLoader(cifar2_train, batch_size=64,shuffle=True)test_loader = torch.utils.data.DataLoader(cifar2_test, batch_size=64,shuffle=True)

数据集使用CIFAR-10数据集中的猫和狗。

CIFAR-10数据集类别

种类       标签

• plane       0

• car           1

• bird         2

• cat           3

• deer         4

• dog          5

• frog         6

• horse       7

• ship         8

• truck        9

可以看到其中cat和dog的标签分别为3和5

借助：

[3,5].index(label)

我们可以将cat标签变为0，dog标签变为1，从而回到二分类问题。

举个例子：

>>> [3,5].index(3)0>>> [3,5].index(5)1

定义模型

参考这篇文章：迁移学习|代码实现

#网络搭建network=models.resnet18(pretrained=True)
for param in network.parameters():    param.requires_grad=False
network.fc=nn.Linear(512,2)#损失函数criterion=nn.CrossEntropyLoss()#优化器optimizer=optim.SGD(network.fc.parameters(),lr=0.01,momentum=0.9)
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")network=network.to(device)

训练模型：

for epoch in range(10):    total_loss = 0    total_correct = 0    for batch in train_loader:   # Get batch        images, labels =batch        images=images.to(device)        labels=labels.to(device)                    optimizer.zero_grad()  #告诉优化器把梯度属性中权重的梯度归零，否则pytorch会累积梯度        preds = network(images)        loss = criterion(preds, labels)        loss.backward()        optimizer.step()                total_loss += loss.item()        _,prelabels=torch.max(preds,dim=1)        total_correct += int((prelabels==labels).sum())    accuracy = total_correct/len(cifar2_train)    print("Epoch:%d  ,  Loss:%f  , Accuracy:%f "%(epoch,total_loss,accuracy))

• Epoch:0  ,  Loss:78.549439  , Accuracy:0.788900

• Epoch:1  ,  Loss:77.828066  , Accuracy:0.801500

• Epoch:2  ,  Loss:66.151785  , Accuracy:0.828100

• Epoch:3  ,  Loss:76.204446  , Accuracy:0.816800

• Epoch:4  ,  Loss:68.886606  , Accuracy:0.828100

• Epoch:5  ,  Loss:71.129405  , Accuracy:0.821200

• Epoch:6  ,  Loss:66.096364  , Accuracy:0.829900

• Epoch:7  ,  Loss:65.504227  , Accuracy:0.827700

• Epoch:8  ,  Loss:76.303878  , Accuracy:0.817100

• Epoch:9  ,  Loss:70.546953  , Accuracy:0.820700

测试模型：

correct=0total=0network.eval()with torch.no_grad():    for batch in test_loader:        imgs,labels=batch        imgs=imgs.cuda()        labels=labels.cuda()                preds=network(imgs)        _,prelabels=torch.max(preds,dim=1)        #print(prelabels.size())        total=total+labels.size(0)        correct=correct+int((prelabels==labels).sum())    #print(total)    accuracy=correct/total    print("Accuracy: ",accuracy)

Accuracy:  0.8025

这里使用的预训练模型是resnet18,我们也可以使用VGG16模型，同时记得改变最后一个全连接层的输出参数，使得其满足我们自己的任务。

除了预训练模型之外，我们还可以对一些超参数进行调整，使最后的效果变得更好！