代码收藏家 基于PyTorch的猫狗二分类模型实现史上最高预测准确率
**
基于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，完美。
最后，今天是2023年3月23日，祝亲们2023年健康、富有、开心！
结束~
作者：软件工程唐神
