编写一个简单的UNet图像分割网络_（三）训练模块

本文链接：https://blog.csdn.net/weixin_66378701/article/details/139729248

以下是一个基于U_Net网络的简易图像分割算法，如有问题请大佬指点一二：

项目目录

三、训练模块

首先我们要明确我们要用于训练的设备是什么，是gpu还是cpu，当前使用cuda训练的技术较为成熟，教程较多，但如果确实没有cuda，也没必要强行给自己创造需求，现在才刚开始学习，没必要几千上万块买一张卡，使用cpu训练一些小模型小数据集也是可以的，或者上云平台

（这里埋个坑，我自己使用的是intel显卡，后续会出intel显卡深度学习环境的搭建过程，官方有介绍文档，但是在安装过程中也存在一定需要补充的点，这里先贴上官方文档链接：Welcome to Intel® Extension for PyTorch* Documentation!，本文代码以cuda及cpu为主，不涉及intel xpu）

训练前的准备工作，我们需要确定训练设备，加载数据集路径，加载权重文件（权重文件就是别人从0训练调试出来的一个模型，我们基于人家调试出来的模型进行训练，训练效果会更好，但是需要确保权重文件的网络结构匹配，网络结构不匹配权重文件不通用）

# train.py
# optim是优化器，nn提供了很多神经网络层和损失函数
from torch import nn,optim
import torch
from torch.utils.data import DataLoader
from data import *  # import * 意思是把data下所有的都导进来
from net import *
from torchvision.utils import  save_image

# 选择执行训练的设备，有cuda就用cuda，没有cuda会使用cpu
device=torch.device('cuda'if torch.cuda.is_available() else 'cpu')
# 权重文件，注意，这里不是随便一个权重文件就拿来用了，权重文件中包含曾经训练过的网络结构
# 把其他网络结构下训练的权重文件直接拿来用的话会导致网络结构不匹配，会报Missing key(s)”和“Unexpected key(s)”的错误
# 不提供权重文件，这里第一次加载的时候就新建个空文件夹叫params就行，里面不用放东西，他会报Fail load weight然后生成一个unet.pth，第二次跑的时候就能成功加载了
weight_path='params/unet.pth'
# 数据集路径
data_path=r'D:\\Softwares\\Codes\\Codes\\Python\\UNet_Demo\\0612_UNet_Demo_test\\VOCdevkit\\VOC2007'

很激动吧，搞了这么久终于能真的把模型跑起来了，但是别急哦，越急越出错哟

目前我们的数据集、网络、以及训练设备都是相互独立的，要用device来训练，就得先把数据集和网络挂载上去，首先执行网络的挂载，数据集挂载在后面

# train.py
if __name__=='__main__':
    # 加载数据集，显卡性能较低的话调低batch_size，shuffle是打乱数据集的选项
    data_loader=DataLoader(MyDataset(data_path),batch_size=4,shuffle=True)
    # 把网络挂载在device上
    net=UNet().to(device)

接下来挂载权重文件咯，从这开始说明多数写在注释中了，保证对应每一步，更便于理解

# 属于train.py的main方法

	# 判断权重文件是否存在，如果存在，就加载进来
    # 一般来说，不建议一切从0开始训练，训练效果较差，训练其他项目时还是挂载个权重文件
    # 再次申明作者不提供权重文件，第一次训练的时候就新建个空文件夹叫params就行，里面不用放东西
    # 他会报Fail load weight然后生成一个unet.pth
    if os.path.exists(weight_path):
        # 使用一些曾在cuda上训练过的权重文件在cpu上进行训练时，必须申明map_location=device
        # 为保证代码兼容性，这里默认加上
        net.load_state_dict(torch.load(weight_path,map_location=device))
        print('Successful load weight')
    else:
        print('Fail load weight')

设置使用的优化器和损失计算方法，这个没什么说的

# 属于train.py的main方法

	# 创建优化器，使用Adam优化器，这里net.parameters()是把网络中的所有参数都拿出来
    opt=optim.Adam(net.parameters())
    # 实际是一个二分类问题，使用二分类交叉熵损失函数BCELoss
    loss_fun=nn.BCELoss()

设置参数开始训练咯

# 属于train.py的main方法

	# epoch一般作为迭代次数，就是训练多少轮，这里直接把epoch用作一个计数器，直接使用死循环无限训练
    # 每50轮保存一次权重文件，随时可以停止训练
    epoch = 1
    # 如果要确定轮数进行训练，可以定义起始轮数Init_Epoch，终止轮数UnFreeze_Epoch
    # 为什么会有起始轮数，因为这是冻结训练的知识点，就是用之前训练的权重文件继续训练，这里再挖个坑，后续增加冻结训练
    # 这里Init_Epoch=1就行，UnFreeze_Epoch是你自己想训练的轮数，用for循环替换掉while
    # Init_Epoch=1，UnFreeze_Epoch=100
    # for epoch in range(Init_Epoch, UnFreeze_Epoch):
    while True:
        # 遍历数据集，enumerate()会生成一个迭代器，迭代器中包含两个元素，第一个元素是索引，第二个元素是数据集
        for i,(image,segment_image) in enumerate(data_loader):
            # 把数据集都挂在device上
            # 这里相当于image=image.to(device)
            # segment_image=segment_image.to(device)
            image,segment_image=image.to(device),segment_image.to(device)

            # 把image传入net，得到输出的图，使用net()默认直接调用forward
            out_image=net(image)
            # 使用loss_fun对预测值和实际值进行比较，计算训练损失
            train_loss=loss_fun(out_image,segment_image)
            # 清空梯度
            opt.zero_grad()
            # 反向计算梯度
            train_loss.backward()
            # 更新梯度
            opt.step()

            # 如果训练到5次，就打印一下损失值
            if i%5==0:
                print(f'{epoch}-{i}-train_loss===>>{train_loss.item()}')
            # 每训练50次，就保存一次权重
            if i%50==0:
                torch.save(net.state_dict(),weight_path)

            # 把原图image，分割图segment_image，预测图out_image拼接在一起，便于观察
            # 这里使用0索引，因为batch_size=4，所以image和segment_image都是4张图，拿出其中第一张
            # 以作者调试过程中参数为例，此时image.shape=(4,3,256,256)，image[0].shape=(3,256,256)，这才是一张图
            _image=image[0]
            _segment_image=segment_image[0]
            _out_image=out_image[0]
            img=torch.stack([_image,_segment_image,_out_image],dim=0)
            # 保存拼接后的图片
            save_image(img,f'{save_path}/{i}.png')

            epoch+=1

对了，这里因为我们保存了生成的图片，所以需要给最开始读取路径下面新加一个保存的路径哟，项目根目录新建一个train_image文件夹

# train.py
save_path='train_image'

可以开始右键train.py开始训练咯，因为没有用别人的预训练权重，我们是从0训练的，所以loss波动特别大很正常

因为我是第二次训练，所以是Successful load weight，第一次运行的话会报Fail load weight，但是此时会看到params文件夹下多了一个unet.pth，重新跑一下就行

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

又阶段性胜利咯！恭喜恭喜！！

下附完整代码

# train.py

# optim是优化器，nn提供了很多神经网络层和损失函数
from torch import nn,optim
import torch
from torch.utils.data import DataLoader
from data import *  # import * 意思是把data下所有的都导进来
from net import *
from torchvision.utils import  save_image

# 选择执行训练的设备，有cuda就用cuda，没有cuda会使用cpu
device=torch.device('cuda'if torch.cuda.is_available() else 'cpu')
# 权重文件，注意，这里不是随便一个权重文件就拿来用了，权重文件中包含曾经训练过的网络结构
# 把其他网络结构下训练的权重文件直接拿来用的话会导致网络结构不匹配，会报Missing key(s)”和“Unexpected key(s)”的错误
# 不提供权重文件，这里第一次加载的时候就新建个空文件夹叫params就行，里面不用放东西，他会报Fail load weight然后生成一个unet.pth
# 第二次跑的时候就能成功加载了
weight_path='params/unet.pth'
# 数据集路径
data_path=r'D:\\Softwares\\Codes\\Codes\\Python\\UNet_Demo\\0612_UNet_Demo_test\\VOCdevkit\\VOC2007'
save_path='train_image'
if __name__=='__main__':
    # 加载数据集，显卡性能较低的话调低batch_size，shuffle是打乱数据集的选项
    data_loader=DataLoader(MyDataset(data_path),batch_size=4,shuffle=True)
    # 把网络挂载在device上
    net=UNet().to(device)
    # 判断权重文件是否存在，如果存在，就加载进来
    # 一般来说，不建议一切从0开始训练，训练效果较差，还是挂载个权重文件
    # 再次申明作者不提供权重文件，第一次训练的时候就新建个空文件夹叫params就行，里面不用放东西
    # 他会报Fail load weight然后生成一个unet.pth
    if os.path.exists(weight_path):
        # 使用一些曾在cuda上训练过的权重文件在cpu上进行训练时，必须申明map_location=device
        # 为保证代码兼容性，这里默认加上
        net.load_state_dict(torch.load(weight_path,map_location=device))
        print('Successful load weight')
    else:
        print('Fail load weight')
    # 创建优化器，使用Adam优化器，这里net.parameters()是把网络中的所有参数都拿出来
    opt=optim.Adam(net.parameters())
    # 实际是一个二分类问题，使用二分类交叉熵损失函数BCELoss
    loss_fun=nn.BCELoss()


    # epoch一般作为迭代次数，就是训练多少轮，这里直接把epoch用作一个计数器，直接使用死循环无限训练
    # 每50轮保存一次权重文件，随时可以停止训练
    epoch = 1
    # 如果要确定轮数进行训练，可以定义起始轮数Init_Epoch，终止轮数UnFreeze_Epoch
    # 为什么会有起始轮数，因为这是冻结训练的知识点，就是用之前训练的权重文件继续训练，这里再挖个坑，后续增加冻结训练
    # 这里Init_Epoch=1就行，UnFreeze_Epoch是你自己想训练的轮数，用for循环替换掉while
    # Init_Epoch=1，UnFreeze_Epoch=100
    # for epoch in range(Init_Epoch, UnFreeze_Epoch):
    while True:
        # 遍历数据集，enumerate()会生成一个迭代器，迭代器中包含两个元素，第一个元素是索引，第二个元素是数据集
        for i,(image,segment_image) in enumerate(data_loader):
            # 把数据集都挂在device上
            # 这里相当于image=image.to(device)
            # segment_image=segment_image.to(device)
            image,segment_image=image.to(device),segment_image.to(device)

            # 把image传入net，得到输出的图，使用net()默认直接调用forward
            out_image=net(image)
            # 使用loss_fun对预测值和实际值进行比较，计算训练损失
            train_loss=loss_fun(out_image,segment_image)
            # 清空梯度
            opt.zero_grad()
            # 反向计算梯度
            train_loss.backward()
            # 更新梯度
            opt.step()

            # 如果训练到5次，就打印一下损失值
            if i%5==0:
                print(f'{epoch}-{i}-train_loss===>>{train_loss.item()}')
            # 每训练50次，就保存一次权重
            if i%50==0:
                torch.save(net.state_dict(),weight_path)

            # 把原图image，分割图segment_image，预测图out_image拼接在一起，便于观察
            # 这里使用0索引，因为batch_size=4，所以image和segment_image都是4张图，拿出其中第一张
            # 以作者调试过程中参数为例，此时image.shape=(4,3,256,256)，image[0].shape=(3,256,256)，这才是一张图
            _image=image[0]
            _segment_image=segment_image[0]
            _out_image=out_image[0]
            img=torch.stack([_image,_segment_image,_out_image],dim=0)
            # 保存拼接后的图片
            save_image(img,f'{save_path}/{i}.png')

            epoch+=1