CV 使用PyTorch构建并训练深度神经网络,详细完整流程

PyTorch训练模型详细完整流程,实操记录

  • Pytorch基本情况
    • 1. 个人用的相关东西的版本
    • 2. 总体流程介绍(自己训练的流程整理的)
  • 开始用PyTorch,按照上述的顺序介绍并实现
    • 1. 写数据集的txt
    • 2. 构建 Dataset 类
    • 3. 设置dataloader(待补充详细方法介绍)
    • 4. 构建模型
    • 5. 加载模型及其参数们
    • 6.训练测试和保存模型

Pytorch基本情况

1. 个人用的相关东西的版本

python==3.6.8

torch==1.9.0+cu102

torchvision==0.10.0+cu102

如需要安装的话,可以在终端这样:

pip install torchvision==0.10.0+cu102

2. 总体流程介绍(自己训练的流程整理的)

  1. 写数据集的txt

写一个保存了数据集路径和标签信息的txt,这里可以分为 train.txt 和 test.txt 。

  1. 构建 Dataset 类

这个类能读取步骤一的txt的内容,并通过 transform 设置Dataset类数据的相关的信息,如toTensor()、等。
并返回一个包含了图片数据和标签的Dataset对象给dataloader。

  1. 设置dataloader

就是接收第二步的Dataset对象,并设置训练数据相关的东西,如:batchsize、等。
后面训练就是直接将这个dataloader喂给model就行了。

  1. 构建模型
  2. 加载模型及其参数们
  3. 训练测试保存模型

开始用PyTorch,按照上述的顺序介绍并实现

1. 写数据集的txt

作用:
写数据集的 txt文件,是把数据集的路径和标签信息记录在txt文件中,为了让 pytorch 后面的 Dataset 类能通过读取txt的方式就能读取数据集。
同时在这一步可以通过写 train.txt 、val.txt 和 test.txt 来区分训练集、验证集和测试集

先说明我的数据集文档情况

DATASET
|—0
|—|—1.bmp
|—|—2.bmp
|—|—3.bmp
|—|—省略号
|—1
|—|—1.bmp
|—|—2.bmp
|—|—3.bmp
|—|—省略号

我只需要把数据集分成训练集和测试集就行了,就只写了 train.txt 和 test.txt。
上代码

'''
         注意这里是这些代码的文件名
project//generate_txt.py
'''
import os
import random

# 指向数据集文件夹DATASET的路径位置
DATA_ROOT = ".//DATASET"
# 0~100,表示训练集和测试集划分的比率,80就是100张图片中有80个划分为训练集
train_test_rate = 80

# 获取DATASET文件夹下文档路径的列表,我这里是只有两个类别 1 和 0
# 就是['DATASET//1', 'DATASET//0']
dir_paths = [os.path.join(DATA_ROOT, p1) for p1 in os.listdir(DATA_ROOT)]
# 建立并以添加的方式打开训练集和测试集的txt
ftrain = open(os.path.join(DATA_ROOT, "train.txt"), "a+")
ftest = open(os.path.join(DATA_ROOT, "test.txt"), "a+")
# 分别遍历各类别的数据集图片
for i, dir_path in enumerate(dir_paths):
	# 因为我的文件夹名字就是label,所以我直接把遍历的index定为label,这里可以根据自己的喜好定
    label = i
    print(f"dirpath:{dir_path}, label:{i}")
    # 获取在某个类别文件夹下的所有图片路径,并保存到一个列表中。
    img_paths = [os.path.join(dir_path, p1) for p1 in os.listdir(dir_path) if p1[-4:].lower().endswith('.bmp')]
    # 遍历某个类别的所有图片路径的列表
    for j, img_path in enumerate(img_paths):
    	# 要保存到txt的内容,就是 图片路径+空格+标签,记得要过行“\n”。
        write_thing = img_path + " " + str(label) + "\n"
        # 随机写到训练集的txt和测试集的txt中
        if random.randint(0, 100) >= 80:
            ftrain.write(write_thing)
        else:
            ftest.write(write_thing)
# 完成前记得关掉ftrain和ftest,养成好习惯
ftest.close()
ftrain.close()

到此,写txt的工作就完成啦。

2. 构建 Dataset 类

作用:
Dataset类就是将txt的内容读取出来,并根据txt里的图片路径去获取数据集图片和标签,并返回一个包含了图片和标签的对象。
怎么做:
其实我们只需要继承pytorch中的Dataset类就可以了,并重写这个类中的两个方法就行了:

get_img_info():
get_img_info就是去读取txt中的路径和标签,这里可以根据自己保存进txt的形式进行读取。
__ getitem __ (self, index):
这个方法的作用是读取一张图片的数据和设置一个标签,就是让一张图片数据和标签对应起来。
可能会有点奇怪,为什么写读取一张图片和其对应标签的函数就能让dataset类返回所有数据集的图片和标签了。这是因为我们是继承了Dataset类,而在Dataset类里面是有一个循环来读取每一个路径下的图片和标签,而读取的方法就是__ getitem __ (index),所以只用写读取一张图片就函数就行了。由于Dataset类中是用index来读取图片的,所以参数不能改变,只能是index。

直接上代码。

'''
     注意这里是这些代码的文件名
project//Dataset.py
'''
from torch.utils.data import Dataset
from PIL import Image


class MyDataset(Dataset):
    def __init__(self, txt_path, transform=None):
        """
        :param data_dir: str, 数据集所在路径
        :param transform: torch.transform,数据预处理
        """
        # data_info存储所有图片路径和标签,方便__getitem__(self, index)用index读取图片
        self.data_info = self.get_img_info(txt_path)
        # 图片的相关参数,如shuffle、等
        self.transform = transform

    # MyDataset类是继承了Dataset类,Dataset类会通过 index 来自动调用__getitem__来读取样本生成对象。
    # 这里就是写 读取每一个样本的方式。
    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

    @staticmethod
    # 这里需要写一个读取txt内容的函数,然后Dataset类会通过__getitem__来读取txt中对应图片路径的图片
    def get_img_info(txt_path):
        data_info = []
        f = open(txt_path, 'r')
        for line in f:
            info = line.split(' ')
            data_info.append((info[0], int(info[1][0])))
        return data_info

    def __len__(self):
        return len(self.data_info)
'''
     注意这里是这些代码的文件名
project//train.py
'''
from project.MyDataset import MyDataset
import torchvision
from torch.utils.data import DataLoader

train_txt = "E://pazhouwork//bupi5//bu5//train.txt"
test_txt = "E://pazhouwork//bupi5//bu5//test.txt"
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])


train_dataset = MyDataset(train_txt, transform)
test_dataset = MyDataset(test_txt, transform)

3. 设置dataloader(待补充详细方法介绍)

作用:
就是设置训练时的一些参数,如batchsize、数据采样方法等。
这里我只设置了batchsize和shuffle。
直接上代码

'''
_cbam_jieyue//train.py
'''
batchsize = 32
shuffle = True
train_loader = DataLoader(train_dataset, batch_size=batchsize, shuffle=shuffle)
test_loader = DataLoader(test_dataset, batch_size=batchsize, shuffle=shuffle)

4. 构建模型

直接上代码,首先是调用系统的ResNet50和ResNet101示例,第二个文件代码是构建自己的模型。有需要的话可以调用系统其他的模型,相信大佬们已经很熟练了。

'''
project//model//resnet.py
'''
from torchvision.models.resnet import resnet50, resnet101
import torch

device="cuda" if torch.cuda.is_available() else "cpu"#GPU加速
# device = "cpu"
print(f"Using {device} device")

resnet50 = resnet50().to(device)
resnet101 = resnet101().to(device)
'''
project//model//model.py
'''

import torch
from torch.nn import Module, Sequential, Flatten, Linear, ReLU, Conv2d, AdaptiveAvgPool2d, BatchNorm2d

'''
具体模型结构可以自己定义哦,下面只是我根据自己设计的模型去写的模型架构,里面的内容只是示例,大佬们不必细究哈
如果想要用系统的模型结构的话,可以直接调用。
'''

device="cuda" if torch.cuda.is_available() else "cpu"#GPU加速
print(f"Using {device} device")

class CONV_BLOCK_INIT(Module):
    def __init__(self, in_channel, out_channel):
        super(CONV_BLOCK_INIT, self).__init__()
        self.conv1 = Conv2d(in_channel, in_channel, kernel_size=1, stride=1, padding='same')
    def forward(self, x):
        x = self.conv1(x)

        return x

class CONV_BLOCK_STATE(Module):
    def __init__(self, in_channel):
        super(CONV_BLOCK_STATE, self).__init__()
        self.conv1 = Conv2d(in_channel, in_channel, kernel_size=1, stride=1, padding='same')

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

class CONV_SHORT_CUT(Module):
    def __init__(self, in_channel):
        super(CONV_SHORT_CUT, self).__init__()
        self.conv1 = Conv2d(in_channel, in_channel, kernel_size=1, stride=1, padding='same')
    def forward(self, x):
        x = self.conv1(x)
        return x

class JUMP_RESIDUAL(Module):
    def __init__(self, in_channel, out_channel):
        super(JUMP_RESIDUAL, self).__init__()
        self.relu = ReLU()
    def forward(self, x):
    	x3 = self.relu(x)
        return x3

class DeepNetwork(Module):#继承nn.Module
    def __init__(self):
        super(DeepNetwork,self).__init__()
        self.flatten=Flatten()
        self.feature_init = Conv2d(3, 64, kernel_size=3, stride=1, padding='same')
        self.state1 = JUMP_RESIDUAL(64, 128)
        self.state2 = JUMP_RESIDUAL(128, 256)
        self.state3 = JUMP_RESIDUAL(256, 512)
        self.state4 = JUMP_RESIDUAL(512, 1024)
        self.avgpool = AdaptiveAvgPool2d((1, 1))
        self.classifier = self._make_linear([1024, 128, 32, 2])
        self.relu = ReLU(inplace=True)
    def forward(self,x):
        x_init = self.feature_init(x)
        x1 = self.state1(x_init)
        x2 = self.state2(x1)
        x3 = self.state3(x2)
        x4 = self.state4(x3)
        x_fea = self.avgpool(x4)
        out = self.classifier(x_fea)
        return out

deep_model=DeepNetwork().to(device)
print(deep_model)

5. 加载模型及其参数们

大概就是这些参数比较重要,还有一些初始化之类的,可能会影响训练结果,后面有时间了再加上去哈。

# 选择模型,填写保存模型名称。
# 这里pretrain=True,就是用预训练好的ResNet50来训练自己的数据集。
model = resnet50(pretrain=True)
save_model_path = "model//resnet50.pth"

epochs = 100
learn_rate = 1e-3
optimizer = optim.Adam(model.parameters(), lr=learn_rate)
criterion = CrossEntropyLoss()
judge = 'f1'

6.训练测试和保存模型

主要是写了一个train函数来进行训练;

写train函数的时候要特别注意下面几个操作的顺序
不能交换!!!不能交换!!!不能交换!!!
我记得之前是有一次交换了不知道哪两个,训练了就跟没训练一样,效果特别差,所以特别拿出来说下。

  1. 获取预测值
    y_pred = model(x)
  2. 优化器
    optimizer.zero_grad()
  3. 获取损失
    loss = criterion(y_pred, y)
  4. 损失反向传播
    loss.backward()
  5. 优化器
    optimizer.step()

并在train函数的末尾加入test函数来测试,获取目前模型的性能,并判断要不要保存;
然后就是参数设置部分。

'''
project//train.py
'''
import torch
from torch import optim
from torch.nn import CrossEntropyLoss
from tqdm import tqdm

from project.MyDataset import MyDataset
import torchvision
from torch.utils.data import DataLoader
from torch.autograd import Variable
from project.model.resnet import resnet50, resnet101

print('''
##########  start to train  #########
''')
def train(model, train_loader, test_loader, epoch, optimizer, criterion, device, judge):
	# 转换成训练模式
    model.train()
    train_loss = 0
    train_len = 0
    # 这个是为了训练的时候有进度条
    loop = tqdm(enumerate(train_loader), total=len(train_loader))
    
    for i, data in loop:
        x, y = data
		# 用GPU训练需要把训练的数据和标签都放进cuda,如下操作
        if device == 'gpu':
            x = Variable(torch.Tensor(x))
            x = x.cuda()
            y = Variable(y)
            y = y.cuda()
        # 获取预测值
        y_pred = model(x)
        # 优化器
        optimizer.zero_grad()
		# 获取损失
        loss = criterion(y_pred, y)
        
        train_loss += loss
        train_len += len(y_pred)
		# 损失反向传播
        loss.backward()
        optimizer.step()

        # 这个是为了训练的时候有进度条,更新进度条信息
        loop.set_description(f'Epoch [{epoch}]')
        loop.set_postfix(loss=train_loss / train_len)

    print(f'finish epoch {epoch}, start to test.')
    test_loss, judge_score = test(model, test_loader, optimizer, criterion, device, judge=judge)
    train_loss = train_loss / train_len
    return train_loss, test_loss, judge_score

def test(model, test_loader, optimizer, criterion, device, judge):
    model.eval()
    test_loss = 0
    count_pred = 0
    count_correct = 0
    count_correct1 = 0
    count_label1 = 0

    with torch.no_grad():
        for i, data in enumerate(test_loader):
            x, y = data

            if device == 'gpu':
                x = Variable(torch.Tensor(x))
                x = x.cuda()
                y = Variable(y)
                y = y.cuda()

            y_pred = model(x)
            optimizer.zero_grad()
            loss = criterion(y_pred, y)
            test_loss += loss

            pred = y_pred.max(1, keepdim=True)[1]
            count_pred += len(pred)

            count_correct += pred.eq(y.view_as(pred)).sum().item()

            for index_label, label in enumerate(y.tolist()):
                if y == 1:
                    count_label1 += 1
                    if pred[index_label][0] == 1:
                        count_correct1 += 1
    accuracy = count_correct/count_pred
    precision = count_pred - count_label1 - count_correct + 2 * count_correct1
    recall = count_correct1/count_label1
    f1_score = 2 * precision * recall / (precision + recall)
    print(f"accuracy:{accuracy} || precision:{precision} || recall:{recall} || f1_score:{f1_score}\n")
    if judge == 'f1':
        judge_score = f1_score
    elif judge == 'accuracy':
        judge_score = accuracy
    elif judge == 'precision':
        judge_score = precision
    elif judge == 'recall':
        judge_score = recall
    else:
        judge_score = False
    return test_loss / count_pred, judge_score

# 来放训练和测试的损失,方便我后面画图用
train_loss_list = []
test_loss_list = []
# 判断要不要保存模型的指标数值
judge_score_max = 0
# 训练epochs次
for epoch in range(1, epochs+1):
	# 训练部分
    train_loss, test_loss, judge_score = train(model, train_loader, test_loader, epoch, optimizer=optimizer, criterion=criterion, device='cpu', judge=judge)
    
    train_loss_list.append(train_loss)
    test_loss_list.append(test_loss)
	
	# 判断要不要保存模型
    if judge_score == False:
        print("judge_score is none.")
    else:
        judge_score_max = max(judge_score, judge_score_max)
        if judge_score_max == judge_score:
            torch.save(model, save_model_path)
            print(f'save model, cause judge_score is {judge_score}')

# 把损失都保存到txt文件中
ftrain = open('record//train.txt', 'a+')
ftest = open('record//test.txt', 'a+')
ftrain.write(f"{item} " for item in train_loss_list)
ftrain.write(f"{item} " for item in test_loss_list)
ftrain.close()
ftest.close()

你可能感兴趣的:(pytorch,深度学习,机器学习)