嘿~全流程带你基于Pytorch手撸图片分类“框架“--HuClassify
Posted Huterox
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了嘿~全流程带你基于Pytorch手撸图片分类“框架“--HuClassify相关的知识，希望对你有一定的参考价值。
文章目录
前言
鸽了两天，从星期二晚上就开始说要发布这篇文章，昨天还发布了一下关于这玩意的文章，今天早上果断删除，因为完整版来喽。好吧，其实看标题有点那啥了，其实我这次做的还是一个非常简单的小dome，一个非常小的分类dome，如果一定要说是框架的话，那么只能叫V0.1版本。使用的神经网络模型也算是非常简单的LeNet。 不过这个玩意我是仿造YoloV5的项目结构来写的，做一个快速了解嘛。
由于整个项目不复杂，而且项目也不是很完善，所以暂时不会考虑上传gayHub。如果要的话，评论区或者私信@我即可。
OK，那么咱们开始吧。
使用
先说说咱们的这个玩意长啥样，怎么式样吧。然后咱们再来说说具体的功能模块是如何实现的。
项目结构
和那个YOLO的项目长得是有点像哈。
那么作为用户，想要愉快玩耍这个项目，那么当然要做的就只有三个地方。
我们先一步一步说。
训练过程
准备数据集与配置
首先第一步当然是准备咱们的数据集啦。
这个数据集很简单，人家YOLO 有 YOLO 数据集，那么我们这里就有HU数据集。
这个数据集准备相当简单。
之后打开咱们的data的配置文件
进入训练
接下来移步到我们的训练文件。train.py
我们直接先找到咱们的超参数设置，这个乍一看和yolo还是有几分相像的。
我这里准备了一个数据集，就是分类数字,1 和 100。
你也可以准备多一点，这个无所谓。
训练显示
在你的参数都设置好了之后，那么你就可以开始你的训练了。
最后，你的结果会在这里
我们可以看到 log.txt
然后我这里设置了要tensorboard,这里咱们看看效果
tensorboard --logdir=runs/train/epx2/logs
这样一个完整的训练过程走完了
使用模型
之后到了使用阶段。
这里有几个参数，设置都挺简单的，我就不在复述 了。
重点是这个。
例如这里，我叫它验证一张图片
编码实现
整个实现过程其实还是很简单的，就是前面调了一会儿。
配置文件
首先由于是V0.1版本嘛，所以读取的文件就是一个字典。
"""
参数设置
"""
import os
# 设置你的类别，训练集与验证集的类别要保持一致
Classes = 2
Classfiy = ["1","100"]
#数据集根目录 注意使用 /
Data_Root = r'F:\\projects\\PythonProject\\MyClassfication\\mydata'
#划分的训练集
Train = "train"
#划分的测试集
Valid = 'valid'
SaveWeights = "runs"
读取配置文件
这里的话，就是工具类里面。
"""
第一个小dome,先从读取python字典开始走起。
"""
from data.ModelConfig import *
class ReadDict(object):
@staticmethod
def ReadModelClasses():
label_dict=dict()
index = 0
#读取分类标签
for label in Classfiy:
label_dict[label]=index
index+=1
return label_dict
@staticmethod
def ReadDataRoot():
return Data_Root
HU数据集解析器
这个是咱们的第一个重点。
这里有两个，一个就是读取文件的嘛，读取数据集。
from data.ModelConfig import *
import os
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms
from utils.ReaderProcess.ReadDict import ReadDict
class MyDataSet(Dataset):
def __init__(self, data_dir, transform=None):
self.label_name = ReadDict.ReadModelClasses()
self.data_info = self.get_img_info(data_dir)
self.transform = transform
def __getitem__(self, index):
path_img, label = self.data_info[index]
img = Image.open(path_img).convert('RGB')
if self.transform is not None:
img = self.transform(img)
return img, label
def __len__(self):
return len(self.data_info)
@staticmethod
def get_img_info(data_dir):
data_info = list()
label_dict=ReadDict.ReadModelClasses()
for root, dirs, _ in os.walk(data_dir): #
# 遍历类别
for sub_dir in dirs:
img_names = os.listdir(os.path.join(root, sub_dir))
img_names = list(filter(lambda x: x.endswith('.jpg'), img_names))
# 遍历图片
for i in range(len(img_names)):
img_name = img_names[i]
path_img = os.path.join(root, sub_dir, img_name)
label = label_dict[sub_dir]
data_info.append((path_img, int(label)))
return data_info
if __name__ == '__main__':
train_dir=Data_Root+"\\\\"+Train
train_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
])
train_data = MyDataSet(data_dir=train_dir, transform=train_transform)
# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=16)
for data in train_data:
img , label = data
print(img.shape)
然后，还有一个，tensor转换器
import torchvision.transforms as transforms
#传说之中的转换器！这里主要有两个，一个是送入神经网络训练的，还有一个是负责验证的。
class TransFormAtions(object):
norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]
train_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
valid_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
上面的参数，是嫖的（norm_mean,norm_std）。
其他
这个都是在utils里面的，那就一块说了吧。
Log.py日志输出
"""
这个主要负责日志，日志输出，后面其实可以在这里实现，模型终短后恢复权重训练
"""
import os
def PrintLog(EXP_Path):
# 保存打印日志
# Open a file
if(os.path.exists(EXP_Path)):
Save_Log_print = EXP_Path+"\\\\log.txt"
return open(file=Save_Log_print,mode='w',encoding='utf-8')
else:
raise Exception("文件异常")
一些模型训练的复杂方法
ModelUtils.py
import random
import numpy as np
import torch
# 设置一个随机种子，这样保证两次随机生成的数是一样的，主要是为了初始化
def set_seed(seed=1):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
模型保存，这个是哪个文件就不用我说了吧。
"""
负责保存模型
"""
from data.ModelConfig import *
import os
import torch
def CreatRun(idx):
Save_Path_Root = SaveWeights + "\\\\train" + "\\\\epx"+str(idx)
if(not os.path.exists(Save_Path_Root)):
os.makedirs(Save_Path_Root)
return Save_Path_Root
else:
return CreatRun(idx+1)
def Save_Model(EXP_Path,weight_best,weight_last):
if(os.path.exists(EXP_Path)):
save_path_root = EXP_Path+"\\\\weights"
if(not os.path.exists(save_path_root)):
os.makedirs(save_path_root)
save_path_best = save_path_root+"\\\\best.pth"
save_path_last = save_path_root+"\\\\last.pth"
torch.save(weight_best,save_path_best)
torch.save(weight_last,save_path_last)
print()
print("best:",save_path_best)
print("last:",save_path_last)
else:
raise Exception("保存地址异常")
if __name__ == '__main__':
Save_Model(0,'1',"best")
使用模型 LeNet
前面说过，深度学习呢，要具备两个条件，一个是天赋，一个是时间。
天赋就是咱们的神经网络结构，这里咱们选择一个简单的。
这里我们选择的网络为LeNet。
结构如下：
这块咱们搞一个2分类的来看看（数据集不好搞）。
图片输入假设是 3 x 32 x 32 batch size 为 4
import torch
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential, CrossEntropyLoss
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self,classes):
super().__init__()
self.feature = Sequential(
nn.Conv2d(3,6,kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2,stride=2),
nn.Conv2d(6,16,5), # A
nn.ReLU(),
nn.MaxPool2d(kernel_size=2,stride=2)
)
self.classifiar = nn.Sequential(
nn.Linear(16*5*5,120), # B
nn.ReLU(),
nn.Linear(120,84),
nn.ReLU(),
nn.Linear(84,classes)
)
def forward(self,x):
x = self.feature(x)
x = x.view(x.size()[0],-1)
x = self.classifiar(x)
return x
if __name__ == '__main__':
net = LeNet(3)
data = torch.randn(4,3,32,32)
out = net(data)
这个在咱们这个项目里面是这个
import torch
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential, CrossEntropyLoss
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self,classes):
super().__init__()
self.feature = Sequential(
nn.Conv2d(3,6,kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2,stride=2),
nn.Conv2d(6,16,5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2,stride=2)
)
self.classifiar = nn.Sequential(
nn.Linear(16*5*5,120), # B
nn.ReLU(),
nn.Linear(120,84),
nn.ReLU(),
nn.Linear(84,classes)
)
def forward(self,x):
x = self.feature(x)
x = x.view(x.size()[0],-1)
x = self.classifiar(x)
return x
def initialize_weights(self):
#参数初始化，随便给点权重，这样的话会加快一点速度（训练）
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.xavier_normal_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight.data, 0, 0.1)
m.bias.data.zero_()
if __name__ == '__main__':
net = LeNet(3)
data = torch.randn(4,3,32,32)
out = net(data)
训练实现
看到红色框哈
之后是咱们的重点。
import argparse
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from utils.DataSet import MyDataSet
from models.LeNet import LeNet
import torch.optim as optim
from utils import ModelUtils
from data.ModelConfig import *
from utils.DataSet.MyDataSet import MyDataSet
from utils.DataSet.TransformAtions import TransFormAtions
from utils.ReaderProcess.ReadDict import ReadDict
import os
from utils import SaveModel
from utils import Log
from torch.utils.tensorboard import SummaryWriter
def train():
# 固定一个随机种子，
# 固定住深度模型训练的过程，使得每次从头开始训练模型初始化方式和数据读取方式保持一致
ModelUtils.set_seed()
#初始化驱动
device=None
if (torch.cuda.is_available()):
if(not opt.device=='cpu'):
div = "cuda:"+opt.device
# 这边后面还得做一个检测，看看有没有坑货，乱输入
device = torch.device(div)
print("\\033[0;31;0m使用GPU训练中：\\033[0m".format(torch.cuda.get_device_name()))
else:
device = torch.device("cpu")
print("\\033[0;31;40m使用CPU训练\\033[0m")
#创建 runs exp 文件
EPX_Path = SaveModel.CreatRun(0)
#日志相关的准备工作
wirter = None
openTensorboard=opt.tensorboardopen
path_board=None
if(openTensorboard):
path_board = EPX_Path+"\\\\logs"
wirter = SummaryWriter(path_board)
fo = Log.PrintLog(EPX_Path)
#准备数据集
transformations = TransFormAtions()
train_data_dir=opt.train_dir
if(not train_data_dir):
train_data_dir = ReadDict.ReadDataRoot()+"\\\\"+Train
if(not os.path.exists(train_data_dir)):
raise Exception("训练集路径错误")
train_data = MyDataSet(data_dir=train_data_dir, transform=transformations.train_transform)
valid_data_dir = opt.valid_dir
if(not valid_data_dir):
valid_data_dir = ReadDict.ReadDataRoot()+"\\\\"+Valid
if(not os.path.exists(valid_data_dir)):
raise Exception("测试集路径错误")
valid_data = MyDataSet(data_dir=valid_data_dir, transform=transformations.valid_transform)
# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=opt.batch_size,num_workers=opt.works, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=opt.batch_size)
#开始进入网络训练
#1 开始初始化网络，设置参数啥的
#1.1 初始化网络
net = LeNet(Classes)
net.initialize_weights()
net = net.to(device)
#1.2选择交叉熵损失函数，做分类问题一般是选择这个损失函数的
criterion = nn.CrossEntropyLoss()
#1.3设置优化器
optimizer = optim.SGD(net.parameters(), lr=opt.lr, momentum=0.09) # 选择优化器
# 设置学习率下降策略,默认的也可以，那就不设置嘛，主要是不断去自动调整学习的那个速度
scheduler 基于pytorch框架实现手写图片的分类
基于昇腾CANN的卡通图像生成可在线体验啦！十分钟带你了解CANN应用开发全流程
PyTorch 深度学习实战 | 基于 ResNet 的花卉图片分类