pytorch快速入门与实战——四、网络训练与测试

专栏目录：pytorch（图像分割UNet）快速入门与实战——零、前言
pytorch快速入门与实战——一、知识准备（要素简介）
pytorch快速入门与实战——二、深度学习经典网络发展
pytorch快速入门与实战——三、Unet实现
pytorch快速入门与实战——四、网络训练与测试

接上文：pytorch快速入门与实战——三、Unet实现
网络实现后读取数据与参数回传。

语义分割实现流程
训练：

根据batch size大小，将数据集中的训练样本和标签读入卷积神经网络。根据实际需要，应先对训练图片及标签进行预处理，如裁剪、数据增强等。这有利于深层网络的的训练，加速收敛过程，同时也避免过拟合问题并增强了模型的泛化能力。

验证：

训练一个epoch结束后，将数据集中的验证样本和标签 读入卷积神经网络，并载入训练权重。根据编写好的语义分割指标进行验证，得到当前训 练过程中的指标分数，保存对应权重。常用一次训练一次验 证的方法更好的监督模型表现。

测试：

所有训练结束后，将数据集中的测试样本和标签读入卷积神经网络，并将保存的最好权重值载入模型，进行测试。 测试结果分为两种，一种是根据常用指标分数衡量网络性能，另一种是将网络的预测结果以 图片的形式保存下来，直观感受分割的精确程度。

1.读取数据

如果你不是做分割，而是分类问题，看这里或者去看其他文章：

ImageFolder
在PyTorch中有一个现成实现的数据读取方法，是torchvision.datasets.ImageFolder，这个api是仿照keras写的，主要是做分类问题，将每一类数据放到同一个文件夹中，比如有10个类别，那么就在一个大的文件夹下面建立10个子文件夹，每个子文件夹里面放的是同一类的数据。

图像分割看下面：

1.1 继承Dataset类

PyTorch读取图片，主要是通过torch.utils.data.Dataset类，导包手法为：

from torch.utils.data import Dataset

然后继承这个类来实现我们自己的数据读取类：非常的简洁
主要是两个方法：类实现方法__init__和取元素方法__getitem__

class myImageDataset(Dataset):
    def __init__(self, inputs_root, labels_root):
        self.files = sorted(glob.glob(f"{inputs_root}\\*.png"))
        self.files_using = sorted(glob.glob(f"{labels_root}\\*.png"))

    def __getitem__(self, index):
        inputs = plt.imread(self.files[index % len(self.files)]) 
        labels = plt.imread(self.files_using[index % len(self.files_using)])
        return inputs, labels

    def __len__(self):
        return len(self.files)

简单解读下：

1. init

inputs_root就是输入数据的根目录路径，labels_root就是标签的路径。
glob是Python自带的方法，glob.glob可以遍历当前目录的文件，上面语句是遍历目录下所有png文件，其中sorted是为了input和label一一对应不乱序。
init的self不用管，传参的时候直接跳过就行，

2. get item

index是一个自增的索引，这个不用管，是方法内部操作。
借着索引取路径，然后读取图片再返回出去就完事了。

1.2 加载器 DataLoader方法

Dataset读取了图片之后如何人工智能似的去按批次处理呢？
pytorch也提供了一个方法：torch.utils.data类里的DataLoader方法。

from torch.utils.data import DataLoader

train_loader = DataLoader(
    dataset=myImageDataset(inputs_root=in_folder + r"\train\inputs",
                         labels_root=in_folder + r"\train\labels"),
    batch_size=16,  # 一批有几个，一般为2的指数
    shuffle=True,
    num_workers=8  # 使用的CPU核心数量
)

通过给我们的dataset塞路径，取数据集，让dataloader去按批次处理，一批处理多少个由batch_size决定，线程数为由参数num_worker决定，其他参数自己去查阅和研究就行。

测试集同理：

test_loader = DataLoader(
    dataset=myImageDataset(inputs_root=in_folder + r"\test\inputs",
                         labels_root=in_folder + r"\test\labels"),
    batch_size=1,
    shuffle=True,
    num_workers=opt.n_cpu
)

2. 其他初始化

2.1 全局变量设置

定义训练批次，每批次多少数据，定义学习率等参数（根据需要自己设置，我复制的师兄的）：

parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=100, help="number of epochs of training")
# parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=16, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=4, help="number of cpu threads to use during batch generation")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=100, help="interval between saving image samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between model checkpoints")
opt = parser.parse_args()

2.2 设置网络、loss函数、优化器optimizer、张量转换tensor

前面莫烦python讲过这个优化器，这里使用adam。
至于张量：张量是多重线性函数，矩阵是张量在一组特定基矢下的表示。
其他的自己去搜吧，也可以看这一篇：浅谈什么是张量tensor

# Initialize net
net = AdUNet()
# Losses
loss = torch.nn.L1Loss()

# 数据传给显卡？
cuda = torch.cuda.is_available()
if cuda:
    net = net.cuda()
    gloss = loss.cuda()

if opt.epoch != 0:
    # Load pretrained models
    net.load_state_dict(torch.load("../saved_models/AdUnet_%d.pkl"))

# Optimizers
optimizer = torch.optim.Adam(net.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))

Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor

2.3 训练

训练大致流程：梯度清零，反向传播，更新学习率

optimizer.zero_grad()  # 梯度归零：step之前要进行梯度归零
loss.backward()  # 进行反向传播求出每个参数的梯度
optimizer.step()  # 更新学习率

具体流程：
插个说明：emmm下方代码中，我从train_loader加载的数据集data里面除了输入inputs和labels还有个路径inputs_labels对不对，这是因为我后面想生成结果图来看，为了给结果图命名设置的一个参数，所以前面我的那个dataset里面的__getitem__也有点不一样，呐这是我的（其实就是取了个路径传出去罢了）：

    def __getitem__(self, index):
        inputs_path = self.files[index % len(self.files)]
        inputs = plt.imread(inputs_path)
        labels = plt.imread(self.files_using[index % len(self.files_using)])
        inputs_name = inputs_path.split("\\")[-1]
        return inputs, labels, inputs_name

2.3.1 读取训练数据，转换训练数据

    for i, data in enumerate(train_loader):
        # 一个batch
        inputs, labels, inputs_path = data
        inputs = inputs.unsqueeze(1)
        labels = labels.unsqueeze(1)
        # 将这些数据转换成Variable类型
        inputs, labels = Variable(inputs), Variable(labels)
        device = torch.device("cuda" if cuda else "cpu")
        inputs = inputs.to(device)
        labels = labels.to(device)

我的灰度图读取进来是一个二维的，也就是（120,240），加上一个batch_size也才三维，我的batchsize为16，所以我读取的inputs是(16,120,240)，但是实际上网络跑起来是四维的，如何做呢？
实际上灰度图是单通道的，也就是(120,240,1)，所以我们要手动给这个通道1加上，在我的调试过程中发现，网络中的通道一般是放在第二位的，所以我们要把他变成(16,1,120,240)，通过unsqueeze解压缩就可以直接实现，要在第二位加就直接inputs.unsqueeze(1)即可。同理第一位加就是inputs.unsqueeze(0)

inputs和labels都转换好了之后要转换成variable类型才能进行反向传播（原因我不知道，照搬的），然后用显卡就把数据传给显卡，不用显卡就传给CPU。

2.3.2 训练网络

• 梯度归零（一开始就是0，为了下一波循环）
• 网络训练
• 得到结果netout
• 计算损失函数loss
• 反向传播
• 更新学习率
  不需要懂，抄就行了。

        optimizer.zero_grad()  # 梯度归零：step之前要进行梯度归零
        net.train()
        netout = net(inputs)

        # Total loss
        gloss = loss(netout, labels)

        gloss.backward()  # 进行反向传播求出每个参数的梯度
        optimizer.step()  # 更新学习率

2.4 测试

初始化评价标准ssim和psnr和rmse。

    total_s = 0  # ssim
    total_p = 0  # psnr
    total_r = 0  # rmse

skimage自带这些库的：

from skimage.measure import compare_ssim as ssim
from skimage.measure import compare_psnr as psnr
from skimage.measure import compare_mse as mse

• 读取测试数据
• 统一数据格式
• Variable转换
• 送到显卡
• 标明是测试（不会参与反向传播，和第一行作用一样，直接加就行）
• 送到网络得出结果netout
• 降到二维（灰度图是二维）
• 与label相比评价图片
• over

with torch.no_grad():
        for inputs, labels, inputs_path in test_loader:
            # 一个batch
            inputs = inputs.unsqueeze(1)
            # labels = labels.unsqueeze(0)
            # 将这些数据转换成Variable类型
            inputs, labels = Variable(inputs), Variable(labels)
            device = torch.device("cuda" if cuda else "cpu")
            inputs = inputs.to(device)
            labels = labels.squeeze(0).cpu().numpy()

            # optimizer.zero_grad()
            net.eval()
            netout = net(inputs)
            img_out = netout.squeeze(1)
            img_out = img_out.squeeze(0)
            img_out = img_out.cpu().numpy()

            # Total loss
            # gloss = loss(netout, labels)
            # print(netout.shape)
            s = ssim(labels, img_out)
            p = psnr(labels, img_out)
            r = sqrt(mse(labels, img_out))

            total_s += float(s.item())
            total_p += float(p.item())
            total_r += float(r)

3 整体代码：

import argparse
import os
import sys

import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader

from main.AdUNet import AdUNet
from main.datasets import *

from skimage.measure import compare_ssim as ssim
from skimage.measure import compare_psnr as psnr
from skimage.measure import compare_mse as mse
from math import sqrt

from other.mkdir import mkdir

mkdir("../saved_models")
parser = argparse.ArgumentParser()
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=100, help="number of epochs of training")
# parser.add_argument("--dataset_name", type=str, default="img_align_celeba", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=16, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=4, help="number of cpu threads to use during batch generation")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument("--sample_interval", type=int, default=100, help="interval between saving image samples")
parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between model checkpoints")
opt = parser.parse_args()

in_folder = r"..\data\pigs"
train_folder = in_folder + "\\train"

# Initialize net
net = AdUNet()
# Losses
loss = torch.nn.L1Loss()

# 数据传给显卡？
cuda = torch.cuda.is_available()
if cuda:
    net = net.cuda()
    gloss = loss.cuda()

if opt.epoch != 0:
    # Load pretrained models
    net.load_state_dict(torch.load("../saved_models/AdUnet_%d.pkl"))

# Optimizers
optimizer = torch.optim.Adam(net.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))

Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor

train_loader = DataLoader(
    dataset=ImageDataset(inputs_root=in_folder + r"\train\inputs",
                         labels_root=in_folder + r"\train\labels"),
    batch_size=opt.batch_size,
    shuffle=True,
    num_workers=opt.n_cpu
)
test_loader = DataLoader(
    dataset=ImageDataset(inputs_root=in_folder + r"\test\inputs",
                         labels_root=in_folder + r"\test\labels"),
    batch_size=1,
    shuffle=True,
    num_workers=opt.n_cpu
)


# ----------
#  Training
# ----------
def train(epoch):
    for i, data in enumerate(train_loader):
        # 一个batch
        inputs, labels, inputs_path = data
        inputs = inputs.unsqueeze(1)
        labels = labels.unsqueeze(1)
        # 将这些数据转换成Variable类型
        inputs, labels = Variable(inputs), Variable(labels)
        device = torch.device("cuda" if cuda else "cpu")
        inputs = inputs.to(device)
        labels = labels.to(device)
        # ------------------
        #  Train net
        # ------------------

        optimizer.zero_grad()  # 梯度归零：step之前要进行梯度归零
        net.train()
        netout = net(inputs)

        # Total loss
        gloss = loss(netout, labels)

        gloss.backward()  # 进行反向传播求出每个参数的梯度
        optimizer.step()  # 更新学习率
##
        # --------------
        #  Log Progress
        # --------------

        sys.stdout.write(
            "\n[Epoch %d/%d] [Batch %d/%d] [D loss: %f]"
            % (epoch, opt.n_epochs, i, len(train_loader), gloss.item())
        )
        with open(r"../data/pigs/loss.txt", "a") as f1:
            f1.write("\n[Epoch %d/%d] [Batch %d/%d] [D loss: %f]"
                     % (epoch, opt.n_epochs, i, len(train_loader), gloss.item()))
            f1.close()


def test():
    total_s = 0  # ssim
    total_p = 0  # psnr
    total_r = 0  # rmse

    with torch.no_grad():
        for inputs, labels, inputs_path in test_loader:
            # 一个batch
            inputs = inputs.unsqueeze(1)
            # labels = labels.unsqueeze(0)
            # 将这些数据转换成Variable类型
            inputs, labels = Variable(inputs), Variable(labels)
            device = torch.device("cuda" if cuda else "cpu")
            inputs = inputs.to(device)
            labels = labels.squeeze(0).cpu().numpy()

            # optimizer.zero_grad()
            net.eval()
            netout = net(inputs)
            img_out = netout.squeeze(1)
            img_out = img_out.squeeze(0)
            img_out = img_out.cpu().numpy()

            # Total loss
            # gloss = loss(netout, labels)
            # print(netout.shape)
            s = ssim(labels, img_out)
            p = psnr(labels, img_out)
            r = sqrt(mse(labels, img_out))

            total_s += float(s.item())
            total_p += float(p.item())
            total_r += float(r)
    print('\n|Epoch %d/%d| |Average SSIM: %f | |Average PSNR: %f| |Average RMSE: %f| '
          % ((epoch + 1), opt.n_epochs, total_s / len(test_loader), total_p / len(test_loader),
             total_r / len(test_loader)))

    with open(r"../data/pigs/test_parameters.txt", "a") as f:
        f.write("|Epoch %d/%d| |Average SSIM: %f | |Average PSNR: %f| |Average RMSE: %f| \r\n"
                % ((epoch + 1), opt.n_epochs, total_s / len(test_loader), total_p / len(test_loader),
                   total_r / len(test_loader)))
        f.close()
    return total_s / len(test_loader), total_p / len(test_loader), total_r / len(test_loader)


if __name__ == '__main__':
    for epoch in range(opt.epoch, opt.n_epochs):
        # 一个epoch
        train(epoch)
        # Save model checkpoints
        # torch.save(net.state_dict(), "saved_models/generator_%d.pkl" % epoch)
        torch.save(net, "../saved_models/AdUnet_%d.pkl" % epoch)
        epoch_s, epoch_p, epoch_r = test()

再加上前面的net类和dataset类，一共三个文件，运行这个文件就行，大功告成！
“呵呵，就这”