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 服务器训练试试，然后上结果。

相关推荐
白光白光几秒前
量子神经网络
人工智能·深度学习·神经网络
GIS 数据栈5 分钟前
每日一书 《基于ArcGIS的Python编程秘笈》
开发语言·python·arcgis
全域观察5 分钟前
如何复制只读模式下的腾讯文档
人工智能·新媒体运营·媒体·内容运营·程序员创富
爱分享的码瑞哥9 分钟前
Python爬虫中的IP封禁问题及其解决方案
爬虫·python·tcp/ip
panpantt32113 分钟前
【参会邀请】第二届大数据与数据挖掘国际会议（BDDM 2024）邀您相聚江城！
大数据·人工智能·数据挖掘
lindsayshuo22 分钟前
jetson orin系列开发版安装cuda的gpu版本的opencv
人工智能·opencv
向阳逐梦22 分钟前
ROS机器视觉入门：从基础到人脸识别与目标检测
人工智能·目标检测·计算机视觉
陈鋆1 小时前
智慧城市初探与解决方案
人工智能·智慧城市
qdprobot1 小时前
ESP32桌面天气摆件加文心一言AI大模型对话Mixly图形化编程STEAM创客教育
网络·人工智能·百度·文心一言·arduino
QQ39575332371 小时前
金融量化交易模型的突破与前景分析
人工智能·金融
热门推荐
01(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明02PyTorch机器学习实现液态神经网络03【HarmonyOS】HUAWEI DevEco Studio 下载地址汇总04玄机平台应急响应—webshell查杀05Coze扣子平台完整体验和实践（附国内和国际版对比）06Ubuntu 20.04使用Livox mid 360 测试 FAST_LIO07【目标跟踪】相机运动补偿08Unity中PICO实现 隔空取物 和 接触抓取物体09RAG 实践- Ollama+RagFlow 部署本地知识库10怎样让音频速度变慢？请跟随以下方法进行操作