深度学习实现猫狗分类

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

一、前期工作

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

1.导入模块

# 数据
import torch
from torchvision.datasets import ImageFolder
import torch.utils.data as Data
from torchvision import transforms
# 模型
import torch.nn as nn
from torchvision.models import vgg16
# 损失函数
import torch.nn as nn
# 优化器
import torch.optim as optim
# 训练
import numpy as np
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
# 工具
from torchsummary import summary
import gc
# 预测
from PIL import Image
import cv2
import 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.4850.4560.406], [0.2290.2240.225]) # 对图像的进行标准化(使用ImageNet数据集的均值和标准差)
])
# train 数据集
trainset = ImageFolder(root='./1-cat-dog/train/',transform = transform) # 构造训练集的Dataset
train_loader = Data.DataLoader(trainset,batch_size = 16,shuffle = True,num_workers = 0# 构造数据集的DataLoader
# val 数据集
valset = ImageFolder(root = './1-cat-dog/val/',transform = transform) # 与训练集一样,构造验证集的Dataset
val_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 = False
model.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(364, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(6464, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(22), stride=(22)),
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(64128, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(128128, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(22), stride=(22)),
        )
        self.block3 = nn.Sequential(
            nn.Conv2d(128256, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(256256, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(256256, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(22), stride=(22)),
        )

        self.block4 = nn.Sequential(
            nn.Conv2d(256512, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(512512, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(512512, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(22), stride=(22)),
        )

        self.block5 = nn.Sequential(
            nn.Conv2d(512512, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(512512, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.Conv2d(512512, kernel_size=(33), stride=(11), padding=(11)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(22), stride=(22)),
        )

        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, 00.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 = 20
train_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()
深度学习实现猫狗分类_第1张图片 png

六、预测

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

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'
深度学习实现猫狗分类_第2张图片 png
# 来个恶搞,预测下白等,hhhha
predict_one_image(image_path='./1-cat-dog/OIP-C.jpg', model=model, transform=transform, classes=classes)
预测结果:dog





'dog'
深度学习实现猫狗分类_第3张图片 png
# 预测我自己用手机拍的一张照片试试
predict_one_image(image_path='./1-cat-dog/gou.jpg', model=model, transform=transform, classes=classes)
预测结果:dog





'dog'
深度学习实现猫狗分类_第4张图片 png
# 预测从bing上下载的一张照片
predict_one_image(image_path='./1-cat-dog/cat_with_hat.jpg', model=model, transform=transform, classes=classes)
预测结果:cat





'cat'
深度学习实现猫狗分类_第5张图片 png
# 预测验证集中的一张照片

predict_one_image(image_path='./1-cat-dog/val/cat/flickr_cat_000003.jpg', model=model, transform=transform, classes=classes)
预测结果:cat





'cat'
深度学习实现猫狗分类_第6张图片 png

七、知识点详解

1. VGG16模型

深度学习实现猫狗分类_第7张图片 vgg-16

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

2. 注意点

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

参考

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

本文由 mdnice 多平台发布

你可能感兴趣的:(深度学习)