本文使用vgg网络实现对猫狗分类。 可以当做图像分类的一个baseline。

一、前期工作

数据：直接到kaggle上下载相应的数据集即可。

1.导入模块

# 数据import torchfrom torchvision.datasets import ImageFolderimport torch.utils.data as Datafrom torchvision import transforms# 模型import torch.nn as nnfrom torchvision.models import vgg16# 损失函数import torch.nn as nn# 优化器import torch.optim as optim# 训练import numpy as npfrom torch.utils.tensorboard import SummaryWriterimport matplotlib.pyplot as plt# 工具from torchsummary import summaryimport gc# 预测from PIL import Imageimport cv2import glob

2.设置硬件环境

如果有GPU则使用GPU并且清空当前GPU显存，否则使用CPU

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')if device == torch.device('cuda'):    gc.collect()    torch.cuda.empty_cache()

二、数据

数据部分主要是构造深度学习模型所需要的“特定数据”，为什么说特定呢？因为模型是很矫情的，一般来说模型对输入图像的尺寸、格式等都有要求，所以我们要按照要求来构造模型所需要的数据集。

# 数据预处理transform = transforms.Compose([    transforms.Resize([224,224]), # 将图像resize到指定大小    transforms.ToTensor(), # 将图像转为tensor    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 对图像的进行标准化（使用ImageNet数据集的均值和标准差）])# train 数据集trainset = ImageFolder(root='./1-cat-dog/train/',transform = transform) # 构造训练集的Datasettrain_loader = Data.DataLoader(trainset,batch_size = 16,shuffle = True,num_workers = 0) # 构造数据集的DataLoader# val 数据集valset = ImageFolder(root = './1-cat-dog/val/',transform = transform) # 与训练集一样，构造验证集的Datasetval_loader = Data.DataLoader(valset,batch_size =2,num_workers = 0) # 构造验证集的DataLoader

# 查看类别与数字的对应关系valset.class_to_idx

{'cat': 0, 'dog': 1}

三、模型

这一部分是模型部分，我们使用pytorch封装好的vgg16模型，并且对其进行微调，因为原始的vgg16模型是预测1000个类别，而我们的任务只有猫和狗两个类别，所以将最后一层的输出调整为2。在训练的时候使用前面预训练模型，这样子加快模型的收敛，而且效果奇佳。

model = vgg16(pretrained = True).to(device)# 看一下原始模型

# 加载预训练模型，并且对模型进行微调weights_path = './vgg16-397923af.pth'model.load_state_dict(torch.load(weights_path,map_location=device))for param in model.parameters():    param.requires_grad = Falsemodel.classifier._modules['6'] = nn.Linear(4096,2)model.to(device)

VGG((features): Sequential((0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(1): ReLU(inplace=True)(2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(3): ReLU(inplace=True)(4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(6): ReLU(inplace=True)(7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(8): ReLU(inplace=True)(9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(11): ReLU(inplace=True)(12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(13): ReLU(inplace=True)(14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(15): ReLU(inplace=True)(16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(18): ReLU(inplace=True)(19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(20): ReLU(inplace=True)(21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(22): ReLU(inplace=True)(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)(24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(25): ReLU(inplace=True)(26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(27): ReLU(inplace=True)(28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(29): ReLU(inplace=True)(30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))(classifier): Sequential((0): Linear(in_features=25088, out_features=4096, bias=True)(1): ReLU(inplace=True)(2): Dropout(p=0.5, inplace=False)(3): Linear(in_features=4096, out_features=4096, bias=True)(4): ReLU(inplace=True)(5): Dropout(p=0.5, inplace=False)(6): Linear(in_features=4096, out_features=2, bias=True))
)

class MY_VGG16(nn.Module):    def __init__(self, init_weights=True):        super(MY_VGG16, self).__init__()

        self.block1 = nn.Sequential(            nn.Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),        )        self.block2 = nn.Sequential(            nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),        )        self.block3 = nn.Sequential(            nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),        )

        self.block4 = nn.Sequential(            nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),        )

        self.block5 = nn.Sequential(            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),            nn.ReLU(),            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),        )

        self.block6 = nn.Sequential(            nn.Linear(in_features=512*7*7, out_features=4096),            nn.ReLU(),            nn.Linear(in_features=4096, out_features=4096),            nn.ReLU(),            nn.Linear(in_features=4096, out_features=2)        )        if init_weights:            self._initialize_weights()

    def _initialize_weights(self):        for m in self.modules():            if isinstance(m, nn.Conv2d):                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')                if m.bias is not None:                    nn.init.constant_(m.bias, 0)            elif isinstance(m, nn.Linear):                nn.init.normal_(m.weight, 0, 0.01)                nn.init.constant_(m.bias, 0)

    def forward(self, x):

        x = self.block1(x)        x = self.block2(x)        x = self.block3(x)        x = self.block4(x)        x = self.block5(x)        x = torch.flatten(x, start_dim=1)        x = self.block6(x)

        return x

model = MY_VGG16(init_weights=True).to(device)

四、损失函数和优化器

这一部分主要是定义损失函数和优化器

# 损失函数和优化器criterion = nn.CrossEntropyLoss()optimizer = optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)

五、训练模型

有了数据、模型、损失函数和优化器之后，我们就可以开始训练模型了。 模型的训练其实就是经过一些列的迭代步骤，让模型学习从“输入”到“输出”的映射函数。

这里我们定义两个函数，train_one_epoch和valid_one_epoch：

• train_one_epoch函数对训练集遍历一遍；
  
• valid_one_epoch函数对验证集遍历一遍；
  

def train_one_epoch(train_loader, model, criterion, optimizer):    model.train()    running_loss = 0.0    metric = 0.0    total = 0.0

    for i, data in enumerate(train_loader):        images, labels = data        images = images.to(device)        labels = labels.to(device)        outs = model(images)

        optimizer.zero_grad()        loss = criterion(outs, labels)        loss.backward()        optimizer.step()

        running_loss += loss.item()        _, preds = torch.max(outs, 1)        metric += (preds == labels).sum().item()        total += labels.shape[0]

    train_loss = running_loss/total    train_accuracy = 100*metric/total # 转成百分比

    return train_loss, train_accuracy    

def valid_one_epoch(valid_loader, model):    model.eval()

    total = 0.0    metric = 0.0

    for i,data in enumerate(valid_loader):        images, labels = data        images = images.to(device)        labels = labels.to(device)

        outs = model(images)        _, preds = torch.max(outs,1)

        total += labels.shape[0]

        metric += (preds == labels).sum().item()

    valid_accuracy = 100*metric/total 

    return valid_accuracy

然后我们就可以开始训练模型了，并且在模型训练完之后保存最终的模型。

epochs = 20train_losses = []train_accuracys = []valid_accuracys = []for epoch in range(epochs):

    train_loss, train_accuracy = train_one_epoch(train_loader= train_loader, model = model, criterion = criterion, optimizer = optimizer)    print(f"epoch:{epoch}, loss:{train_loss}, accuracy:{train_accuracy}")    train_losses.append(train_loss)    train_accuracys.append(train_accuracy)

    valid_accuracy = valid_one_epoch(valid_loader= val_loader, model = model)    print(f"valid accuracy:{valid_accuracy}")    valid_accuracys.append(valid_accuracy)

torch.save(model,'test.pth')

epoch:0, loss:0.020580440697570642, accuracy:87.91666666666667
valid accuracy:100.0
epoch:1, loss:0.004143970409252991, accuracy:99.58333333333333
valid accuracy:100.0
epoch:2, loss:0.0024526457864946375, accuracy:100.0
valid accuracy:100.0
epoch:3, loss:0.0018572715828971316, accuracy:99.79166666666667
valid accuracy:100.0
epoch:4, loss:0.001885829681608205, accuracy:99.79166666666667
valid accuracy:100.0
epoch:5, loss:0.0012724352748288462, accuracy:100.0
valid accuracy:100.0
epoch:6, loss:0.0012308882627015312, accuracy:99.58333333333333
valid accuracy:100.0
epoch:7, loss:0.000938531729237487, accuracy:100.0
valid accuracy:100.0
epoch:8, loss:0.0008737363134666035, accuracy:100.0
valid accuracy:100.0
epoch:9, loss:0.0007390241080429405, accuracy:100.0
valid accuracy:100.0
epoch:10, loss:0.000971730300807394, accuracy:99.79166666666667
valid accuracy:100.0
epoch:11, loss:0.0006812206517982607, accuracy:100.0
valid accuracy:100.0
epoch:12, loss:0.000887786266199934, accuracy:100.0
valid accuracy:100.0
epoch:13, loss:0.0006422517117850173, accuracy:100.0
valid accuracy:100.0
epoch:14, loss:0.0005755920026179714, accuracy:100.0
valid accuracy:100.0
epoch:15, loss:0.0005269992495110879, accuracy:100.0
valid accuracy:100.0
epoch:16, loss:0.0006754447094863281, accuracy:100.0
valid accuracy:100.0
epoch:17, loss:0.0004891647525558559, accuracy:100.0
valid accuracy:100.0
epoch:18, loss:0.0004677536303158073, accuracy:100.0
valid accuracy:100.0
epoch:19, loss:0.00048221670779942843, accuracy:100.0
valid accuracy:100.0

训练完成之后，我们来可视化一下训练过程中的误差变化。

plt.figure(figsize=(13,8))plt.subplot(1,3,1)plt.plot(range(epochs), train_losses)plt.title('Loss')plt.xlabel('Epoch')plt.ylabel('Loss')

plt.subplot(1,3,2)plt.plot(range(epochs),train_accuracys)plt.title('Train Accuracy')plt.xlabel('Epoch')plt.ylabel('Accuracy')

plt.subplot(1,3,3)plt.plot(range(epochs),valid_accuracys)plt.title('Valid Accuracy')plt.xlabel('Epoch')plt.ylabel('Accuracy')

plt.savefig('loss_accuracy_changing_with_epoch.png')plt.show()

六、预测

有了最终的模型，我们来预测一下，先写一个函数封装一下。

classes = list(valset.class_to_idx)model = torch.load('./test.pth')

def predict_one_image(image_path, model, transform, classes):

    test_img = cv2.imread(image_path)    plt.imshow(test_img)    test_img = Image.fromarray(test_img)

    test_img = transform(test_img)    img = test_img.to(device).unsqueeze(0)

    model.eval()    output = model(img)

    _,pred = torch.max(output,1)    pred_class = classes[pred]    print(f'预测结果：{pred_class}')    return pred_class

# 预测训练集中的某张照片predict_one_image(image_path='./1-cat-dog/train/dog/flickr_dog_000003.jpg', model=model, transform=transform, classes=classes)

预测结果：dog'dog'

# 来个恶搞，预测下白等，hhhhapredict_one_image(image_path='./1-cat-dog/OIP-C.jpg', model=model, transform=transform, classes=classes)

预测结果：dog'dog'

# 预测我自己用手机拍的一张照片试试predict_one_image(image_path='./1-cat-dog/gou.jpg', model=model, transform=transform, classes=classes)

预测结果：dog'dog'

# 预测从bing上下载的一张照片predict_one_image(image_path='./1-cat-dog/cat_with_hat.jpg', model=model, transform=transform, classes=classes)

预测结果：cat'cat'

# 预测验证集中的一张照片

predict_one_image(image_path='./1-cat-dog/val/cat/flickr_cat_000003.jpg', model=model, transform=transform, classes=classes)

预测结果：cat'cat'

七、知识点详解

1. VGG16模型

vgg16的网络架构图如上所知，可以看到其总共有6个子块，之所以叫做vgg16是因为其中带有可学习参数的网络层一共有16层。

2. 注意点

• 在使用预训练模型的时候要对模型进行微调，以使得模型适应自己的任务；
  
• 在进行预测时候，也要对图片进行相应的预处理，尤其是图像大小
  

参考

【1】https://www.cnblogs.com/vvlj/p/14141826.html
【2】https://neurohive.io/en/popular-networks/vgg16/