Pytorch：图像语义分割-基于VGG19的FCN8s实现
Pytorch: 图像语义分割-基于VGG19的FCN8s语义分割网络实现
Copyright: Jingmin Wei, Pattern Recognition and Intelligent System, School of Artificial and Intelligence, Huazhong University of Science and Technology
Pytorch教程专栏链接
文章目录
-
-
- Pytorch: 图像语义分割-基于VGG19的FCN8s语义分割网络实现
-
- 数据准备
- 数据集和数据加载器构建
- 可视化预处理效果
- FCN网络搭建
- 网络模型训练与验证
-
本教程不商用，仅供学习和参考交流使用，如需转载，请联系本人。
上节介绍的是基于 ResNet-101 网络搭建的 FCN 模型。这节将基于 VGG19 网络，搭建、训练和测试自己的图像全卷积语义分割网络。
使用 VOC2012 对网络进行训练和验证。每个数据集约有 1000 1000 1000 张图片，并且图像之间的尺寸不完全相同，数据集有 21 21 21 类需要学习的目标类别。
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import PIL
import copy
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as Data
from torchvision import transforms
from torchvision.models import vgg19 # FCN的backbone
训练前，由于网络算力要求较高，我们将代码和数据集放到服务器端。
举个例子：首先打开服务器，之后在 VSCode 打开远程资源管理器，使用 ssh 命令连接。
ssh yourserver # yourserver为你自己的服务器的地址
在 Windows Terminal 中 scp 传输代码文件和压缩数据集，数据集传输完后还需要解压
# Terminal端操作
scp "E:\\Jupyter WorkSpace\\PytorchLearn\\torch_ssn_segmentation.ipynb" yourserver:/home/mist/ssn
scp "E:\\Jupyter WorkSpace\\PytorchLearn\\VOC2012.rar" yourserver:/home/mist/ssn/data
# Server端操作
cd ~/ssn/data
unrar x VOC2012.rar
# 模型加载选择GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
print(torch.cuda.device_count())
print(torch.cuda.get_device_name(0))
cuda
1
NVIDIA GeForce RTX 3080
数据准备
首先需要去网上下载 VOC2012 数据集。
定义标注的每类对应的名称和颜色：
# 物体对应类别
classes = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair',
'cow', 'dining table', 'dog', 'horse', 'motorbike',
'person', 'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
# 对应类别的RGB值
colormap = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
[128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0],
[64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128],
[192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0], [0, 64, 128]]
数据预处理需要对每张图像做如下操作：
-
将原始图像和标机图像所对应的图片路径一一对应。
-
将图像统一切分为固定尺寸时，需要保持原始图像和其对应的标记好的图像在切分后，每个像素也仍然是一一对应的，所以需要对原始图像和目标的标记图像从相同位置进行切分。因此，在切分之前还需要过滤掉尺寸小于给定切分尺寸的图像。
-
对原始图像进行数据标准化。
-
针对标记好的 RGB 图像，将 RGB 的值对应的类重新定义，把 3D 的 RGB 图像转化为一个二维数据，并且数组中每个位置的取值对应着图像在该像素点的类别。
完成上述预处理操作，需要定义下面几个辅助函数：
image2label 将一张标记好的图像 y 转化为类别标签图像。该函数对应于上面的第 4 4 4 步。
# 给定一个标好的图片y，将像素值对应的物体类别找出来，转为二维标签矩阵
def image2label(image, colormap):
# 将标签转化为每个像素值为一类数据
cm2lbl = np.zeros(256**3)
for i, cm in enumerate(colormap):
cm2lbl[ (cm[0]*256 + cm[1]*256 + cm[2]) ] = i
# 对一张图像进行转换
image = np.array(image, dtype='int64')
ix = (image[:, :, 0]*256 + image[:, :, 1]*256 + image[:, :, 2])
image2 = cm2lbl[ix]
return image2
rand_crop 函数对应第 2 2 2 步，完成对图像 X 和标签图像 y 随机裁剪，随机裁剪后原图像和标签图像的每个像素一一对应。high 和 width 用来指定图像裁剪后的高和宽。
# 随机裁剪数据X和对应标好的图y
def rand_crop(data, label, high, width):
im_width, im_high = data.size
# 生成图像随机点的位置
left = np.random.randint(0, im_width-width)
top = np.random.randint(0, im_high-high)
right = left + width
bottom = top + high
data = data.crop([left, top, right, bottom])
label = label.crop([left, top, right, bottom])
return data, label
img_transforms 对应第 3 3 3 步。定义了对数据 X 和标签 y 做随机裁剪，数据 X 的标准化，以及对应标签 y 转为二维标签图的处理。最后输出原始图像 X 和类别标签 y 的张量数据类型。
# 一个batch图像的转换操作
def img_transforms(data, label, high, width, colormap):
data, label = rand_crop(data, label, high, width) # 调用数据X和标签y随机裁剪
data_augment = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485,0.456,0.406],
std=[0.229,0.224,0.225])])
data = data_augment(data) # 数据X变换
label = torch.from_numpy(image2label(label, colormap)) # 标签y转为二维标签图
return data, label
read_image_path 对应第 1 1 1 步，从给定的文件路中定义出对应的原始图像和标记好的目标图像的存储路径列表。即 data(X), label(y) 。
# 定义列出需要读取的数据路径的函数
def read_image_path(root='./data/VOC2012/ImageSets/Segmentation/train.txt'):
# 保存指定路径下的所有需要读取的图像的文件路径
image = np.loadtxt(root, dtype=str)
n = len(image)
data, label = [None] * n, [None] * n
for i, file_name in enumerate(image):
# 图像X和标签图像y的文件路径
data[i] = './data/VOC2012/JPEGImages/%s.jpg' % (file_name)
label[i] = './data/VOC2012/SegmentationClass/%s.png' % (file_name)
return data, label
数据集和数据加载器构建
为了将数据定义为数据加载器 Data.DataLoader() 方法可以接受的数据格式，定义好上述辅助函数后，需要定义一个继承类，继承自 torch.utils.data.Dataset 类，并修改定义自己的数据操作，从而得到 DataLoader 可以接受的数据格式。
这一步也是很多 PyTorch 处理数据集时的基本操作。
其中的 protected 方法 _filter 用来过滤掉图像尺寸小于固定切分尺寸的样本，对应第 2 2 2 步。
class MyDataset(Data.Dataset):
def __init__(self, data_root, high, width, im_transform, colormap):
# 初始化方法
# 文件路径, 裁剪后的宽高, 数据增强方式, 颜色图
super(MyDataset, self).__init__()
self.data_root = data_root
self.high = high
self.width = width
self.im_transform = im_transform
self.colormap = colormap
data_list, label_list = read_image_path(root=self.data_root)
self.data_list = self._filter(data_list)
self.label_list = self._filter(label_list)
def _filter(self, images):
# 定义一个protected方法，返回一个列表生成式
return [im for im in images if (PIL.Image.open(im).size[1] > self.high and
PIL.Image.open(im).size[0] > self.width)]
def __getitem__(self, index):
# 多态重写getitem方法
# Dataloader为惰性迭代器
# 迭代某张图片时才会调用该方法对某张图片进行图像增强
img_name = self.data_list[index]
label_name = self.label_list[index] # 文件名
img = PIL.Image.open(img_name)
label = PIL.Image.open(label_name).convert('RGB') # 打开对应图像
img, label = self.im_transform(img, label, self.high, self.width, self.colormap)
return img, label
def __len__(self):
return len(self.data_list) # 获得数据集大小
下面首先建立 Dataset，读取原始数据 X 和对应的标签数据 y，然后建立 DataLoader ，每个 batch 包含 8 8 8 张图像。
# 定义输入图像的高宽
high, width = 320, 480
# 读取数据，定义数据集
voc_train = MyDataset('./data/VOC2012/ImageSets/Segmentation/train.txt',
high, width, img_transforms, colormap) # 训练集
voc_val = MyDataset('./data/VOC2012/ImageSets/Segmentation/val.txt',
high, width, img_transforms, colormap) # 验证集
每个随机梯度下降时计算 8 8 8 张图的局部平均梯度( b a t c h _ s i z e = 8 batch\_size=8 batch_size=8 ) ，打乱数据，设置多进程为 8 8 8 (Windows OS 必须设为 0 0 0 ！)，并使用锁页内存加快迭代速度。
# 根据数据集，创建数据加载器
train_loader = Data.DataLoader(voc_train, batch_size=8, shuffle=True,
num_workers=8, pin_memory=True)
val_loader = Data.DataLoader(voc_val, batch_size=8, shuffle=True,
num_workers=8, pin_memory=True)
# 检查训练数据集的第一个batch的样本维度是否正确
for step, (b_x, b_y) in enumerate(train_loader):
if step > 0:
break
# 输出尺寸，以及数据类型
print('one X_train batch\'s shape:', b_x.shape)
print('one y_train batch\'s shape:', b_y.shape)
one X_train batch's shape: torch.Size([8, 3, 320, 480])
one y_train batch's shape: torch.Size([8, 320, 480])
可视化预处理效果
下面对一个 batch 的图像和其标签进行可视化，以检查数据预处理效果，可视化需要定义两个辅助函数。
inv_normalize_image 将标准化的图像进行逆标准化操作，转为能够可视化的 0 − 1 0-1 0−1 区间。
label2image 是将二维的类别标签矩阵转为三维的图像分割后的数据，是image2label的逆操作。不同的类别转化为特定的 RGB 值。
# 将标准化后的图像X转为0-1区间
def inv_normalize_image(data):
rgb_mean = np.array([0.485, 0.456, 0.406])
rgb_std = np.array([0.229, 0.224, 0.225])
data = data.astype('float32') * rgb_mean + rgb_std
return data.clip(0, 1)
# 从预测的二维标签矩阵y转为图像
def label2image(prelabel, colormap):
h, w = prelabel.shape
prelabel = prelabel.reshape(h*w, -1)
image = np.zeros((h*w, 3), dtype='int32')
for i in range(len(colormap)):
index = np.where(prelabel == i) # 标签对应的colormap索引
image[index, :] = colormap[i]
return image.reshape(h, w, 3)
# 可视化一个训练batch的图像，检查数据预处理效果
b_x_numpy = b_x.data.numpy()
b_x_numpy = b_x_numpy.transpose(0, 2, 3, 1) # rgb->bgr
b_y_numpy = b_y.data.numpy()
plt.figure(figsize=(16, 6))
for i in range(4):
plt.subplot(2, 4, i+1)
plt.imshow(inv_normalize_image(b_x_numpy[i])) # X可视化
plt.axis('off')
plt.subplot(2, 4, i+5)
plt.imshow(label2image(b_y_numpy[i], colormap)) # y可视化
plt.axis('off')
plt.subplots_adjust(wspace=0.1, hspace=0.1) # 调整间距
plt.show()
FCN网络搭建
上一节预处理 FCN 网络的 backbone 是 ResNet-101。这里考虑到训练时间和效率，搭建全卷积语义分割网络时，基础的 backbone 是 VGG19，可以直接从 torchvision 库中导入模型.
from torchsummary import summary
model_vgg19 = vgg19(pretrained=True)
# 只使用特征层，不使用平均池化层和全连接层
backbone = model_vgg19.features
summary(backbone, input_size=(3, high, width), device='cpu') # [3, 320, 480]
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /home/mist/pytorch/c10/core/TensorImpl.h:1156.)
return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 320, 480] 1,792
ReLU-2 [-1, 64, 320, 480] 0
Conv2d-3 [-1, 64, 320, 480] 36,928
ReLU-4 [-1, 64, 320, 480] 0
MaxPool2d-5 [-1, 64, 160, 240] 0
Conv2d-6 [-1, 128, 160, 240] 73,856
ReLU-7 [-1, 128, 160, 240] 0
Conv2d-8 [-1, 128, 160, 240] 147,584
ReLU-9 [-1, 128, 160, 240] 0
MaxPool2d-10 [-1, 128, 80, 120] 0
Conv2d-11 [-1, 256, 80, 120] 295,168
ReLU-12 [-1, 256, 80, 120] 0
Conv2d-13 [-1, 256, 80, 120] 590,080
ReLU-14 [-1, 256, 80, 120] 0
Conv2d-15 [-1, 256, 80, 120] 590,080
ReLU-16 [-1, 256, 80, 120] 0
Conv2d-17 [-1, 256, 80, 120] 590,080
ReLU-18 [-1, 256, 80, 120] 0
MaxPool2d-19 [-1, 256, 40, 60] 0
Conv2d-20 [-1, 512, 40, 60] 1,180,160
ReLU-21 [-1, 512, 40, 60] 0
Conv2d-22 [-1, 512, 40, 60] 2,359,808
ReLU-23 [-1, 512, 40, 60] 0
Conv2d-24 [-1, 512, 40, 60] 2,359,808
ReLU-25 [-1, 512, 40, 60] 0
Conv2d-26 [-1, 512, 40, 60] 2,359,808
ReLU-27 [-1, 512, 40, 60] 0
MaxPool2d-28 [-1, 512, 20, 30] 0
Conv2d-29 [-1, 512, 20, 30] 2,359,808
ReLU-30 [-1, 512, 20, 30] 0
Conv2d-31 [-1, 512, 20, 30] 2,359,808
ReLU-32 [-1, 512, 20, 30] 0
Conv2d-33 [-1, 512, 20, 30] 2,359,808
ReLU-34 [-1, 512, 20, 30] 0
Conv2d-35 [-1, 512, 20, 30] 2,359,808
ReLU-36 [-1, 512, 20, 30] 0
MaxPool2d-37 [-1, 512, 10, 15] 0
================================================================
Total params: 20,024,384
Trainable params: 20,024,384
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.76
Forward/backward pass size (MB): 729.49
Params size (MB): 76.39
Estimated Total Size (MB): 807.64
----------------------------------------------------------------
之前的教程我们也介绍过 VGG 的网络结构，这里可以看出，VGG19 的特征提取层通过 5 5 5 个 MaxPooling 层将图像尺寸缩小到了原来的 1 32 \frac{1}{32} 321 ，即图像使用最大值池化进行下采样，分别在 MaxPool2d-5 缩小到原来的 1 2 \frac{1}{2} 21 ，MaxPool2d-10 缩小到原来的 1 4 \frac{1}{4} 41 ，MaxPool2d-19 缩小到原来的 1 8 \frac{1}{8} 81 ，MaxPool2d-28 缩小到原来的 1 16 \frac{1}{16} 161 ，MaxPool2d-37 缩小到原来的 1 32 \frac{1}{32} 321 。
使用 VGG19 backbone 作为全卷积语义分割网络的下采样层，而在 VGG19 特征层之后，网络将增添上采样层。即增加新的转置卷积层，并完成部分层的特征逐点相加操作。最终将特征映射的尺寸逐渐恢复到原图的大小，特征映射的数量逐渐恢复到原类别的数量。
下图展示了不同种类的 FCN 网络语义分割操作方法，其中 FCN-32s 就是将最后的卷积或池化结果通过转置卷积，直接将特征映射的尺寸扩大 32 32 32 倍进行输出，而 FCN-16s 则是联合前面一次的结果将特征映射进行 16 16 16 倍的放大输出，而 FCN-8s 是联合前面两次的结果，通过转置卷积将特征映射的尺寸进行 8 8 8 倍的放大输出．在 FCN-8s 中将进行以下的操作步骤：
-
将最后一层的特征映射 P5(在 VGG19 中是第 5 5 5 个最大值池化层)通过转置卷积扩大 2 2 2 倍，得到新的特征映射 T5，并和 pool4 的特征映射 P4 相加可得到 T5+P4 。
-
将 T5+P4 通过转置卷积扩大 2 2 2 倍得到 T4，然后与 pool3 的特征映射 P3 相加得到 T4+P3 。
-
通过转置卷积，将特征映射 T4+P3 的尺寸扩大 8 8 8 倍，得到和输入形状一样大的结果．
接下来根据论文，基于 VGG19 的下采样 backbone，搭建 FCN-8s，用于图像的语义分割：
class FCN8s(nn.Module):
def __init__(self, num_classes, backbone):
super(FCN8s, self).__init__()
self.num_classes = num_classes # 分隔的总类别
self.backbone = backbone.to(device) # 下采样使用vgg19.features
# 定义上采样的层操作
self.relu = nn.ReLU(inplace=True)
self.deconv1 = nn.ConvTranspose2d(in_channels=512, out_channels=512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
self.bn1 = nn.BatchNorm2d(512)
self.deconv2 = nn.ConvTranspose2d(512, 256, 3, stride=2, padding=1, dilation=1, output_padding=1)
self.bn2 = nn.BatchNorm2d(256)
self.deconv3 = nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, dilation=1, output_padding=1)
self.bn3 = nn.BatchNorm2d(128)
self.deconv4 = nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, dilation=1, output_padding=1)
self.bn4 = nn.BatchNorm2d(64)
self.deconv5 = nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, dilation=1, output_padding=1)
self.bn5 = nn.BatchNorm2d(32)
self.classifier = nn.Conv2d(32, num_classes, 1)
# VGG19的MaxPooling所在层，用于逐点相加
self.pooling_layers = {'4': 'maxpool_1', '9': 'maxpool_2', '18': 'maxpool_3', '27': 'maxpool_4', '36': 'maxpool_5'}
def forward(self, x):
output = {}
# 对图像做下采样，并hook出pooling层特征
for name, layer in self.backbone._modules.items():
# 从第一层开始获取图像下采样特征
x = layer(x)
# 如果是pooling层，则保存到output中
if name in self.pooling_layers:
output[self.pooling_layers[name]] = x
P5 = output['maxpool_5'] # size=(N, 512, x.H/32, x.W/32)
P4 = output['maxpool_4'] # size=(N, 512, x.H/16, x.W/16)
P3 = output['maxpool_3'] # size=(N, 512, x.H/8, x.W/8)
# 对特征做转置卷积，即上采样，放大到原来大小
T5 = self.relu(self.deconv1(P5)) # size=(N, 512, x.H/16, x.W/16)
T5 = self.bn1(T5 + P4) # 特征逐点相加
T4 = self.relu(self.deconv2(T5)) # size=(N, 256, x.H/8, x.W/8)
T4 = self.bn2(T4 + P3)
T3 = self.bn3(self.relu(self.deconv3(T4))) # size=(N, 128, x.H/4, x.W/4)
T2 = self.bn4(self.relu(self.deconv4(T3))) # size=(N, 64, x.H/2, x.W/2)
T1 = self.bn5(self.relu(self.deconv5(T2))) # size=(N, 32, x.H, x.W)
score = self.classifier(T1) # 最后一层卷积输出, size=(N, num_classes, x.H, x.W)
return score
# 使用backbone为VGG19
fcn8s = FCN8s(21, backbone).to(device)
# 输入图像的尺寸应该是32的整数倍
summary(fcn8s.cpu(), input_size=(3, high, width), device='cpu')
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 320, 480] 1,792
ReLU-2 [-1, 64, 320, 480] 0
Conv2d-3 [-1, 64, 320, 480] 36,928
ReLU-4 [-1, 64, 320, 480] 0
MaxPool2d-5 [-1, 64, 160, 240] 0
Conv2d-6 [-1, 128, 160, 240] 73,856
ReLU-7 [-1, 128, 160, 240] 0
Conv2d-8 [-1, 128, 160, 240] 147,584
ReLU-9 [-1, 128, 160, 240] 0
MaxPool2d-10 [-1, 128, 80, 120] 0
Conv2d-11 [-1, 256, 80, 120] 295,168
ReLU-12 [-1, 256, 80, 120] 0
Conv2d-13 [-1, 256, 80, 120] 590,080
ReLU-14 [-1, 256, 80, 120] 0
Conv2d-15 [-1, 256, 80, 120] 590,080
ReLU-16 [-1, 256, 80, 120] 0
Conv2d-17 [-1, 256, 80, 120] 590,080
ReLU-18 [-1, 256, 80, 120] 0
MaxPool2d-19 [-1, 256, 40, 60] 0
Conv2d-20 [-1, 512, 40, 60] 1,180,160
ReLU-21 [-1, 512, 40, 60] 0
Conv2d-22 [-1, 512, 40, 60] 2,359,808
ReLU-23 [-1, 512, 40, 60] 0
Conv2d-24 [-1, 512, 40, 60] 2,359,808
ReLU-25 [-1, 512, 40, 60] 0
Conv2d-26 [-1, 512, 40, 60] 2,359,808
ReLU-27 [-1, 512, 40, 60] 0
MaxPool2d-28 [-1, 512, 20, 30] 0
Conv2d-29 [-1, 512, 20, 30] 2,359,808
ReLU-30 [-1, 512, 20, 30] 0
Conv2d-31 [-1, 512, 20, 30] 2,359,808
ReLU-32 [-1, 512, 20, 30] 0
Conv2d-33 [-1, 512, 20, 30] 2,359,808
ReLU-34 [-1, 512, 20, 30] 0
Conv2d-35 [-1, 512, 20, 30] 2,359,808
ReLU-36 [-1, 512, 20, 30] 0
MaxPool2d-37 [-1, 512, 10, 15] 0
ConvTranspose2d-38 [-1, 512, 20, 30] 2,359,808
ReLU-39 [-1, 512, 20, 30] 0
BatchNorm2d-40 [-1, 512, 20, 30] 1,024
ConvTranspose2d-41 [-1, 256, 40, 60] 1,179,904
ReLU-42 [-1, 256, 40, 60] 0
BatchNorm2d-43 [-1, 256, 40, 60] 512
ConvTranspose2d-44 [-1, 128, 80, 120] 295,040
ReLU-45 [-1, 128, 80, 120] 0
BatchNorm2d-46 [-1, 128, 80, 120] 256
ConvTranspose2d-47 [-1, 64, 160, 240] 73,792
ReLU-48 [-1, 64, 160, 240] 0
BatchNorm2d-49 [-1, 64, 160, 240] 128
ConvTranspose2d-50 [-1, 32, 320, 480] 18,464
ReLU-51 [-1, 32, 320, 480] 0
BatchNorm2d-52 [-1, 32, 320, 480] 64
Conv2d-53 [-1, 21, 320, 480] 693
================================================================
Total params: 23,954,069
Trainable params: 23,954,069
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.76
Forward/backward pass size (MB): 972.07
Params size (MB): 91.38
Estimated Total Size (MB): 1065.21
----------------------------------------------------------------
根据 torchsummary 结果可得，网络输入数据为 3 3 3 通道的 320 × 480 320\times480 320×480 的 RGB 图像，输出为 21 21 21 通道的 320 × 480 320\times480 320×480 的特征映射，该特征映射可以通过 F.log_softmax() 转为预测的类别。在 320 × 480 320\times480 320×480 的输出特征图中，每个取值就对应着相应的像素位置的预测类别。
这里提一句，上节教程我们也讲解了 U-Net 和 SegNet 的网络搭建。如果你想使用别的语义分割模型，把上面的类和对象定义换成 U-Net / SegNet 即可。
网络模型训练与验证
使用数据对网络模型训练一定的次数，并输出训练过程中最优的网络模型。
def train_model(model, loss_func, optimizer, train_dataloder, val_dataloader, num_epochs=25):
'''模型，损失函数，优化器，训练集和验证集，训练轮数'''
# 初始化变量
best_model_weights = copy.deepcopy(model.state_dict()) # 深拷贝最好模型
best_loss = 1e10
train_loss_all = []
val_loss_all = []
since = time.time()
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs-1))
print('-'*10)
train_loss, val_loss = 0.0, 0.0
train_num, val_num = 0, 0
# 训练，mini-batch(stochastic) gradient decesent
model.train()
for step, (X_train, y_train) in enumerate(train_dataloder):
optimizer.zero_grad() # 清空过往梯度
X_train = X_train.float().to(device)
y_train = y_train.long().to(device)
output = model(X_train)
# output = F.log_softmax(output, dim=1) # 映射->像素级类别
pre_lab = torch.argmax(output, 1) # 预测
loss = loss_func(output, y_train) # 损失
loss.backward() # 梯度反向传播
optimizer.step() # 根据梯度更新参数
train_loss += loss.item() * len(y_train)
train_num += len(y_train)
# 计算一个epoch中整个训练集的loss
train_loss_all.append(train_loss / train_num)
print('{} Train Loss: {:.4f}'.format(epoch, train_loss_all[-1]))
# 验证，计算整体在验证集上的损失
model.eval()
for step, (X_val, y_val) in enumerate(val_dataloader):
X_val = X_val.float().to(device)
y_val = y_val.long().to(device)
output = model(X_val)
# output = F.log_softmax(output, dim=1) # 映射->像素级类别
pre_lab = torch.argmax(output, 1) # 预测
loss = loss_func(output, y_val) # 损失
val_loss += loss.item() * len(y_val)
val_num += len(y_val)
# 计算一个epoch中整个验证集的loss
val_loss_all.append(val_loss / val_num)
print('{} Val Loss: {:.4f}'.format(epoch, val_loss_all[-1]))
# 保存最好的网络参数
if val_loss_all[-1] < best_loss:
best_loss = val_loss_all[-1]
best_model_weights = copy.deepcopy(model.state_dict())
# 计算每个epoch的训练时间
time_use = time.time() - since
print('Epoch {} complete in {:.0f}m {:.0f}s'.format(num_epochs-1, time_use // 60, time_use % 60))
train_process = pd.DataFrame(data={'epoch': range(num_epochs),
'train_loss_all': train_loss_all,
'val_loss_all': val_loss_all})
# 返回最佳模型
model.load_state_dict(best_model_weights)
return model, train_process
接下来定义学习率，训练轮次，损失函数和优化算法，对模型开始训练。
由于损失函数使用 NLLLoss + log_softmax 的组合，代码出现矩阵维度不一致问题，添加断点发现维度是一致的，所以就去掉了 log_softmax，且损失函数使用交叉熵。
LR = 0.0003
# negative log likelihood loss, 当程序使用log_softmax时使用该损失函数
# loss_func = nn.NLLLoss().to(device)
loss_func = nn.CrossEntropyLoss().to(device)
# 使用综合动量法和自适应学习率的Adam算法，并定义权重衰减
optimizer = torch.optim.Adam(fcn8s.parameters(), lr=LR, weight_decay=1e-4)
# 对数据训练epoch轮
fcn8s, train_process = train_model(fcn8s.cuda(), loss_func, optimizer,
train_loader, val_loader, num_epochs=30)
Epoch 0/29
----------
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /home/mist/pytorch/c10/core/TensorImpl.h:1156.)
return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
0 Train Loss: 2.5066
0 Val Loss: 1.6618
Epoch 29 complete in 0m 33s
Epoch 1/29
----------
1 Train Loss: 1.9521
1 Val Loss: 1.6669
Epoch 29 complete in 1m 5s
Epoch 2/29
----------
2 Train Loss: 1.6134
2 Val Loss: 1.4367
Epoch 29 complete in 1m 36s
Epoch 3/29
----------
3 Train Loss: 1.3749
3 Val Loss: 1.8534
Epoch 29 complete in 2m 7s
Epoch 4/29
----------
4 Train Loss: 1.2205
4 Val Loss: 1.4224
Epoch 29 complete in 2m 38s
Epoch 5/29
----------
5 Train Loss: 1.1137
5 Val Loss: 1.0534
Epoch 29 complete in 3m 10s
Epoch 6/29
----------
6 Train Loss: 1.0206
6 Val Loss: 1.2075
Epoch 29 complete in 3m 41s
Epoch 7/29
----------
7 Train Loss: 0.9866
7 Val Loss: 0.9980
Epoch 29 complete in 4m 12s
Epoch 8/29
----------
8 Train Loss: 0.9737
8 Val Loss: 1.2905
Epoch 29 complete in 4m 43s
Epoch 9/29
----------
9 Train Loss: 0.9126
9 Val Loss: 0.9246
Epoch 29 complete in 5m 14s
Epoch 10/29
----------
10 Train Loss: 0.8783
10 Val Loss: 0.8917
Epoch 29 complete in 5m 45s
Epoch 11/29
----------
11 Train Loss: 0.8586
11 Val Loss: 0.8941
Epoch 29 complete in 6m 16s
Epoch 12/29
----------
12 Train Loss: 0.8292
12 Val Loss: 0.9219
Epoch 29 complete in 6m 48s
Epoch 13/29
----------
13 Train Loss: 0.7993
13 Val Loss: 0.9292
Epoch 29 complete in 7m 19s
Epoch 14/29
----------
14 Train Loss: 0.7781
14 Val Loss: 0.8786
Epoch 29 complete in 7m 50s
Epoch 15/29
----------
15 Train Loss: 0.7882
15 Val Loss: 0.8466
Epoch 29 complete in 8m 21s
Epoch 16/29
----------
16 Train Loss: 0.7443
16 Val Loss: 0.8487
Epoch 29 complete in 8m 52s
Epoch 17/29
----------
17 Train Loss: 0.7285
17 Val Loss: 1.3754
Epoch 29 complete in 9m 23s
Epoch 18/29
----------
18 Train Loss: 0.7263
18 Val Loss: 0.8057
Epoch 29 complete in 9m 54s
Epoch 19/29
----------
19 Train Loss: 0.7105
19 Val Loss: 0.8222
Epoch 29 complete in 10m 26s
Epoch 20/29
----------
20 Train Loss: 0.6825
20 Val Loss: 0.8035
Epoch 29 complete in 10m 57s
Epoch 21/29
----------
21 Train Loss: 0.6409
21 Val Loss: 0.8084
Epoch 29 complete in 11m 28s
Epoch 22/29
----------
22 Train Loss: 0.6428
22 Val Loss: 0.8263
Epoch 29 complete in 11m 59s
Epoch 23/29
----------
23 Train Loss: 0.6256
23 Val Loss: 0.8310
Epoch 29 complete in 12m 30s
Epoch 24/29
----------
24 Train Loss: 0.6020
24 Val Loss: 0.8015
Epoch 29 complete in 13m 1s
Epoch 25/29
----------
25 Train Loss: 0.5812
25 Val Loss: 0.8111
Epoch 29 complete in 13m 32s
Epoch 26/29
----------
26 Train Loss: 0.5710
26 Val Loss: 0.8397
Epoch 29 complete in 14m 4s
Epoch 27/29
----------
27 Train Loss: 0.5678
27 Val Loss: 0.8663
Epoch 29 complete in 14m 35s
Epoch 28/29
----------
28 Train Loss: 0.5677
28 Val Loss: 0.8077
Epoch 29 complete in 15m 6s
Epoch 29/29
----------
29 Train Loss: 0.5297
29 Val Loss: 0.8142
Epoch 29 complete in 15m 37s
# 保存网络
torch.save(fcn8s, 'fcn8s.pt')
通过折线图可视化训练过程中损失函数的变化情况。
# 可视化模型训练过程
plt.figure(figsize=(10, 6))
plt.plot(train_process.train_loss_all, 'ro-', label='Train Loss')
plt.plot(train_process.epoch, train_process.val_loss_all, 'bs-', label='Val Loss')
plt.legend()
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.show()
接下来使用训练好的网络，从验证集中获取一个 batch 的图像，对其进行语义分割，将得到的结果和人工标注的 Label 进行对比。
# 验证集中获得一个batch的数据
for step, (b_x, b_y) in enumerate(val_loader):
if step > 0:
break
# 可视化预测效果
fcn8s.eval()
b_x = b_x.float().to(device)
b_y = b_y.long().to(device)
output = fcn8s(b_x)
output = F.log_softmax(output, dim=1)
pre_lab = torch.argmax(output, 1)
b_x_numpy = b_x.cpu().data.numpy()
b_x_numpy = b_x_numpy.transpose(0, 2, 3, 1)
b_y_numpy = b_y.cpu().data.numpy()
pre_lab_numpy = pre_lab.cpu().data.numpy()
# 可视化验证集的图像
plt.figure(figsize=(16, 10))
for i in range(4):
plt.subplot(3, 4, i+1)
plt.imshow(inv_normalize_image(b_x_numpy[i]))
plt.axis('off')
plt.title(str(i+1))
plt.subplot(3, 4, i+5)
plt.imshow(label2image(b_y_numpy[i], colormap))
plt.axis('off')
plt.title(str(i+5))
plt.subplot(3, 4, i+9)
plt.imshow(label2image(pre_lab_numpy[i], colormap))
plt.axis('off')
plt.title(str(i+9))
plt.subplots_adjust(wspace=0.05, hspace=0.05)
plt.show()
1 − 4 1-4 1−4 为原始 RGB 图， 5 − 8 5-8 5−8 为人工标注的语义分割图， 9 − 12 9-12 9−12 为全卷积语义分割网络对图像的分割结果。从对比可以看出，网络能分割出一些目标，但是分割的精度并不高，并且分类的类别也出现一定的误差。
模型整体仍有提升空间，这与我们使用的基础网络 backbone VGG19 的深度不够，以及训练集的数据较少都有关系。
如果可以的话，我们也能把模型的 backbone 换成 ResNet 101，或者模型整体换成上一节代码中的 U-Net 和 SegNet，重新训练，对比不同模型的分割结果。
其他的优化方向，包括但不限于 K 折交叉验证，提前停止训练，修改图像增强方式，使用更多的图像作为训练集。使用 Stacking 集成学习结合不同模型的分割结果，图像开运算优化边界。篇幅原因，这里不再赘述。