第14篇 Fast AI深度学习课程——超分辨与图像分割

一、超分辨

超分辨指的是由低分辨率的图片获得高分辨率的图片。

1. 准备数据集

数据集无需标注，将图像进行降采样即可获得配对的高低分辨率的图像。同样不需数据标注的应用场景还有：图像旋转、图像去噪、黑白图像着色等。

为定义合适的Dataset，我们去fastai.dataset.py中找到一个和需求相接近的。其中有一个FilesDataset，其接受文件名，数据集的输入x为图像。我们继承该类，并覆写输出函数get_y()和类别函数get_c():

class MatchedFilesDataset(FilesDataset):
    def __init__(self, fnames, y, transform, path):
        self.y=y
        assert(len(fnames)==len(y))
        super().__init__(fnames, transform, path)
    def get_y(self, i): return open_image(os.path.join(self.path, self.y[i]))
    def get_c(self): return 0

然后在设置对数据集所进行的变换时，指定输入图像x的分辨率为y的一半。

tfms = tfms_from_model(arch, sz_lr, tfm_y=TfmType.PIXEL, 
          aug_tfms=aug_tfms, sz_y=sz_hr)
datasets = ImageData.get_ds(MatchedFilesDataset, (trn_x,trn_y), 
               (val_x,val_y), tfms, path=PATH_TRN)
md = ImageData(PATH, datasets, bs, num_workers=16, classes=None)

在对图像进行变换时，可能会做维度顺序的调整、图像归一化等，这些操作的参数会被存储，后续可通过dataset.denorm()函数，对dataset中的图像进行还原，以用于展示。

2. 网络模型

有两种方式可以实现超分辨：一种是先进行升采样，然后使用一系列的卷积，逐步生成图像；另一种是先提取图像特征，然后升采样。课程中使用的是第二种，因为其计算量会小很多。

现考虑特征的提取。由于输入图像和输出图像很接近，在此我们使用ResBlock做特征提取，这样可以保持图像主体不变，而仅针对用于图像分辨率增强的细节进行学习。定义如下：

def conv(ni, nf, kernel_size=3, actn=True):
    layers = [nn.Conv2d(ni, nf, kernel_size, padding=kernel_size//2)]
    if actn: layers.append(nn.ReLU(True))
    return nn.Sequential(*layers)


class ResSequential(nn.Module):
    def __init__(self, layers, res_scale=1.0):
        super().__init__()
        self.res_scale = res_scale
        self.m = nn.Sequential(*layers)

    def forward(self, x):
        x = x + self.m(x) * self.res_scale
        return x

        
def res_block(nf):
    return ResSequential(
        [conv(nf, nf), conv(nf, nf, actn=False)],
        0.1)

其中ResBlock的结构如下图所示。其与之前的ResBlock的主要不同之处在于去除了Batch Norm层。理由是为了尽量保证图像与原图相符，不能使像素(无论是直通的数据、还是残差数据)过分偏离原图的分布。而Batch Norm层强行改变了整块数据的分布。

图 1.ResBlock结构

查看ResSequential的定义，其与Pytorch的Sequential的差别在于多了一个残差的缩放因子。之所以这样做，是因为在训练过程中存在如下现象：当所使用的训练数据的batch过大时，早期训练过程极易产生很大的损失函数。而对ResBlcok的残差部分使用小于1的因子进行缩放，可减缓这一现象，使得训练过程更为稳定。直觉上来看，这毫无道理，因为可通过对卷积核乘上一个系数达到同样的效果。这一trick能够有效，可能与计算的离散化有关。

最终所构建的网络结构如下：

图 2.超分辨网络结构

值得说明的是其中的升采样模块。升采样的方法也有很多，之前接触到的有两种：一种是使用转置卷积的方法；另一种是使用最近邻升采样，然后用1x1的卷积进行优化。在这里并未使用转置卷积和近邻插值。使用转置卷积进行升采样的缺点如下：其存在许多不必要的计算，因为升采样时会插入许多0；其会引入不存在的结构特性，因为在卷积边缘和内部，实际起作用的像素数是不同的，这种不同经过一步一步的扩张，会引入原图中所不存在的结构特性。

图 3.转置卷积图示

课程中所用的升采样法是像素交叉法。如一个特征切片大小为nxn，要升采样2倍，则沿特征方向按4倍来扩展输出维度；每4个特征切片的像素交叉排列，再恢复原特征维的长度，但特征切片变成了2nx2n。

图 4.升采样图示

综上，升采样模块的代码如下：

def upsample(ni, nf, scale):
    layers = []
    for i in range(int(math.log(scale,2))):
        layers += [conv(ni, nf*4), nn.PixelShuffle(2)]
    return nn.Sequential(*layers)

然而，进行上述操作后，所得结果会出现棋盘现象。出现的原因如下：在升采样之初，所有的特征切片的随机初始化都是相互独立的。这就会导致最终所得的用于像素交叉的特征，各切片间存在系统性差异。这样交织排列后自然就出现了棋盘现象。解决方法就是在初始化时，随机初始化nf个特征切片，然后复制到其他切片上。

3. 损失函数

使用生成图片与真实高分辨率的图片的像素差值的平方和作为损失函数，会导致图像的模糊(还没想明白)。因此考虑使用特征差异，来描述生成图片和图片真值之间的差异。

同风格迁移中的做法一样，使用VGG提取图像特征，然后计算图像的内容损失。首先找到做MaxPool之前的网络层的激活值，舍弃最后若干层接触域过大、分辨率较低的特征。

4. 并行执行

使用nn.DataParallel封装模型，并指定所要使用的GPU序号。

m = to_gpu(SrResnet(64, scale))
m = nn.DataParallel(m, [0,2])
learn = Learner(md, SingleModel(m), opt_fn=optim.Adam)

模型会被封装进一个Module属性中，这和单GPU时不同，也导致两种情形下数据存储的不同。若要在单个GPU上加载从多个GPU上存储下来的模型，有两种方法：一种是在多GPU存储时，存储m.module；第二种是指明同一个GPU的id，使用DataParallel进行save和load。

5. 步进式超分辨

如果我们已经通过训练获得可将图片扩大2倍的网络Net-2，如何得到一个可将图片扩大4倍的网络Net-4呢？两个网络相比较，Net-4要比Net-4前者多一个upsample层。可以使用如下语句将Net-2的系数导入Net-4的对应层，而对Net-4中多出来的层，进行随机初始化，并训练。

learn.model.load_state_dict(t, strict=False)

其中strict=False表示有多少层就加载多少层。

二、图像分割

1. 数据说明

本部分的数据来源于Kaggle上的Carnava竞赛。本例中仅使用该数据集中的训练数据，其中包含近320辆汽车的16种角度下的图片以及掩膜。涉及的文件包括train_mask.csv、train/、train_mask/。

下载数据后，首先将图片及其分类掩码整理为统一样式(文件格式、尺寸等)。然后在划分训练集和验证集时，注意不要将同一辆汽车的不同角度的图片划到不同集合上。注意在做数据修饰时，设置tfm_y=TfmType.CLASS，这样旋转图片时，不会对掩膜图片进行插值。

1. 使用基于ImageNet预训练的ResNet34模型

类似于在目标识别一课中的做法，将分割问题视为分类问题，在ResNet34网络的基础上，添加定制的头部。
由于最终要获得和原图同样大小的掩膜图片，而ResNet34输出的特征切片的尺寸要小于原图，因此可考虑使用一系列的升采样层作为附加的头部。

class StdUpsample(nn.Module):
    def __init__(self, nin, nout):
        super().__init__()
        self.conv = nn.ConvTranspose2d(nin, nout, 2, stride=2)
        self.bn = nn.BatchNorm2d(nout)
        
    def forward(self, x): return self.bn(F.relu(self.conv(x)))

flatten_channel = Lambda(lambda x: x[:,0])
simple_up = nn.Sequential(
    nn.ReLU(),
    StdUpsample(512,256),
    StdUpsample(256,256),
    StdUpsample(256,256),
    StdUpsample(256,256),
    nn.ConvTranspose2d(256, 1, 2, stride=2),
    flatten_channel
)

为什么是5层升采样呢？对于ResNet34，其卷积部分会将224x224的图像转化为7x7的特征，降采样5次共32倍。

然后指定优化器、设置损失函数，即可构建模型：

models = ConvnetBuilder(resnet34, 0, 0, 0, custom_head=simple_up)
learn = ConvLearner(md, models)
learn.opt_fn=optim.Adam
learn.crit=nn.BCEWithLogitsLoss()
learn.metrics=[accuracy_thresh(0.5)]

使用尺寸为128x128图片集。训练附加层的参数，可得到96.3%的准确率。然而，这一准确率并不意味着分割效果可被接受：

图 5.准确率为96.3%时分割结果

然后解除网络的冻结状态，微调`ResNet34`层系数，可得`99.1%`的分割准确度。由下图所示的分割结果中可见，车身的后视镜部分略有缺失。

图 6.准确率为99.1%时分割结果

采用步进式的训练策略，依次在512x512、1024x1024尺寸的数据集上，使用前次训练的参数为初始值，训练模型，最终可得分割准确率达99.8%的网络。

为什么网络结构未变，而输入图片的尺寸可以变化呢？原因在于在ConvBuilder()中，将ResNet34的主干部分抽离了出来，不包含全连接层，仅包含卷积部分。对于卷积部分，不需要指定图片的大小。

2. U-Net

考虑上述模型。ResNet34会将一幅224x224的图像，缩减为尺寸为7x7的特征集。而这7x7的特征集则是后续升采样模块的起始。从接触域的角度来说，采用上述方法，我们实际上是从一个很粗糙的特征集出发，想得到一个精细的分割。这是有难度的。如果我们能够利用在卷积过程中产生的精细粒度(接触域)不同的特征，那么有理由相信，能够省力地得到一个效果还可以的分割结果。

U-Net就是利用了这一思想，其在升采样过程中，将升采样的输入特征，和降采样过程中对应步骤的激活函数，沿特征维进行拼接(有可能需要进行尺寸的裁剪)，以使得升采样过程可以得到足够精细的图片信息。其网络结构如下。其正得名于下图所示的U字构型。

图 7.U-Net网络结构

本课将实现类U-Net的网络结构。

为使用已有的预训练的ResNet34网络，首先使用get_base()函数，获取ResNet34网络的主干。

f = resnet34
cut,lr_cut = model_meta[f]
def get_base():
    layers = cut_model(f(True), cut)
    return nn.Sequential(*layers)

查看Fast.AI的源码可知，model_meta中存储了各个模型的主干网络的结束索引。对于ResNet34而言，其主干网络截止于全连接层之前，包含了由各卷积层和ResBlock构成的降采样部分。

然后考虑使用前几课中用到的钩子技术，将ResNet34中，每个跨立度为2的卷积层的输入保存下来，以拼接到升采样过程中对应的特征上。

class SaveFeatures():
    features=None
    def __init__(self, m): self.hook = m.register_forward_hook(self.hook_fn)
    def hook_fn(self, module, input, output): self.features = output
    def remove(self): self.hook.remove()

定义U-Net Block，以实现U-Net的升采样策略：

class UnetBlock(nn.Module):
    def __init__(self, up_in, x_in, n_out):
        super().__init__()
        up_out = x_out = n_out//2
        self.x_conv  = nn.Conv2d(x_in,  x_out,  1)
        self.tr_conv = nn.ConvTranspose2d(up_in, up_out, 2, stride=2)
        self.bn = nn.BatchNorm2d(n_out)
        
    def forward(self, up_p, x_p):
        up_p = self.tr_conv(up_p)
        x_p = self.x_conv(x_p)
        cat_p = torch.cat([up_p,x_p], dim=1)
        return self.bn(F.relu(cat_p))

其中up_in为升采样模块的输入特征up_p的维度，x_in为从降采样模块接引来的特征x_p的维度，n_out为输出特征的维度。x_conv是对x_p进行操作的卷积模块，tr_conv是对up_in进行转置卷积的模块。由于最后要将本模块特征与降采样模块特征拼接，因此，设置x_conv和tr_conv的输出特征的维度各为n_out/2。

由此定义网络模型：

class Unet34(nn.Module):
    def __init__(self, rn):
        super().__init__()
        self.rn = rn
        self.sfs = [SaveFeatures(rn[i]) for i in [2,4,5,6]]
        self.up1 = UnetBlock(512,256,256)
        self.up2 = UnetBlock(256,128,256)
        self.up3 = UnetBlock(256,64,256)
        self.up4 = UnetBlock(256,64,256)
        self.up5 = nn.ConvTranspose2d(256, 1, 2, stride=2)
        
    def forward(self,x):
        x = F.relu(self.rn(x))
        x = self.up1(x, self.sfs[3].features)
        x = self.up2(x, self.sfs[2].features)
        x = self.up3(x, self.sfs[1].features)
        x = self.up4(x, self.sfs[0].features)
        x = self.up5(x)
        return x[:,0]
    
    def close(self):
        for sf in self.sfs: sf.remove()

最后对目标识别做一点说明。在第8、9课中，我们得到的用于目标识别的网络实际上对小目标的识别效果不好。可以考虑使用U-Net的策略进行改进。事实上，已经有文献提出了相应的方法，并将之命名为特征金字塔网络(FPN，Feature Pyramid network)，其核心思想与U-Net并无不同。

一些有用的链接

• 课程wiki: 本节课程的一些相关资源，包括课程笔记、课上提到的博客地址等。
• Perceptual Losses for Real-Time Style Transfer and Super-Resolution: Justin的讲风格迁移和超分辨的论文，Justin就是那个CS231n的小哥。
• 一篇关于如何处理结构化数据的博客。
• fastai.model.DynamicUnet: 类U-Net网络的实现代码，可对U-Net Block中的特征维度灵活设置。