pytorch 实现语义分割 PSPNet

hawk2014bj2024-05-03 11:23

语意分割是指一张图片上包含多个物体，通过语义分割可以识别物体分类、物体名称、像素识别的任务。和物体检测不同，他不会将物体框出来，而是根据像素的归属把物体标注出来。PSPNet 的输入是一张图片，例如300500，那么输出就是一个 300 500 数组，数组中的值就是分类索引，如果是 20 中分类，分类索引就是 0-19。

PSPNet 的输入和输出，如下图

PSPNet 物体检测流程

  1. 预处理：图像调整为 475*475， 对颜色标准化。
  2. 图像输入神经网络，输出 21 * 475 * 475数据，每个数据就是当前像素分类概率
  3. 根据概率最高的分类生成图像
  4. 将图像还原成原有尺寸

数据准备

  1. 准备数据

  2. 数据增强

  3. 创建 Dataset 类

  4. 创建 Dataloader

    准备数据

    导入软件包

    import os.path as osp
    from PIL import Image

    import torch.utils.data as data

    def make_datapath_list(rootpath):
    """
    创建用于学习、验证的图像数据和标注数据的文件路径列表变量

     Parameters
 ----------
 rootpath : str
     指向数据文件夹的路径

 Returns
 -------
 ret : train_img_list, train_anno_list, val_img_list, val_anno_list
     保存了指向数据的路径列表变量
 """

 #创建指向图像文件和标注数据的路径的模板
 imgpath_template = osp.join(rootpath, 'JPEGImages', '%s.jpg')
 annopath_template = osp.join(rootpath, 'SegmentationClass', '%s.png')

 #训练和验证，分别获取相应的文件 ID（文件名）
 train_id_names = osp.join(rootpath + 'ImageSets/Segmentation/train.txt')
 val_id_names = osp.join(rootpath + 'ImageSets/Segmentation/val.txt')

 #创建指向训练数据的图像文件和标注文件的路径列表变量
 train_img_list = list()
 train_anno_list = list()

 for line in open(train_id_names):
     file_id = line.strip()  #删除空格和换行
     img_path = (imgpath_template % file_id)  #图像的路径
     anno_path = (annopath_template % file_id)  #标注数据的路径
     train_img_list.append(img_path)
     train_anno_list.append(anno_path)

 #创建指向验证数据的图像文件和标注文件的路径列表变量
 val_img_list = list()
 val_anno_list = list()

 for line in open(val_id_names):
     file_id = line.strip()  #删除空格和换行符
     img_path = (imgpath_template % file_id)  #图像的路径
     anno_path = (annopath_template % file_id)  #标注数据的路径
     val_img_list.append(img_path)
     val_anno_list.append(anno_path)

 return train_img_list, train_anno_list, val_img_list, val_anno_list

    #确认执行结果，获取文件路径列表
    rootpath = "./data/VOCdevkit/VOC2012/"

    train_img_list, train_anno_list, val_img_list, val_anno_list = make_datapath_list(
    rootpath=rootpath)

    print(train_img_list[0])
    print(train_anno_list[0])

    创建 dataset
    #首先导入数据处理类和数据增强类
    from utils.data_augumentation import Compose, Scale, RandomRotation, RandomMirror, Resize, Normalize_Tensor

    class DataTransform():
    """
    图像和标注的预处理类。训练和验证时分别采取不同的处理方法
    将图像的尺寸调整为input_size x input_size
    训练时进行数据增强处理

     Attributes
 ----------
 input_size : int
     指定调整图像尺寸的大小
 color_mean : (R, G, B)
     指定每个颜色通道的平均值
 color_std : (R, G, B)
     指定每个颜色通道的标准差
 """

 def __init__(self, input_size, color_mean, color_std):
     self.data_transform = {
         'train': Compose([
             Scale(scale=[0.5, 1.5]),  #图像的放大
             RandomRotation(angle=[-10, 10]),  #旋转
             RandomMirror(), #随机镜像
             Resize(input_size),  #调整尺寸(input_size)
             Normalize_Tensor(color_mean, color_std)  #颜色信息的正规化和张量化
         ]),
         'val': Compose([
             Resize(input_size),  #调整图像尺寸(input_size))
             Normalize_Tensor(color_mean, color_std) #颜色信息的正规化和张量化
         ])
     }

 def __call__(self, phase, img, anno_class_img):
     """
     Parameters
     ----------
     phase : 'train' or 'val'
         指定预处理的执行模式。
     """
     return self.data_transform[phase](img, anno_class_img)

    class VOCDataset(data.Dataset):
    """
    用于创建VOC2012的Dataset的类，继承自PyTorch的Dataset类

     Attributes
 ----------
 img_list : 
     保存了图像路径列表
 anno_list :
     保存了标注路径列表n
 phase : 'train' or 'test'
     保存了标注路径列表
 transform : object
     预处理类的实例
 """

 def __init__(self, img_list, anno_list, phase, transform):
     self.img_list = img_list
     self.anno_list = anno_list
     self.phase = phase
     self.transform = transform

 def __len__(self):
     '''返回图像的张数'''
     return len(self.img_list)

 def __getitem__(self, index):
     '''
     获取经过预处理的图像的张量形式的数据和标注
     '''
     img, anno_class_img = self.pull_item(index)
     return img, anno_class_img

 def pull_item(self, index):
     ''''获取图像的张量形式的数据和标注'''

     # 1.读入图像数据
     image_file_path = self.img_list[index]
     img = Image.open(image_file_path)   #[高度][宽度][颜色RGB]

     # 2.读入标注图像数据
     anno_file_path = self.anno_list[index]
     anno_class_img = Image.open(anno_file_path)   #[ 高度 ][ 宽度 ]

     # 3.进行预处理操作
     img, anno_class_img = self.transform(self.phase, img, anno_class_img)

     return img, anno_class_img

    #确认执行结果

    #(RGB)颜色的平均值和均方差
    color_mean = (0.485, 0.456, 0.406)
    color_std = (0.229, 0.224, 0.225)

    #(RGB)颜色的平均值和均方差
    train_dataset = VOCDataset(train_img_list, train_anno_list, phase="train", transform=DataTransform(
    input_size=475, color_mean=color_mean, color_std=color_std))

    val_dataset = VOCDataset(val_img_list, val_anno_list, phase="val", transform=DataTransform(
    input_size=475, color_mean=color_mean, color_std=color_std))

    #读取数据的示例
    print(val_dataset.getitem(0)[0].shape)
    print(val_dataset.getitem(0)[1].shape)
    print(val_dataset.getitem(0))

    #创建数据加载器

    batch_size = 8

    train_dataloader = data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True)

    val_dataloader = data.DataLoader(
    val_dataset, batch_size=batch_size, shuffle=False)

    #集中保存到字典型变量中
    dataloaders_dict = {"train": train_dataloader, "val": val_dataloader}

    #确认执行结果
    batch_iterator = iter(dataloaders_dict["val"]) #转换为迭代器
    imges, anno_class_imges = next(batch_iterator) #取出第一个元素
    print(imges.size()) # torch.Size([8, 3, 475, 475])
    print(anno_class_imges.size()) # torch.Size([8, 3, 475, 475])

    import numpy as np
    import matplotlib.pyplot as plt
    %matplotlib inline

    每次执行都会改变

    ·读取图像数据

    index = 0
    imges, anno_class_imges = train_dataset.getitem(index)

    #图像的显示
    img_val = imges
    img_val = img_val.numpy().transpose((1, 2, 0))
    plt.imshow(img_val)
    plt.show()

    标注图像的显示

    anno_file_path = train_anno_list[0]
    anno_class_img = Image.open(anno_file_path) # [高度][宽度][颜色RGB]
    p_palette = anno_class_img.getpalette()

    anno_class_img_val = anno_class_imges.numpy()
    anno_class_img_val = Image.fromarray(np.uint8(anno_class_img_val), mode="P")
    anno_class_img_val.putpalette(p_palette)
    plt.imshow(anno_class_img_val)
    plt.show()

    读取图像数据

    index = 0
    imges, anno_class_imges = val_dataset.getitem(index)

    画像的表示

    img_val = imges
    img_val = img_val.numpy().transpose((1, 2, 0))
    plt.imshow(img_val)
    plt.show()

    标注图像的显示

    anno_file_path = train_anno_list[0]
    anno_class_img = Image.open(anno_file_path) # [高度][宽度][颜色RGB]
    p_palette = anno_class_img.getpalette()

    anno_class_img_val = anno_class_imges.numpy()
    anno_class_img_val = Image.fromarray(np.uint8(anno_class_img_val), mode="P")
    anno_class_img_val.putpalette(p_palette)
    plt.imshow(anno_class_img_val)
    plt.show()

PSPNet 网络构建

PSPNet网络包括 Feature（Encoder）、Pyramid Pooling、Decoder、AuxLoss 四个模块构成。

# 实现 PSPNet 网络
# 导入软件包
import torch
import torch.nn as nn
import torch.nn.functional as F

class PSPNet(nn.Module):
    def __init__(self, n_classes):
        super(PSPNet, self).__init__()

        #参数设置
        block_config = [3, 4, 6, 3]  # resnet50
        img_size = 475
        img_size_8 = 60  #设为img_size的1/8

        #创建组成子网络的四个模块
        self.feature_conv = FeatureMap_convolution()
        self.feature_res_1 = ResidualBlockPSP(
            n_blocks=block_config[0], in_channels=128, mid_channels=64, out_channels=256, stride=1, dilation=1)
        self.feature_res_2 = ResidualBlockPSP(
            n_blocks=block_config[1], in_channels=256, mid_channels=128, out_channels=512, stride=2, dilation=1)
        self.feature_dilated_res_1 = ResidualBlockPSP(
            n_blocks=block_config[2], in_channels=512, mid_channels=256, out_channels=1024, stride=1, dilation=2)
        self.feature_dilated_res_2 = ResidualBlockPSP(
            n_blocks=block_config[3], in_channels=1024, mid_channels=512, out_channels=2048, stride=1, dilation=4)

        self.pyramid_pooling = PyramidPooling(in_channels=2048, pool_sizes=[
            6, 3, 2, 1], height=img_size_8, width=img_size_8)

        self.decode_feature = DecodePSPFeature(
            height=img_size, width=img_size, n_classes=n_classes)

        self.aux = AuxiliaryPSPlayers(
            in_channels=1024, height=img_size, width=img_size, n_classes=n_classes)

    def forward(self, x):
        x = self.feature_conv(x)
        x = self.feature_res_1(x)
        x = self.feature_res_2(x)
        x = self.feature_dilated_res_1(x)

        output_aux = self.aux(x)  #将Feature模块中转到Aux模块

        x = self.feature_dilated_res_2(x)

        x = self.pyramid_pooling(x)
        output = self.decode_feature(x)

        return (output, output_aux)
Feature 子网络

Feature 包括 5 个子网络，FeatureMap、两个ResidualBlockPSP和两个 dilated 版的 ResidualBlockPSP。

# Feature Map Convolution 子网络
class conv2DBatchNormRelu(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, bias):
        super(conv2DBatchNormRelu, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels,
                              kernel_size, stride, padding, dilation, bias=bias)
        self.batchnorm = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        #inplase指定不保存输入数据，直接计算输出结果，达到节约内存的目的

    def forward(self, x):
        x = self.conv(x)
        x = self.batchnorm(x)
        outputs = self.relu(x)

        return outputs

class FeatureMap_convolution(nn.Module):
    def __init__(self):
        '''创建网络结构'''
        super(FeatureMap_convolution, self).__init__()

        # #卷积层1
        in_channels, out_channels, kernel_size, stride, padding, dilation, bias = 3, 64, 3, 2, 1, 1, False
        self.cbnr_1 = conv2DBatchNormRelu(
            in_channels, out_channels, kernel_size, stride, padding, dilation, bias)

        #卷积层2
        in_channels, out_channels, kernel_size, stride, padding, dilation, bias = 64, 64, 3, 1, 1, 1, False
        self.cbnr_2 = conv2DBatchNormRelu(
            in_channels, out_channels, kernel_size, stride, padding, dilation, bias)

        #卷积层3
        in_channels, out_channels, kernel_size, stride, padding, dilation, bias = 64, 128, 3, 1, 1, 1, False
        self.cbnr_3 = conv2DBatchNormRelu(
            in_channels, out_channels, kernel_size, stride, padding, dilation, bias)

        #最大池化层
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

    def forward(self, x):
        x = self.cbnr_1(x)
        x = self.cbnr_2(x)
        x = self.cbnr_3(x)
        outputs = self.maxpool(x)
        return outputs


class ResidualBlockPSP(nn.Sequential):
    def __init__(self, n_blocks, in_channels, mid_channels, out_channels, stride, dilation):
        super(ResidualBlockPSP, self).__init__()

        #设置bottleNeckPSP
        self.add_module(
            "block1",
            bottleNeckPSP(in_channels, mid_channels,
                          out_channels, stride, dilation)
        )

        #循环设置bottleNeckIdentifyPSP
        for i in range(n_blocks - 1):
            self.add_module(
                "block" + str(i+2),
                bottleNeckIdentifyPSP(
                    out_channels, mid_channels, stride, dilation)
            )

class conv2DBatchNorm(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, bias):
        super(conv2DBatchNorm, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels,
                              kernel_size, stride, padding, dilation, bias=bias)
        self.batchnorm = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        outputs = self.batchnorm(x)

        return outputs

class bottleNeckPSP(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels, stride, dilation):
        super(bottleNeckPSP, self).__init__()

        self.cbr_1 = conv2DBatchNormRelu(
            in_channels, mid_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)
        self.cbr_2 = conv2DBatchNormRelu(
            mid_channels, mid_channels, kernel_size=3, stride=stride, padding=dilation, dilation=dilation, bias=False)
        self.cb_3 = conv2DBatchNorm(
            mid_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)

        #跳跃链接
        self.cb_residual = conv2DBatchNorm(
            in_channels, out_channels, kernel_size=1, stride=stride, padding=0, dilation=1, bias=False)

        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        conv = self.cb_3(self.cbr_2(self.cbr_1(x)))
        residual = self.cb_residual(x)
        return self.relu(conv + residual)

class bottleNeckIdentifyPSP(nn.Module):
    def __init__(self, in_channels, mid_channels, stride, dilation):
        super(bottleNeckIdentifyPSP, self).__init__()

        self.cbr_1 = conv2DBatchNormRelu(
            in_channels, mid_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)
        self.cbr_2 = conv2DBatchNormRelu(
            mid_channels, mid_channels, kernel_size=3, stride=1, padding=dilation, dilation=dilation, bias=False)
        self.cb_3 = conv2DBatchNorm(
            mid_channels, in_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        conv = self.cb_3(self.cbr_2(self.cbr_1(x)))
        residual = x
        return self.relu(conv + residual)
Pyramid Pooling池化
class PyramidPooling(nn.Module):
    def __init__(self, in_channels, pool_sizes, height, width):
        super(PyramidPooling, self).__init__()

        #在forward中使用的图像尺寸
        self.height = height
        self.width = width

        #各个卷积层输出的通道数
        out_channels = int(in_channels / len(pool_sizes))

       #生成每个卷积层
        # 该实现方法非常"耿直"，虽然笔者很想用 for 循环来编写这段代码，但最后还是决定优先以容易理解的方                     式编写
        # pool_sizes: [6, 3, 2, 1]
        self.avpool_1 = nn.AdaptiveAvgPool2d(output_size=pool_sizes[0])
        self.cbr_1 = conv2DBatchNormRelu(
            in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)

        self.avpool_2 = nn.AdaptiveAvgPool2d(output_size=pool_sizes[1])
        self.cbr_2 = conv2DBatchNormRelu(
            in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)

        self.avpool_3 = nn.AdaptiveAvgPool2d(output_size=pool_sizes[2])
        self.cbr_3 = conv2DBatchNormRelu(
            in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)

        self.avpool_4 = nn.AdaptiveAvgPool2d(output_size=pool_sizes[3])
        self.cbr_4 = conv2DBatchNormRelu(
            in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)

    def forward(self, x):

        out1 = self.cbr_1(self.avpool_1(x))
        out1 = F.interpolate(out1, size=(
            self.height, self.width), mode="bilinear", align_corners=True)

        out2 = self.cbr_2(self.avpool_2(x))
        out2 = F.interpolate(out2, size=(
            self.height, self.width), mode="bilinear", align_corners=True)

        out3 = self.cbr_3(self.avpool_3(x))
        out3 = F.interpolate(out3, size=(
            self.height, self.width), mode="bilinear", align_corners=True)

        out4 = self.cbr_4(self.avpool_4(x))
        out4 = F.interpolate(out4, size=(
            self.height, self.width), mode="bilinear", align_corners=True)

        #最后将结果进行合并，指定dim=1按通道数的维数进行合并
        output = torch.cat([x, out1, out2, out3, out4], dim=1)

        return output

Decoder和 AuxLoss的实现

class DecodePSPFeature(nn.Module):
    def __init__(self, height, width, n_classes):
        super(DecodePSPFeature, self).__init__()

        #在forward中使用的图像尺寸
        self.height = height
        self.width = width

        self.cbr = conv2DBatchNormRelu(
            in_channels=4096, out_channels=512, kernel_size=3, stride=1, padding=1, dilation=1, bias=False)
        self.dropout = nn.Dropout2d(p=0.1)
        self.classification = nn.Conv2d(
            in_channels=512, out_channels=n_classes, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        x = self.cbr(x)
        x = self.dropout(x)
        x = self.classification(x)
        output = F.interpolate(
            x, size=(self.height, self.width), mode="bilinear", align_corners=True)

        return output

class AuxiliaryPSPlayers(nn.Module):
    def __init__(self, in_channels, height, width, n_classes):
        super(AuxiliaryPSPlayers, self).__init__()

       #在forward中使用的图像尺寸
        self.height = height
        self.width = width

        self.cbr = conv2DBatchNormRelu(
            in_channels=in_channels, out_channels=256, kernel_size=3, stride=1, padding=1, dilation=1, bias=False)
        self.dropout = nn.Dropout2d(p=0.1)
        self.classification = nn.Conv2d(
            in_channels=256, out_channels=n_classes, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        x = self.cbr(x)
        x = self.dropout(x)
        x = self.classification(x)
        output = F.interpolate(
            x, size=(self.height, self.width), mode="bilinear", align_corners=True)

        return output

# 定义模型
net = PSPNet(n_classes=21)
net
 #生成伪数据
batch_size = 2
dummy_img = torch.rand(batch_size, 3, 475, 475)

#计算
outputs = net(dummy_img)
print(outputs)
训练
# 导入软件包
import random
import math
import time
import pandas as pd
import numpy as np

import torch
import torch.utils.data as data
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
import torch.optim as optim

# 初始设定
# Setup seeds
torch.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)

from utils.dataloader import make_datapath_list, DataTransform, VOCDataset

# 创建文件路径列表
rootpath = "./data/VOCdevkit/VOC2012/"
train_img_list, train_anno_list, val_img_list, val_anno_list = make_datapath_list(
    rootpath=rootpath)

# Dataset作成
#(RGB) 颜色的平均值和均方差
color_mean = (0.485, 0.456, 0.406)
color_std = (0.229, 0.224, 0.225)

train_dataset = VOCDataset(train_img_list, train_anno_list, phase="train", transform=DataTransform(
    input_size=475, color_mean=color_mean, color_std=color_std))

val_dataset = VOCDataset(val_img_list, val_anno_list, phase="val", transform=DataTransform(
    input_size=475, color_mean=color_mean, color_std=color_std))

#生成DataLoader
batch_size = 8

train_dataloader = data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True)

val_dataloader = data.DataLoader(
    val_dataset, batch_size=batch_size, shuffle=False)

#集中保存到字典型变量中
dataloaders_dict = {"train": train_dataloader, "val": val_dataloader}


from utils.pspnet import PSPNet

# 制作网络模型
#使用ADE20K数据集中事先训练好的模型，ADE20K的分类数量是150
net = PSPNet(n_classes=150)

#载入 ADE20K 中事先训练好的参数
state_dict = torch.load("./weights/pspnet50_ADE20K.pth")
net.load_state_dict(state_dict)

#将分类用的卷积层替换为输出数量为21的卷积层
n_classes = 21
net.decode_feature.classification = nn.Conv2d(
    in_channels=512, out_channels=n_classes, kernel_size=1, stride=1, padding=0)

net.aux.classification = nn.Conv2d(
    in_channels=256, out_channels=n_classes, kernel_size=1, stride=1, padding=0)

#对替换的卷积层进行初始化。由于激励函数是Sigmoid，因此使用Xavier进行初始化


def weights_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.xavier_normal_(m.weight.data)
        if m.bias is not None:  #如果bias存在
            nn.init.constant_(m.bias, 0.0)


net.decode_feature.classification.apply(weights_init)
net.aux.classification.apply(weights_init)


print('网络设置完毕 ：成功的载入了事先训练完毕的权重')


#设置损失函数
class PSPLoss(nn.Module):
    """#设置损失函数"""

    def __init__(self, aux_weight=0.4):
        super(PSPLoss, self).__init__()
        self.aux_weight = aux_weight  #aux_loss的权重

    def forward(self, outputs, targets):
        """
        损失函数的计算。

        Parameters
        ----------
        outputs : PSPNet的输出(tuple)
            (output=torch.Size([num_batch, 21, 475, 475]), output_aux=torch.Size([num_batch, 21, 475, 475]))。

        targets : [num_batch, 475, 4755]
            正解的标注信息

        Returns
        -------
        loss :张量
            损失值
        """

        loss = F.cross_entropy(outputs[0], targets, reduction='mean')
        loss_aux = F.cross_entropy(outputs[1], targets, reduction='mean')

        return loss+self.aux_weight*loss_aux


criterion = PSPLoss(aux_weight=0.4)


# #由于使用的是微调，因此要降低学习率く
optimizer = optim.SGD([
    {'params': net.feature_conv.parameters(), 'lr': 1e-3},
    {'params': net.feature_res_1.parameters(), 'lr': 1e-3},
    {'params': net.feature_res_2.parameters(), 'lr': 1e-3},
    {'params': net.feature_dilated_res_1.parameters(), 'lr': 1e-3},
    {'params': net.feature_dilated_res_2.parameters(), 'lr': 1e-3},
    {'params': net.pyramid_pooling.parameters(), 'lr': 1e-3},
    {'params': net.decode_feature.parameters(), 'lr': 1e-2},
    {'params': net.aux.parameters(), 'lr': 1e-2},
], momentum=0.9, weight_decay=0.0001)


# #设置调度器
def lambda_epoch(epoch):
    max_epoch = 30
    return math.pow((1-epoch/max_epoch), 0.9)


scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_epoch)

#创建对模型进行训练的函数


def train_model(net, dataloaders_dict, criterion, scheduler, optimizer, num_epochs):

    #确认GPU是否可用
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("使用的设备 ：", device)

    #将网络载入GPU中
    net.to(device)

    #如果网络相对固定，开启高速处理选项
    torch.backends.cudnn.benchmark = True

    #图像的张数
    num_train_imgs = len(dataloaders_dict["train"].dataset)
    num_val_imgs = len(dataloaders_dict["val"].dataset)
    batch_size = dataloaders_dict["train"].batch_size

    #设置迭代计数器
    iteration = 1
    logs = []

    # multiple minibatch
    batch_multiplier = 3

    # epochのループ
    for epoch in range(num_epochs):

        #epoch的循环
        t_epoch_start = time.time()
        t_iter_start = time.time()
        epoch_train_loss = 0.0  #epoch的损失和
        epoch_val_loss = 0.0  #epoch的损失和

        print('-------------')
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-------------')

        #对每轮epoch进行训练和验证的循环
        for phase in ['train', 'val']:
            if phase == 'train':
                net.train()  #将模式设置为训练模式
                scheduler.step()  #更新最优化调度器
                optimizer.zero_grad()
                print('（train）')

            else:
                if((epoch+1) % 5 == 0):
                    net.eval()   #将模型设置为验证模式
                    print('-------------')
                    print('（val）')
                else:
                    #每5轮进行1次验证
                    continue

           #从数据加载器中读取每个小批量并进行循环
            count = 0  # multiple minibatch
            for imges, anno_class_imges in dataloaders_dict[phase]:
                #如果小批量的尺寸是1，批量正规化处理会报错，因此需要避免
                if imges.size()[0] == 1:
                    continue

                #如果GPU可用，将数据传输到GPU中
                imges = imges.to(device)
                anno_class_imges = anno_class_imges.to(device)

                
                #使用multiple minibatch对参数进行更新
                if (phase == 'train') and (count == 0):
                    optimizer.step()
                    optimizer.zero_grad()
                    count = batch_multiplier

                #正向传播计算
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = net(imges)
                    loss = criterion(
                        outputs, anno_class_imges.long()) / batch_multiplier

                    #训练时采用反向传播
                    if phase == 'train':
                        loss.backward() #梯度的计算
                        count -= 1  # multiple minibatch

                        if (iteration % 10 == 0):  #每10次迭代显示一次loss
                            t_iter_finish = time.time()
                            duration = t_iter_finish - t_iter_start
                            print('迭代 {} || Loss: {:.4f} || 10iter: {:.4f} sec.'.format(
                                iteration, loss.item()/batch_size*batch_multiplier, duration))
                            t_iter_start = time.time()

                        epoch_train_loss += loss.item() * batch_multiplier
                        iteration += 1

                    #验证时
                    else:
                        epoch_val_loss += loss.item() * batch_multiplier

        #每个epoch的phase的loss和正解率
        t_epoch_finish = time.time()
        print('-------------')
        print('epoch {} || Epoch_TRAIN_Loss:{:.4f} ||Epoch_VAL_Loss:{:.4f}'.format(
            epoch+1, epoch_train_loss/num_train_imgs, epoch_val_loss/num_val_imgs))
        print('timer:  {:.4f} sec.'.format(t_epoch_finish - t_epoch_start))
        t_epoch_start = time.time()

        #保存日志
        log_epoch = {'epoch': epoch+1, 'train_loss': epoch_train_loss /
                     num_train_imgs, 'val_loss': epoch_val_loss/num_val_imgs}
        logs.append(log_epoch)
        df = pd.DataFrame(logs)
        df.to_csv("log_output.csv")

    #保存最终的网络
    torch.save(net.state_dict(), 'weights/pspnet50_' +
               str(epoch+1) + '.pth')

# 学习和验证的实现
num_epochs = 30
train_model(net, dataloaders_dict, criterion, scheduler, optimizer, num_epochs=num_epochs)

至此， 网络搭建完成，CPU 训练太慢了，5.1 假期找个 GPU 服务器训练试试，然后上结果。

相关推荐
千天夜2 分钟前
使用UDP协议传输视频流！（分片、缓存）
python·网络协议·udp·视频流
Debroon3 分钟前
RuleAlign 规则对齐框架：将医生的诊断规则形式化并注入模型，无需额外人工标注的自动对齐方法
人工智能
测试界的酸菜鱼6 分钟前
Python 大数据展示屏实例
大数据·开发语言·python
羊小猪~~9 分钟前
神经网络基础--什么是正向传播？？什么是方向传播？？
人工智能·pytorch·python·深度学习·神经网络·算法·机器学习
AI小杨11 分钟前
【车道线检测】一、传统车道线检测：基于霍夫变换的车道线检测史诗级详细教程
人工智能·opencv·计算机视觉·霍夫变换·车道线检测
晨曦_子画15 分钟前
编程语言之战：AI 之后的 Kotlin 与 Java
android·java·开发语言·人工智能·kotlin
道可云17 分钟前
道可云人工智能&元宇宙每日资讯｜2024国际虚拟现实创新大会将在青岛举办
大数据·人工智能·3d·机器人·ar·vr
人工智能培训咨询叶梓26 分钟前
探索开放资源上指令微调语言模型的现状
人工智能·语言模型·自然语言处理·性能优化·调优·大模型微调·指令微调
zzZ_CMing26 分钟前
大语言模型训练的全过程：预训练、微调、RLHF
人工智能·自然语言处理·aigc
newxtc27 分钟前
【旷视科技-注册/登录安全分析报告】
人工智能·科技·安全·ddddocr
热门推荐
01【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总02(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明03组基轨迹建模 GBTM的介绍与实现（Stata 或 R）04【AIGC】重塑未来的科技巨轮05全面解析：构建基于深度学习的安全帽检测系统（UI界面+YOLO代码+数据集）06【经验分享】Ubuntu22.04安装微信(linux官方版)07基于YOLOv10深度学习的CT扫描图像肾结石智能检测系统【python源码+Pyqt5界面+数据集+训练代码】深度学习实战、目标检测08Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO09RAG 实践- Ollama+RagFlow 部署本地知识库10红米手机使用google play