本文使用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。在训练的时候使用前面预训练模型,这样子加快模型的收敛,而且效果奇佳。