基于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),
    # 将数据转换为tensor
    transforms.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 = data
        inputs,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
        # 梯度清0
        optimizer.zero_grad()
        # 计算梯度
        loss.backward()
        # 修改权值
        optimizer.step()
def test():
    model.eval()
    correct = 0
    for i, data in enumerate(test_loader):
        # 获得数据和对应的标签
        inputs, labels = data
        inputs,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)))    #损失率+准确率为1

    correct = 0
    for i, data in enumerate(train_loader):
        # 获得数据和对应的标签
        inputs, labels = data
        inputs,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, models


model = 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_t


def 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_cn


def 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 summary
        summary(model, input_size=(3, 224, 224), device="cpu")

    model.to(device)
    return model


def process_img(path_img):

    # hard code
    norm_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 --> img
    img_rgb = Image.open(path_img).convert('RGB')

    # img --> tensor
    img_tensor = img_transform(img_rgb, inference_transform)
    img_tensor.unsqueeze_(0)        # chw --> bchw
    img_tensor = img_tensor.to(device)

    return img_tensor, img_rgb


if __name__ == "__main__":
    num_classes=2
    # config
    path_state_dict = os.path.join(BASE_DIR, "model.pth")
    path_img = os.path.join(BASE_DIR,  "data-predict", "dog.jpg")

    # 1/5 load img
    img_tensor, img_rgb = process_img(path_img)

    # 2/5 load model
    model = 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 visualization
    plt.imshow(img_rgb)
    plt.title("predict:{}".format(pred_str))
    plt.text(5, 45, "top {}:{}".format(1, pred_str), bbox=dict(fc='yellow'))
    plt.show()

    演示一下：

     

     多形象，OK，完美。 

最后，今天是2022年的最后一天，2022.12.31 21:00，祝亲们2023年健康、富有、开心！

结束~