基于Pytorch的猫狗分类

无偿分享~
猫狗二分类文件下载地址
在下一章说

猫狗分类这个真是困扰我好几天，找了好多资料都是以TensorFlow的猫狗分类，但我们要求的是以pytorch的猫狗分类。刚开始我找到了也运行成功了觉得可以了，最后看了一眼实践要求傻眼了，老师要pytorch，我却弄了TensorFlow，主要是当时不懂觉得这俩一样，之后的寻找中慢慢发现这俩都是环境，不一样。之后就又找，找了好几天，可辛苦了，网上大部分以都是TensorFlow的猫狗分类，很少有pytorch。不过，之后的之后弄出来了。这个过程学到了很多东西，写个文章记录一下。

我的用了GPU，你们都试试安装GPU吧，很简单，我以前以为很难，其实不是。安装GPU你要记住，先安装cuda，用cuda的版本去安装pytorch，之后~~~So easy！不懂得也可以问我。

好了，正题：

我的软件：pycharm专业版，原来是社区版，但是他创建文件那里选项太少，直接转战专业版yyds,我这个是网上找的破解版（这几天探索过程我发现Visual Studio Code也可以运行python，当时觉得挺好，就弄了弄那个，可是之后运行不了好像什么tensboard版本不对，我索性放弃直接用pycharm这个专业版的）

pytorch有两个，一个CPU，一个GPU（原来一个电脑可以有两个pytorch，命名不一样就行）

pytorch中的python：3.7，torch：1.2.0。

cuda:10.0(有点低，但起码我版本都匹配，能用)

下面上代码，代码里面有讲解，自己去悟悟吧，加油，发现问题解决问题确实是学习的好方法（因为哥哥我受益匪浅）

先上目录

其中data-predict是存放预测图片的，我就放了两张~

一：数据准备

网上用的都是当年猫狗大战比赛的数据集，解压之后如下图所示

我们都是用train里面的图片，有25000张，我觉得有点多，就创建了一个小点的数据集，就是上面目录中的Smalldata。你们去网上找这个25000的数据集都有百度网盘下载，（别去原网站下载，麻烦，还不一定成功，连我这么一个坚持为王的人都放弃了~）

data.py

import os, shutil
# 下载的kaggle数据集路径
original_dataset_dir = '/pythonProject3/猫狗分类/Bigdata'
# 新的小数据集放置路径
base_dir = '/pythonProject3/猫狗分类/Smalldata'
os.mkdir(base_dir)
train_dir = os.path.join(base_dir, 'train')
os.mkdir(train_dir)
test_dir = os.path.join(base_dir, 'test')
os.mkdir(test_dir)train_cats_dir = os.path.join(train_dir, 'cats')
os.mkdir(train_cats_dir)
train_dogs_dir = os.path.join(train_dir, 'dogs')
os.mkdir(train_dogs_dir)
test_cats_dir = os.path.join(test_dir, 'cats')
os.mkdir(test_cats_dir)
test_dogs_dir = os.path.join(test_dir, 'dogs')
os.mkdir(test_dogs_dir)fnames = ['cat.{}.jpg'.format(i) for i in range(200)]
for fname in fnames:src = os.path.join(original_dataset_dir, fname)dst = os.path.join(train_cats_dir, fname)shutil.copyfile(src, dst)fnames = ['cat.{}.jpg'.format(i) for i in range(300, 400)]
for fname in fnames:src = os.path.join(original_dataset_dir, fname)dst = os.path.join(test_cats_dir, fname)shutil.copyfile(src, dst)fnames = ['dog.{}.jpg'.format(i) for i in range(200)]
for fname in fnames:src = os.path.join(original_dataset_dir, fname)dst = os.path.join(train_dogs_dir, fname)shutil.copyfile(src, dst)fnames = ['dog.{}.jpg'.format(i) for i in range(300, 400)]
for fname in fnames:src = os.path.join(original_dataset_dir, fname)dst = os.path.join(test_dogs_dir, fname)shutil.copyfile(src, dst)print('total training cat images:', len(os.listdir(train_cats_dir)))
print('total training dog images:', len(os.listdir(train_dogs_dir)))
print('total test cat images:', len(os.listdir(test_cats_dir)))
print('total test dog images:', len(os.listdir(test_dogs_dir)))
演示一下：言多必失，接着来~

二：训练和模型创建，对了里面还有读取数据

train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")  #判断GPU是否可用
# 数据预处理 数据增强
transform = transforms.Compose([# 对图像进行随机的裁剪crop以后再resize成固定大小（224*224）transforms.RandomResizedCrop(224),# 随机旋转20度（顺时针和逆时针）transforms.RandomRotation(20),# 随机水平翻转transforms.RandomHorizontalFlip(p=0.5),# 将数据转换为tensortransforms.ToTensor()
])# 读取数据
root = 'Smalldata'   #root是数据集目录
# 获取数据的路径，使用transform增强变化
train_dataset = datasets.ImageFolder(root + '/train', transform)
test_dataset = datasets.ImageFolder(root + '/test', transform)
# 导入数据
# 每个批次8个数据，打乱
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, shuffle=True)
# 类别名称
classes = train_dataset.classes
# 类别编号
classes_index = train_dataset.class_to_idx
print("类别名称",classes)
print("类别编号",classes_index)
# models.下有很多pytorch提供的训练好的模型
model = models.vgg16(pretrained=True)
# 我们主要是想调用vgg16的卷积层，全连接层自己定义，覆盖掉原来的
# 如果想只训练模型的全连接层（不想则注释掉这个for）
for param in model.parameters():param.requires_grad = False
# 构建新的全连接层
# 25088：卷阶层输入的是25088个神经元，中间100是自己定义的，输出类别数量2
model.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 100),torch.nn.ReLU(),torch.nn.Dropout(p=0.5),torch.nn.Linear(100, 2)# 这里可以加softmax也可以不加)
model=model.to(device)      #将模型发送到GPU上
print("使用GPU:",next(model.parameters()).device)  # 输出：cuda:0
LR = 0.0001
# 定义代价函数
entropy_loss = nn.CrossEntropyLoss()   #损失函数
# 定义优化器
optimizer = optim.SGD(model.parameters(), LR, momentum=0.9)
print("开始训练~")
def train():model.train()for i, data in enumerate(train_loader):# 获得数据和对应的标签inputs, labels = datainputs,labels=inputs.to(device),labels.to(device)  #将数据发送到GPU上# 获得模型预测结果，（64，10）out = model(inputs)# 交叉熵代价函数out(batch,C),labels(batch)loss = entropy_loss(out, labels).to(device)  #别忘了损失函数也要发到GPU# 梯度清0optimizer.zero_grad()# 计算梯度loss.backward()# 修改权值optimizer.step()
def test():model.eval()correct = 0for i, data in enumerate(test_loader):# 获得数据和对应的标签inputs, labels = datainputs,labels=inputs.to(device),labels.to(device)# 获得模型预测结果out = model(inputs)# 获得最大值，以及最大值所在的位置_, predicted = torch.max(out, 1)# 预测正确的数量correct += (predicted == labels).sum()print("Test acc: {:.2f}".format(correct.item() / len(test_dataset)))print("Test loss:{:.2f}".format(1-correct.item() / len(test_dataset)))    #损失率+准确率为1correct = 0for i, data in enumerate(train_loader):# 获得数据和对应的标签inputs, labels = datainputs,labels=inputs.to(device),labels.to(device)# 获得模型预测结果out = model(inputs)# 获得最大值，以及最大值所在的位置_, predicted = torch.max(out, 1)# 预测正确的数量correct += (predicted == labels).sum()print("Train acc: {:.2f}".format(correct.item() / len(train_dataset)))print("Train loss:{:.2f}".format(1-correct.item() / len(train_dataset)))
for epoch in range(0,10):print('epoch:', epoch)train()test()
torch.save(model.state_dict(), 'model.pth')
print("~结束训练")
演示一下：

因为刚开始这个我弄得CPU没GPU，这个我改了很多次，用了一点Vgg16,属于CNN，就那个放到GPU上跑的代码我弄了一天才弄好，原来是这个：model=model.to(device) 弄错了，第二个model不是VGG16,可是当时就是不行，哎，浪费了哥一天时间。

三：预测（随便取一张猫狗图片可以识别出来是cat还是dog）

predict.py

import torch
import numpy as np
from PIL import Image
from torchvision import transforms, modelsmodel = models.vgg16(pretrained=True)
# 构建新的全连接层
model.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 100),torch.nn.ReLU(),torch.nn.Dropout(p=0.5),torch.nn.Linear(100, 2))# 载入训练好的模型，里面保存的是模型的参数
model.load_state_dict(torch.load('model.pth'))# 预测模式
model.eval()label = np.array(['cat', 'dog'])# 数据预处理
transform = transforms.Compose([transforms.Resize(224),transforms.ToTensor()
])def predict(image_path):# 打开图片img = Image.open(image_path)# 数据处理，再增加一个维度，在第0维度增加1维，变成一个4维的数据（原先是3维，宽高和维度3）img = transform(img).unsqueeze(0)# 预测得到结果outputs = model(img)# 1表示第1个维度（有2种可能的值，猫0狗1），获得最大值所在位置（猫和苟哪一个可能性更大），第0个维度是每个批次的图片数量（1）_, predicted = torch.max(outputs, 1)# 转化为类别名称print(label[predicted.item()])predict('data-predict/cat.jpg')
predict('data-predict/dog.jpg')
演示一下：没图片，那我那个matplotlib不是白学了，绝对不行，整上~（可能我有一点点强迫症吧）

   我始终觉得这个不够完美验证出来不显示图片，所以就查找资料升级了一下。

四：升级版预测

升级版predict.py

import os
os.environ['NLS_LANG'] = 'SIMPLIFIED CHINESE_CHINA.UTF8'
import time
import json
import torch
import torchvision.transforms as transforms
from PIL import Image
from matplotlib import pyplot as plt
import torchvision.models as models
import   torchsummary
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")def img_transform(img_rgb, transform=None):"""将数据转换为模型读取的形式:param img_rgb: PIL Image:param transform: torchvision.transform:return: tensor"""if transform is None:raise ValueError("找不到transform！必须有transform对img进行处理")img_t = transform(img_rgb)return img_tdef load_class_names(p_clsnames, p_clsnames_cn):"""加载标签名:param p_clsnames::param p_clsnames_cn::return:"""with open(p_clsnames, "r") as f:class_names = json.load(f)with open(p_clsnames_cn, encoding='UTF-8') as f:  # 设置文件对象class_names_cn = f.readlines()return class_names, class_names_cndef get_model(path_state_dict, num_classes, vis_model=False):"""创建模型，加载参数:param path_state_dict::return:"""model = models.vgg16(num_classes=num_classes)model.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 100),torch.nn.ReLU(),torch.nn.Dropout(p=0.5),torch.nn.Linear(100, 2))pretrained_state_dict = torch.load(path_state_dict)model.load_state_dict(pretrained_state_dict)model.eval()if vis_model:from torchsummary import summarysummary(model, input_size=(3, 224, 224), device="cpu")model.to(device)return modeldef process_img(path_img):# hard codenorm_mean = [0.485, 0.456, 0.406]norm_std = [0.229, 0.224, 0.225]inference_transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop((224, 224)),transforms.ToTensor(),transforms.Normalize(norm_mean, norm_std),])# path --> imgimg_rgb = Image.open(path_img).convert('RGB')# img --> tensorimg_tensor = img_transform(img_rgb, inference_transform)img_tensor.unsqueeze_(0)        # chw --> bchwimg_tensor = img_tensor.to(device)return img_tensor, img_rgbif __name__ == "__main__":num_classes=2# configpath_state_dict = os.path.join(BASE_DIR, "model.pth")path_img = os.path.join(BASE_DIR,  "data-predict", "dog.jpg")# 1/5 load imgimg_tensor, img_rgb = process_img(path_img)# 2/5 load modelmodel = get_model(path_state_dict,num_classes, True)with torch.no_grad():time_tic = time.time()outputs = model(img_tensor)time_toc = time.time()# 4/5 index to class names_, pred_int = torch.max(outputs.data, 1)_, top1_idx = torch.topk(outputs.data, 1, dim=1)#pred_idx = int(pred_int.cpu().numpy())if pred_idx == 0:pred_str= str("cat")print("img: {} is: {}".format(os.path.basename(path_img), pred_str))else:pred_str = str("dog")print("img: {} is: {}".format(os.path.basename(path_img), pred_str))print("time consuming:{:.2f}s".format(time_toc - time_tic))# 5/5 visualizationplt.imshow(img_rgb)plt.title("predict:{}".format(pred_str))plt.text(5, 45, "top {}:{}".format(1, pred_str), bbox=dict(fc='yellow'))plt.show()

演示一下：