深度学习实现猫狗分类

本文使用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.485, 0.456, 0.406], [0.229, 0.224, 0.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。在训练的时候使用前面预训练模型，这样子加快模型的收敛，而且效果奇佳。