13.10 语义分割 & 全卷积网络

语义分割是对图像的每个像素分类

全卷积网络采用卷积神经网络实现从图像像素到像素类别的转换，全卷积网络将中间层特征的高和宽转换回输入图像的尺寸（引入转置卷积实现的）。
最终的类别预测与输入图像在像素上一一对应。

全卷积网络模型模型框架

1. 使用卷积神经网络提取图像特征。
2. 1x1卷积层将通道数转换为类别数。
3. 转置卷积层将特征图的高和宽转换为输入图像的大小。使最终模型的输入和输出高和宽相同。
4. 最终输出通道包含了该空间位置像素的类别预测。

主代码

我的torch和torchvision版本

import torch
torch.__version__

‘1.7.1+cu110’

import torchvision
torchvision.__version__

‘0.8.2+cu110’

%matplotlib inline
import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from d2l import torch as d2l
from tools import train_tool,data_tool
import time

'''
    用ImageNet数据集上预训练的ResNet-18模型来提取图像特征,并将该网络记为pretrained_net。
    ResNet-18模型的最后几层包括全局平均汇聚层和全连接层,全卷积网络中不需要它们。
'''
pretrained_net = torchvision.models.resnet18(pretrained=True)
# 输出最后3层
list(pretrained_net.children())[-3:]

输出：

[Sequential(
   (0): BasicBlock(
     (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
     (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     (relu): ReLU(inplace=True)
     (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
     (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     (downsample): Sequential(
       (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
       (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     )
   )
   (1): BasicBlock(
     (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
     (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     (relu): ReLU(inplace=True)
     (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
     (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
   )
 ),
 AdaptiveAvgPool2d(output_size=(1, 1)),
 Linear(in_features=512, out_features=1000, bias=True)]

# 创建全卷积网络Net,它复制ResNet18中的预训练层，去掉最后的两层：分别是全局平均汇聚层和全连接层
net = nn.Sequential(*list(pretrained_net.children())[:-2])

# 测试预训练模型的输出shape
X = torch.rand(size=(1,3,320,480))
net(X).shape # 变成了通道为512，shape为(10,18)，h,w缩小到原来的1/32

'''为模型添加：1x1卷积层 和 全卷积层'''
num_classes=21
# 使用1x1卷积层，将输出通道转换为该数据集的类别数21
net.add_module('1x1_conv',nn.Conv2d(512,num_classes,kernel_size=1))
# 输入和输出通道不改变，但要把输出结果的h,w变回原始尺寸。
'''
    输入:(batch_size,num_classes，10,18)
    计算过程：
            输出矩阵的高 = (输入的高-1) x 高方向的步长 - 2 x 高方向的填充 + 卷积核的高
                       =（10-1)x32 - 2x16 + 64
                       = 320
            输出矩阵的高 = 480
            成功将卷积提取的宽高减少到1/32的特征还原回了输入的尺寸
'''
net.add_module('transpose_conv',nn.ConvTranspose2d(num_classes,num_classes,kernel_size=64,padding=16,stride=32))

'''初始化转置卷积层，用双线性插值进行上采样放大图像'''
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (torch.arange(kernel_size).reshape(-1, 1),
    torch.arange(kernel_size).reshape(1, -1))
    filt = (1 - torch.abs(og[0] - center) / factor) * \
           (1 - torch.abs(og[1] - center) / factor)
    weight = torch.zeros((in_channels, out_channels,kernel_size, kernel_size))
    weight[range(in_channels), range(out_channels), :, :] = filt
    return weight

# 构建一个将输入的高和宽分别放大两倍的转置卷积层，并将该卷积核用bilinear_kernal函数初始化
conv_trans = nn.ConvTranspose2d(3, 3, kernel_size=4, padding=1, stride=2,bias=False)
'''
    卷积层的参数是双线性插值的权重。
    上面只是让转置卷积后的输出形状放大，不能保证图片内容不变。
    转置矩阵的值使用双线性插值的权重，这样就可以用于将图片放大，不打乱图片的排布

'''
conv_trans.weight.data.copy_(bilinear_kernel(3, 3, 4)); # 分号是防止它输出

# 读取图像，展示上采样结果
img = torchvision.transforms.ToTensor()(d2l.Image.open('../data/images/cat_dog.jpg'))
X = img.unsqueeze(0)
Y = conv_trans(X)
out_img = Y[0].permute(1, 2, 0).detach()

d2l.set_figsize()
print('input image shape:', img.permute(1, 2, 0).shape)
d2l.plt.imshow(img.permute(1, 2, 0));
print('output image shape:', out_img.shape)
d2l.plt.imshow(out_img);

# 用双线性插值上采样的参数 初始化 模型转置卷积层的参数
W = bilinear_kernel(num_classes, num_classes, 64)
net.transpose_conv.weight.data.copy_(W);

batch_size, crop_size = 32, (320, 480)
train_iter, test_iter = data_tool.load_data_voc(batch_size, crop_size)

# 用通道的索引表示类别，该数据集共21个标签，那么就用21个通道
def loss(inputs, targets):
    '''
    交叉熵计算之后大概变成了(样本，h=21维的通道，h，w)
    第一个mean(1)：表示按照类别维度求和取平均值，求出这21个类别的平均损失
    第二个mean(1)：表示按照样本求平均值，在类别平均损失基础上，对每个样本求和取平均值
    得到每个样本的每个类别的平均损失
    '''
    return F.cross_entropy(inputs, targets, reduction='none').mean(1).mean(1)

num_epochs, lr, wd, devices = 5, 0.001, 1e-3, d2l.try_all_gpus()
trainer = torch.optim.SGD(net.parameters(),lr=lr,weight_decay=wd)

'''开始计时'''
start_time = time.time()
# 开始训练
train_tool.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)
# d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)

'''计时结束'''
end_time = time.time()
run_time = end_time - start_time
# 将输出的秒数保留两位小数
if int(run_time)<60:
    print(f'{round(run_time,2)}s')
else:
    print(f'{round(run_time/60,2)}minutes')

工具代码

放在tools目录下：

data_tool.py

import torch
import torchvision
from d2l import torch as d2l
import os

VOC_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]]
#@save
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
'diningtable', 'dog', 'horse', 'motorbike', 'person',
'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']

'''
   筒装法得到一个列表
   索引：3维RGB值转换的一维索引值 colormap[0] * 256 + colormap[1]) * 256 + colormap[2]
   值：标签索引 0-20
'''
def voc_colormap2label():
    """
    标签颜色 和 标签类别长度都是21
    计算3维RGB值对应的一维索引。用的是筒装法(索引是3维RGB转成一维所计算的值，这21个像素对应的值是标签的索引 0-20)，
    """
    colormap2label = torch.zeros(256 ** 3, dtype=torch.long)
    for i, colormap in enumerate(VOC_COLORMAP):
        # RGB值与一维索引相对应
        colormap2label[(colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
    return colormap2label

def voc_label_indices(colormap, colormap2label):
    """将VOC标签中的RGB值映射到它们的类别索引"""
    colormap = colormap.permute(1, 2, 0).numpy().astype('int32')
    # 输入rgb值，计算3维RGB值对应的一维索引
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256 + colormap[:, :, 2])
    # 访问列表得到像素标签对应的标签索引
    return colormap2label[idx]

'''
    预处理数据：之前是缩放图像使其符合模型的输入形状。然而在语义分割中，这样做还需要再映射回原始输入尺寸，这样做对语义分割来说不好。
    具体的方法是使用图像增广中的随机裁减。要裁减输入图像和标签相同的区域。
'''
def voc_rand_crop(feature,label,height,width):
    # 随机裁减特征图像和标签图像
    rect = torchvision.transforms.RandomCrop.get_params(
        feature,(height,width)
    )
    feature = torchvision.transforms.functional.crop(feature,*rect)
    label = torchvision.transforms.functional.crop(label,*rect)
    return feature,label

'''自定义语义分割数据集类'''
class VOCSegDataset(torch.utils.data.Dataset):
    def __init__(self,is_train,crop_size,voc_dir):
        self.transform = torchvision.transforms.Normalize(
            mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
        self.crop_size = crop_size
        # 得到列表形式的输入图片和标签图片
        features,labels = read_voc_images(voc_dir,is_train=is_train)
        # 过滤掉图像尺寸小于随机裁减所指定的输出尺寸，然后对输入图片进行归一化
        self.features = [self.normalize_image(feature)
                         for feature in self.filter(features)]
        self.labels = self.filter(labels)
        #  用筒装法 计算3维RGB值对应的一维索引
        self.colormap2label = voc_colormap2label()
        print('read ' + str(len(self.features)) + ' examples')

    # 把图像标准化的函数
    def normalize_image(self,img):
        return self.transform(img.float() / 255)

    # 过滤函数：将图像尺寸小于随机裁减所指定的输出尺寸的图像过滤掉
    def filter(self,imgs): # img.shape[0]应该是批次
        return [img for img in imgs if (
            img.shape[1]>=self.crop_size[0] and
            img.shape[2]>=self.crop_size[1]
        )]
    '''
        可以让对象通过索引来访问元素如：
        vsd = VOCSegDataset()
        vsd[0]
    '''
    def __getitem__(self, idx):
        # 返回idx索引下的裁减后的feature,label
        feature,label = voc_rand_crop(self.features[idx],self.labels[idx], *self.crop_size)
        return (feature,voc_label_indices(label,self.colormap2label))

    #  可以通过Len得到对象的长度
    def __len__(self):
        return len(self.features)


'''
    ImageSets/Segmentation路径:用于训练和测试样本的文本文件
    JPEGImages: 图像样本
    SegmentationClass： 标签，标签也是图像格式，其尺寸和它所标注的输入图像的尺寸相同。颜色相同的像素属于同一个语义类别
'''
def read_voc_images(voc_dir,is_train=True):
    '''读取 voc_dir/ImageSets/Segmentation/train.txt或val.txt'''
    txt_fname = os.path.join(voc_dir,'ImageSets','Segmentation',
                             'train.txt' if is_train else 'val.txt')
    # mode = torchvision.io.image.ImageReadMode.RGB
    with open(txt_fname,'r') as f:
        # split()方法是将内容按空格分割为一个列表 images
        images = f.read().split()
    features,labels = [],[]
    for i,fname in enumerate(images):
        # 把图像样本放到features中
        features.append(torchvision.io.read_image(os.path.join(voc_dir,'JPEGImages',f'{fname}.jpg')))
        # 读取标签，加载到labels中，mode是torchvision.io.image.ImageReadMode.RGB
        labels.append(torchvision.transforms.ToTensor()(d2l.Image.open(
            os.path.join(voc_dir,'SegmentationClass',f'{fname}.png')).convert('RGB')))
    return features, labels

'''整合如上的所有函数，定义一个函数加载并读取语义分割数据集'''
def load_data_voc(batch_size,crop_size):
   '''返回训练数据集dataloader 和 测试数据集dataloader'''
   # 下载数据并返回文件目录
   voc_dir = d2l.download_extract('voc2012',os.path.join('VOCdevkit', 'VOC2012'))
   num_workers = d2l.get_dataloader_workers()
   train_iter = torch.utils.data.DataLoader(
       VOCSegDataset(True,crop_size,voc_dir),batch_size,
       shuffle=True,drop_last = True,num_workers=num_workers
   )
   test_iter = torch.utils.data.DataLoader(
       VOCSegDataset(False,crop_size,voc_dir),batch_size,
       shuffle=True,drop_last = True,num_workers=num_workers
   )
   return train_iter,test_iter

train_tool.py

from torch import nn
from d2l import torch as d2l
import torch

def train_batch_ch13(net, X, y, loss, trainer, devices):
    """Train for a minibatch with mutiple GPUs (defined in Chapter 13).

    Defined in :numref:`sec_image_augmentation`"""
    if isinstance(X, list):
        # Required for BERT fine-tuning (to be covered later)
        X = [x.to(devices[0]) for x in X]
    else:
        X = X.to(devices[0])
    y = y.to(devices[0])
    net.train()
    trainer.zero_grad()
    pred = net(X)
    l = loss(pred, y)
    l.sum().backward()
    trainer.step()
    train_loss_sum = l.sum()
    train_acc_sum = d2l.accuracy(pred, y)
    return train_loss_sum, train_acc_sum

def train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs,
               devices=d2l.try_all_gpus()):
    """Train a model with mutiple GPUs (defined in Chapter 13).

    Defined in :numref:`sec_image_augmentation`"""
    timer, num_batches = d2l.Timer(), len(train_iter)
    animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs], ylim=[0, 1],
                            legend=['train loss', 'train acc', 'test acc'])
    net = nn.DataParallel(net, device_ids=devices).to(devices[0])
    '''改动'''
    best_test_acc = 0
    for epoch in range(num_epochs):
        # Sum of training loss, sum of training accuracy, no. of examples,
        # no. of predictions
        metric = d2l.Accumulator(4)
        for i, (features, labels) in enumerate(train_iter):
            timer.start()
            l, acc = train_batch_ch13(
                net, features, labels, loss, trainer, devices)
            metric.add(l, acc, labels.shape[0], labels.numel())
            timer.stop()
            if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
                animator.add(epoch + (i + 1) / num_batches,
                             (metric[0] / metric[2], metric[1] / metric[3],
                              None))
        test_acc = d2l.evaluate_accuracy_gpu(net, test_iter)
        # 此处保存的是最高准确率的一轮
        if test_acc > best_test_acc:
            best_test_acc = test_acc
            # 保存准确率最高的一轮训练参数
            torch.save(net.state_dict(), 'best_segment_net_weights.pth')

        animator.add(epoch + 1, (None, None, test_acc))
    print(f'loss {metric[0] / metric[2]:.3f}, train acc '
          f'{metric[1] / metric[3]:.3f}, test acc {test_acc:.3f}, best test acc {best_test_acc:.3f}')
    print(f'{metric[2] * num_epochs / timer.sum():.1f} examples/sec on '
          f'{str(devices)}')