Pytorch实战训练第一篇-------猫狗分类
0、数据集准备
本次数据集使用Kaggle上的开源数据集，这是网站链接：猫狗数据集
本次数据集共：857.48MB，分为train和test两个文件夹，其中train中有25000张图片，test中有12500张图片。其中有一点需要注意的是，这个train文件夹中的猫和狗的照片没有分开打包，所以我们需要进一步对train中的图片进行处理。
1、对train中的数据进行分割
- 获取目录下的所有.jpg文件，并转化为一个列表
files = glob.glob(os.path.join(path, "*.jpg"))
- 对图片进行随机排序，以保证对train文件进行划分为训练集和验证集时具有随机性
num_of_images = len(files)
shuffle = np.random.permutation(num_of_images) # 对图片进行随机排序
- 对数据进行划分，前2000张图片为验证集，后23000图片为训练集。并放入对应的目录中，同时在train和valid目录中再分别创建两个目录，分别为cat和dog，将数据分别放入对应的目录中。
# 划分验证集
for i in shuffle[:2000]:
folder = files[i].split('\\')[-1].split('.')[0] # 看属于哪个文件，就创建哪个文件，第一个创建完成之后，其他相同的就直接放进去
image = files[i].split('\\')[-1] # 图片原文件名
print(folder, image)
os.rename(files[i], os.path.join(path, "valid", folder, image))
# 划分训练集
for i in shuffle[2000:]:
folder = files[i].split("\\")[-1].split(".")[0]
image = files[i].split("\\")[-1]
os.rename(files[i], os.path.join(path, "train", folder, image))
- 数据分割源码：
import numpy as np
import os
import glob
# 文件路径
path = "./archive/dc/train/"
# 预处理，把图片按照名称分类
files = glob.glob(os.path.join(path, "*.jpg"))
print("一共{}张图片".format(len(files))) # 一共25000张图片
num_of_images = len(files)
shuffle = np.random.permutation(num_of_images) # 对图片进行随机排序
# 创建目录
os.mkdir(os.path.join(path, "train"))
os.mkdir(os.path.join(path, "valid"))
for t in ["train/", "valid/"]:
for folder in ["dog/", "cat/"]:
os.mkdir(os.path.join(path, t, folder))
# 划分验证集
for i in shuffle[:2000]:
folder = files[i].split('\\')[-1].split('.')[0] # 看属于哪个文件，就创建哪个文件，第一个创建完成之后，其他相同的就直接放进去
image = files[i].split('\\')[-1] # 图片原文件名
print(folder, image)
os.rename(files[i], os.path.join(path, "valid", folder, image))
# 划分训练集
for i in shuffle[2000:]:
folder = files[i].split("\\")[-1].split(".")[0]
image = files[i].split("\\")[-1]
os.rename(files[i], os.path.join(path, "train", folder, image))
2、定义网络模型
- 这个神经网络是我在网上随便找了一个框架进行了搭建，顺便稍微修改了一下。训练之后的模型框架如下图所示。在这顺便给大家推荐一个非常好用的查看神经网络模型的网址：NETRON，该网页的详细说明可以查看这篇文章：NETRON说明。
- 因为最后进行Flatten()的时候，数据大小需要计算，比较麻烦，这里给大家介绍一个小技巧，可以将Flatten()后面的线性层注销，然后用一个数据进行测试，这样就可以直接得到数据的大小。
if __name__ =="__main__":
model=Model()
input=torch.ones((32,3,224,224))
output=model(input)
print(output.shape)
- 神经网络模型搭建源码：
from torch.nn import Sequential,Conv2d,MaxPool2d,Linear,BatchNorm2d,ReLU,Flatten,BatchNorm1d
from torch import nn
import torch
# 搭建神经网络
class Model(nn.Module):
def __init__(self):
super(Model,self).__init__()
self.model=Sequential(
Conv2d(3, 32, (7,7),stride=(2,2),padding=2),
BatchNorm2d(32),
ReLU(),
MaxPool2d(2),
Conv2d(32, 32, 5, padding=2),
BatchNorm2d(32),
ReLU(),
MaxPool2d(2),
Conv2d(32, 64, 5, padding=2),
BatchNorm2d(64),
ReLU(),
MaxPool2d(2),
Flatten(),
Linear(10816, 64), # 易错，展开后长度发生改变
BatchNorm1d(64),
ReLU(),
Linear(64, 2)
)
def forward(self,x):
output=self.model(x)
return output
if __name__ =="__main__":
model=Model()
input=torch.ones((32,3,224,224))
output=model(input)
print(output.shape)
3、训练模型
- 定义训练模型所用设备，我这里使用的是GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("使用{}运行".format(device))
- 数据转化。
因为Pytorch需要传入Tensor类型的数据，所以需要进行相关处理，并且数据集中的图片大小不一，需要进行大小转换，这里还用到的归一化，方便计算，其中的mean和std数值是常用于处理图片数据的，可以理解为默认是这个值。
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
# Image数据常用归一化数值
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
- 数据集准备和加载。
因为数据类型是图片，并且在相对应的目录下，这里使用ImageFolder，直接可以加载为dataset，并且其目录对应下的文件自动打上目录名标签。例如：dog目录下的文件全部为dog标签。这里使用DataLoader加载数据，batch_size设置为32，这个参数根据你电脑配置进行对应的设置，设置过大会导致负载报错。
# 数据集准备 ImageFolder可以加载root下的所有文件，返回一个数据集对象
train = ImageFolder(train_path, data_transform)
valid = ImageFolder(valid_path, data_transform)
print(train.classes) # ['cat', 'dog']
print(train.class_to_idx) # {'cat': 0, 'dog': 1} 顺序取决于文件夹的顺序
# DataLoader加载数据
train_data = DataLoader(train, batch_size=32, shuffle=True, drop_last=True)
valid_data = DataLoader(valid, batch_size=32, shuffle=True, drop_last=True)
- 参数设置、损失函数及优化器设置。
需要注意的是这里的模型实例化、损失函数均需要转化为你所要运行的device状态（我这里是cuda），否则就会报错。
# 模型实例化
model = Model()
model.to(device)
# 参数设置
# 学习率
learning_rate = 1e-2
# 训练次数
total_train_step = 0
# 测试次数
total_valid_step = 0
# 训练轮数
epoch = 20
# 创建损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device)
# 优化器
optimizer = torch.optim.SGD(params=model.parameters(), lr=learning_rate)
- 模型训练及保存最后一轮训练结果，其中进行数据训练时，需要将imgs和targets转换到device状态。
imgs = imgs.to(device)
targets = targets.to(device)
- train_model源码：
import torch
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from model import *
# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("使用{}运行".format(device))
# 数据路径
train_path = "./archive/dc/train/train"
valid_path = "./archive/dc/train/valid"
# 数据转换
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
# Image数据常用归一化数值
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 数据集准备 ImageFolder可以加载root下的所有文件，返回一个数据集对象
train = ImageFolder(train_path, data_transform)
valid = ImageFolder(valid_path, data_transform)
print(train.classes) # ['cat', 'dog']
print(train.class_to_idx) # {'cat': 0, 'dog': 1} 顺序取决于文件夹的顺序
# 查看数据集大小
train_data_size = len(train)
valid_data_size = len(valid)
print("训练集大小为：{}".format(train_data_size))
print("测试集大小为：{}".format(valid_data_size))
# DataLoader加载数据
train_data = DataLoader(train, batch_size=32, shuffle=True, drop_last=True)
valid_data = DataLoader(valid, batch_size=32, shuffle=True, drop_last=True)
# 模型实例化
model = Model()
model.to(device)
# 参数设置
# 学习率
learning_rate = 1e-2
# 训练次数
total_train_step = 0
# 测试次数
total_valid_step = 0
# 训练轮数
epoch = 20
# 创建损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device)
# 优化器
optimizer = torch.optim.SGD(params=model.parameters(), lr=learning_rate)
for i in range(epoch):
print("============第{}轮训练开始==========".format(i + 1))
# 训练开始
model.train()
for data in train_data:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = model(imgs)
# print(imgs)
# print(targets)
# 计算损失值
loss = loss_fn(outputs, targets)
# 优化器优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
# if total_train_step%100==0:
print({"训练次数：{}，Loss：{}".format(total_train_step, loss.item())})
# 验证开始
model.eval()
total_valid_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in valid_data:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = model(imgs)
loss = loss_fn(outputs, targets)
total_valid_loss = total_valid_loss + loss
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy = total_accuracy + accuracy
print("验证集上的Loss为：{}".format(total_valid_loss))
print("验证集上的准确率为 {}".format(total_accuracy / valid_data_size))
# 保存最后一轮的训练结果
if i == epoch - 1:
torch.save(model, "model{}.pth".format(i + 1))
print("模型已保存！")
4、测试模型及将分类后的图片加载到对应的目录下
- 定义键值函数
因为之前训练时 0 表示cat，1表示dog，模型预测最终也只能输出0或1，我需要将分类后的图片加载到对应目录下，所以我需要进行键值转换。
def get_Key(num):
if num == 0:
return "cat"
else:
return "dog"
- 定义数据类型
由于我需要的到文件的路径来进行文件路径的修改，以达到最终放入对应目录的效果，所以我在这定义了一个新的数据类型，返回图片信息和文件路径。
# 定义数据类
class CustomImageFolder():
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.image_paths = os.listdir(root_dir) # 存储所有的文件名
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
img_path = os.path.join(self.root_dir, self.image_paths[idx])
image = Image.open(img_path).convert(("RGB"))
if self.transform:
image = self.transform(image)
return image, img_path
- 测试源码：
import os
from torchvision import transforms
from model import *
from PIL import Image
from torch.utils.data import DataLoader, TensorDataset
def get_Key(num):
if num == 0:
return "cat"
else:
return "dog"
# 定义数据类
class CustomImageFolder():
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.image_paths = os.listdir(root_dir) # 存储所有的文件名
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
img_path = os.path.join(self.root_dir, self.image_paths[idx])
image = Image.open(img_path).convert(("RGB"))
if self.transform:
image = self.transform(image)
return image, img_path
# 数据类型转换
data_trans = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 模型实例化
model = torch.load("./model5.pth", map_location="cuda")
if __name__ == '__main__':
# 定义文件路径
dir_pth = "./archive/dc/test"
# 创建一个数据集
dataset = CustomImageFolder(root_dir=dir_pth, transform=data_trans)
data_size = len(dataset)
print("测试集大小为：{}".format(data_size)) # 测试集大小为：2000
# 加载数据集
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=True, num_workers=4)
# 创建文件夹
os.mkdir(os.path.join(dir_pth, "cat"))
os.mkdir(os.path.join(dir_pth, "dog"))
for i, batch in enumerate(dataloader):
inputs, paths = batch[0].to("cuda"), batch[1] # inputs是图像数据，paths是对应路径
outputs = model(inputs)
for output, path in zip(outputs, paths):
result = output.argmax().item()
class_name = get_Key(result)
if result == 0:
image_name = path.split("\\")[-1]
os.rename(path, os.path.join(dir_pth, "cat", image_name))
else:
image_name = path.split('\\')[-1]
os.rename(path, os.path.join(dir_pth, "dog", image_name))
print("第{}轮已结束".format(i + 1))
5、总结
- 网络模型不够优秀，最终的验证集正确率只有85%左右（没附图主要是因为我忘记截图了。。。）
- 在测试时，将数据进行了分类，最终可以看到对应目录下的图片，由于正确率确实不够高，导致还是会有肉眼可见的分类错误。分了两个目录，还有多余的是因为设置了drop_last=True(还有一个原因是中途调试了一下，导致部分数据被我删了，实际测试数据集大小好像是8500)。
- 光cat这个目录第一页就已经能看到dog了。。。。
- 此次实战中没有对损失率及准确率进行画图，下次可以考虑加上，毕竟看起来6一点。
9万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
