数据集

首先，当然还是要说说数据集问题，我采用的是最经典的MNIST识别手写数字，原因嘛，一是简单，二是正好任务需要。由于我的任务只需要1～5数字，所以我的数据集是更简单版！！5分类问题。下载是从kaggle里面搜索的。
kaggle

train（代码详解）

导入库

from pathlib import Path
from hpp import *
from model import *
import torch
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader

这里肯定是在后面写代码的时候需要什么导入什么的，hpp是我自己写的一些函数和类的py文件（这个文件的库文件导入自己看，我就不放出来了），model是模型搭建文件，torch需要去官网下载。troch

获取当前文件目录路径

FILE=Path(__file__).resolve()
ROOT=FILE.parents[0]

调用main函数

main()

编写main函数

def main():
    #data
    data=ROOT/"dataset"
    train_data,train_label,test_data,test_label=get_train_test_data(data)

    #model
    model=Module()
    
    #train
    train(model,train_data,train_label,test_data,test_label)

先是获取到data数据集的文件路径，然后进入自己写的get_train_test_data函数里面（看下面），拿到数据集的数据路径和标签之后，创建model（后面会说），将模型和数据标签扔进train函数里面去训练。

编写get_train_test_data

def get_train_test_data(data_path):
    train_ratio=0.8
    test_ratio=1-train_ratio
    train_data,test_data,train_label,test_label=[],[],[],[]
    
    files=[]
    files+=glob.glob(str(data_path/"**"/"*.*"),recursive=True)
    files_labels=[os.path.dirname(label) for label in files]
    labels_index=[os.path.basename(label) for label in files_labels]
    index=[i for i in range(len(labels_index))]
    random.seed(0)
    random.shuffle(index)
    
    for i in range(0,int(len(index)*train_ratio)):
        train_data.append(files[index[i]])
        train_label.append(int(labels_index[index[i]]))
        
    for i in range(int(len(index)*train_ratio),len(index)):
        test_data.append(files[index[i]])
        test_label.append(int(labels_index[index[i]]))
        
    return train_data,train_label,test_data,test_label

根据train_ratio设置train和test的划分比例，通过glob获取数据集里面的所有文件，可以查找一下这个库是怎么用的（这个地方需要自己去改）,由于我的文件label序号就是文件的名字，所以根据dirname,basename去获取下标，然后通过random去打乱数据集。

编写train

def train(model,train_data,train_label,test_data,test_label,epochs=10,batch_size=256,imgsz=80):
    model.train()
    #loss
    loss_fn=torch.nn.CrossEntropyLoss()
    
    #optimizer
    optimizer=torch.optim.Adam(model.parameters(),lr=0.0005)
    
    #loader
    train_loader=dataloader(train_data,train_label)
    train_size=len(train_loader)
    train_loader=DataLoader(dataset=train_loader,batch_size=batch_size,num_workers=4)
    all_step=len(train_loader)
    writer=SummaryWriter("logs")

    #train
    print("------------------train------------------")
    for i in range(epochs):
        all_acc=0
        all_loss=0
        step=0
        for img,it_label in train_loader:
            output=model(img)
            loss=loss_fn(output,it_label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            acc=(output.argmax(1)==it_label).sum()
            print(acc)
            all_acc+=acc
            print(all_acc)
            all_loss+=loss
            step+=1
            print(f"已完成{step/all_step:.1f},当前loss:{loss}")
        accuracy=all_acc/train_size
        writer.add_scalar("train_loss",all_loss,epochs)
        writer.add_scalar("train_acc",accuracy,epochs)
        test_acc=test(model,test_data,test_label)
        print(f"epochs[{i+1}/{epochs}],train_loss:{all_loss:.3f}|train_acc:{accuracy*100:.2f}%|test_acc:{test_acc*100:.2f}%")
        writer.add_scalar("test_acc",test_acc,epochs)
        dir_path = 'save/train'  
        if not os.path.exists(dir_path): 
            os.makedirs(dir_path)
        torch.save(model.state_dict(),f"save/train/Module_{i}.pth") 
        torch.onnx.export(model,img,f"save/train/Module_{i}.onnx")
        print("模型已保存")       
    writer.close()

设置损失函数和优化器，加载数据看后面），开启tensorboard记录，然后开始训练，从加载器里面读取数据标签，将图片放进模型里面，得到预测结果，经过损失函数计算损失率，优化器优化梯度，损失函数反向传播，计算准确率，打印结果，并且tensorboard记录，进入test函数测试（下面），最后保存模型，如果要用c++就转成onnx模型。

test

def test(model,data,label,batch_size=256,imgsz=80):
    model.eval()
    
    print("---------------test----------------")
    #dataloader
    test_loader=dataloader(data,label)
    test_size=len(test_loader)
    test_loader=DataLoader(dataset=test_loader,batch_size=batch_size,num_workers=4)
    all_step=len(test_loader)
     
    #test
    with torch.no_grad():
        all_acc=0
        for img,it_label in test_loader:
            output=model(img)
            acc=(output.argmax(1)==it_label).sum()
            all_acc+=acc
            accuracy=all_acc/test_size
    return accuracy

获取数据同上，无梯度变化，去测试准确率，并且返回准确率。

dataloader

class dataloader(Dataset):
    def __init__(self,data_path,label):
        self.data_path=data_path
        self.label=label
        
        self.imgs=[cv2.imread(data) for data in self.data_path]
        
    def __len__(self):
        return len(self.imgs)
    
    def __getitem__(self, index) :
        img=torch.tensor(self.imgs[index],dtype=torch.float32)
        #改变通道位置
        img=img.permute(2,0,1)
        label=torch.tensor(self.label[index],dtype=torch.long)
        return img,label

主要目的是根据路径通过opencv读取到图片，并且注意要将BGR改为RGB，同时将读取图片和标签都转化成tensor类型，这里必须要有__len__和__getitem__。

model

最开始的时候，我也不知道网络该怎么搭建，从何搭建，所以我选择从我之前学习的yolo下手，学习它的代码（所以我的代码有点yolo的影子，可能是迷你版到偏移了），但是实际效果，我只能说一言难尽，可能是我写的有问题吧。（一定是，以后再研究研究）（现在发现是数据问题）后面我想到了最开始学习搭建的一个基础网络结构，巻积加下采样加激活函数，三个最简单的CNN结构的堆叠。抱着试试看的态度，我开始了我的model。

import torch
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear, ReLU


class module(nn.Module):
    def __init__(self):
        super(module,self).__init__()
        self.sequential = Sequential(
            Conv2d(3, 32, 5, 1, 2),
            MaxPool2d(2),
            ReLU(),
            Conv2d(32, 32, 5, 1, 2),
            MaxPool2d(2),
            ReLU(),
            Conv2d(32, 64, 5, 1, 2),
            MaxPool2d(2),
            ReLU(),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self,x):
        x =self.sequential(x)
        return x

损失函数和优化器（多试）

损失函数确定是交叉熵函数，优化器的话，我建议是先SGD再Adam，再考虑其他的。

层数

然后想着要不要改变层数，先是增加层数，只能说效果很崩溃，发现增加竟然效果不好！！突然意识到，会不会我减少一层效果反而会更好呢？结果就是，虽然降低了1%的准确率，但是收敛速率提高了，准确率在20轮的时候提高了。感觉可能是有点退化的那个意思了。

Batch_size

我最开始的大小设置的是128,后面尝试了256，64，最后还是感觉256比较合适。

下采样

为什么会想到改变下采样呢？因为yolo在后面版本的时候就没有使用了，想着应该有道理，所以改用了1*1巻积核，步长为2达到下采样的效果。然后也想过空洞巻积，效果不太好，损失增大了，而且运行速率降低很多。

学习率

从0.01到0.00001尝试。

最后的model代码

import torch
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear, ReLU,Softmax,BatchNorm2d,SiLU,ELU


class Module(nn.Module):
    def __init__(self):
        super(Module,self).__init__()
        self.sequential = Sequential(
            Conv2d(3, 64, 5, 1, 2),
            Conv2d(64,64,1,2),
            BatchNorm2d(64),
            Softmax(),
            Conv2d(64, 128, 5, 1, 2),
            Conv2d(128,128,1,2),
            BatchNorm2d(128),
            Softmax(),
            Flatten(),
            Linear(51200, 64),
            Linear(64,8)
        )

    def forward(self,x):
        x =self.sequential(x)
        return x

detect

导入库，得到绝对路径

from pathlib import Path
import torch
from model import *
from hpp import *
from torch.utils.data import DataLoader
import time
FILE=Path(__file__).resolve()
ROOT=FILE.parents[0]
run("save/train/Module_2.pth",ROOT/"dataset",80)

run函数

def run(weights,source,imgsz):
    #load_model
    load_model=torch.load(weights)
    model=Module()
    model.load_state_dict(load_model)
    model.eval()
    
    #dataloader
    datasets=LoadImages(source)
    
    label_list=["负样本","1","2","3","4","5","烧饼","基地"]
    with torch.no_grad():
        for path,im,img in datasets:
            start=time.time()
            im=im.unsqueeze(0)
            output=model(im)
            _,predicted=torch.max(output,1)
            classIndex=predicted.item()
            print(label_list[classIndex])
            print(time.time()-start)

先加载模型，加载图片（下面），无梯度检测，计算时间，记得要添加一个维度，也就是batch的位置，获取预测结果。

LoadImage函数

class LoadImages:
    def __init__(self,path,imgsz=80,stride=32,auto=True,transforms=None, vid_stride=1):
        self.path,self.imgsz,self.auto=path,imgsz,auto
        self.files = []
        files=[]
        files+=glob.glob(str(path/"**"/"*.*"),recursive=True)
        index=[i for i in range(len(files))]
        random.seed(42)
        random.shuffle(index)
        
        for i in range(0,int(len(index))):
            self.files.append(files[index[i]])
        images = [x for x in self.files if x.split(".")[-1].lower() in IMG_FORMATS]
        self.nf=len(images)
        
    def __iter__(self):
        self.count=0
        return self
    def __next__(self):
        if self.count==self.nf-1:
            raise StopIteration

        self.count+=1
        path=self.files[self.count]
        img=cv2.imread(path)
        assert img is not None, f"Image Not Found {img}"
        im = letterbox(img, self.imgsz, auto=self.auto)[0]  # padded resize
        im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        im = np.ascontiguousarray(im)  # contiguous
        cv2.imshow("img",img)
        cv2.waitKey(0)
        img=torch.tensor(img,dtype=torch.float32)
        im=torch.tensor(im,dtype=torch.float32)
        return self.path,im,img

跟上面差不多（讲不动了，自己看吧），主要还是通道要改，shape要改，最后转成tensor类型。

letterbox

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    """Resizes and pads image to new_shape with stride-multiple constraints, returns resized image, ratio, padding."""
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)

对图片做一定处理。

总结

自己搭建模型去训练数据集并没有那么的困难，如果数据集简单的话，一个简单网络能达到的效果就会很不错，这个网络是达到96%，然后调参之后是98%，后面我也打算去做其他的尝试，尽量更高。