★★★ 本文源自AI Studio社区精品项目，【点击此处】查看更多精品内容 >>>

---

Paddle复现CPS: 基于交叉伪监督的半监督语义分割

1. CPS半监督语义分割简介

• 参考代码：TorchSemiSeg
• 参考论文：Semi-Supervised Semantic Segmentation with Cross Pseudo Supervision
• CPS解读参考：[CVPR 2021] CPS: 基于交叉伪监督的半监督语义分割

1.1 项目背景

• 不同于图像分类任务，数据的标注对于语义分割任务来说是比较困难而且成本高昂的，需要为图像的每一个像素标注一个标签，包括一些特别细节的物体，如电线杆等。但是对于获取RGB数据是比较简单，如何利用大量的无标注数据去提高模型的性能，便是半监督语义分割领域研究的问题

• 作者为半监督语义分割任务设计了一种非常简洁而又性能很好的算法(cross pseudo supervision, CPS)，训练时，使用两个相同结构、但是不同初始化的网络，添加约束使得两个网络对同一样本的输出是相似的。具体来说，当前网络产生的one-hot pseudo label，会作为另一路网络预测的目标，这个过程可以用cross entropy loss监督，就像传统的全监督语义分割任务的监督一样。该算法在在两个benchmark (PASCAL VOC, Cityscapes) 都取得了SOTA的结果。本项目复现的为在0.5倍Cityscapes数据集数据量上的结果

1.2 CPS算法

• 如上图所示，CPS的设计非常的简洁。训练时，使用两个网络$f({\theta }_1)$ 和 $f({\theta }_2)$。这样对于同一个输入图像X，可以有两个不同的输出P1和P2。通过argmax操作得到对应的one-hot标签Y1和Y2。类似于self-training中的操作，将这两个伪标签作为监督信号。例如用Y2作为P1的监督，Y1作为P2的监督，并用cross entropy loss约束

• 对于这两个网络，使用相同的结构，例如DeepLabV3+，但是不同的初始化。使用kaiming_normal进行两次随机初始化，而没有对初始化的分布做特定的约束。当然了，如果设计特定的初始化，没准CPS的效果会更好

• 在测试的时候，我们只使用其中一个网络进行inference，所以不增加任何测试/部署时候的开销

2. 核心代码

• 半监督语义分割任务的网络一般是已有网络，特别之处在于训练方式，同时由于是半监督学习，在数据加载部分也和监督学习不同，同时还会有cutmix这种batch transform的trick用上，出于篇幅考虑，本项目仅展示复现的网络、优化器设置以及体现CPS思想的训练过程的部分核心代码

2.1 网络结构

• 本项目复现CPS算法所用的子网络为DeepLabV3+，backbone为resnet50，不同于PaddleSeg所用的resnet_ssld系列。且将resnet50作为backbone时，在layer4修改了卷积的步长和空洞大小，关键代码如下:

def _make_layer__nostride_dilate(self, block, norm_layer, planes, blocks, stride=1, bn_eps=1e-5, bn_momentum=0.1):
    downsample = nn.Sequential(
        nn.Conv2D(self.inplanes, planes * block.expansion, 
                    kernel_size=1, stride=1, bias_attr=False),
            norm_layer(planes * block.expansion, epsilon=bn_eps,
                    momentum=bn_momentum), 
    )

    dilate = 2
    layers = []
    layers.append(block(self.inplanes, planes, stride, norm_layer, bn_eps,
                        bn_momentum, downsample, dilate))

    self.inplanes = planes * block.expansion
    for i in range(1, blocks):
        dilate *= 2
        layers.append(block(self.inplanes, planes,
                            norm_layer=norm_layer, bn_eps=bn_eps,
                            bn_momentum=bn_momentum, dilate=dilate))

    return nn.Sequential(*layers)

• CPS的大网络中包含两个Deeplabv3+子网络，对于同一个输入有不同的输出。论文源码中是通过step参数来控制其中一个网络的前向计算，但是对于paddle多卡分布式训练来说，有条件的前向计算会报错，所以复现代码改为每次都进行两次前向计算，增加一些计算的开销来解决这个问题，大网络代码如下:

class Network(nn.Layer):
    def __init__(self, num_classes, pretrained_model=None):
        super(Network, self).__init__()
        self.backbone1 = resnet50(pretrained_model, norm_layer=layers.SyncBatchNorm,
                                  bn_eps=config.bn_eps,
                                  bn_momentum=config.bn_momentum,
                                  deep_stem=True, stem_width=64, as_backbone=True)
                                  
        self.backbone2 = resnet50(pretrained_model, norm_layer=layers.SyncBatchNorm,
                                  bn_eps=config.bn_eps,
                                  bn_momentum=config.bn_momentum,
                                  deep_stem=True, stem_width=64, as_backbone=True)
        
        self.branch1 = DeepLabV3P(num_classes=num_classes, backbone=self.backbone1, align_corners=True)
        self.branch2 = DeepLabV3P(num_classes=num_classes, backbone=self.backbone2, align_corners=True)

    def forward(self, data):
        if not self.training:
            pred1 = self.branch1(data)
            return pred1

        return self.branch1(data), self.branch2(data)

2.2 优化器初始化

• CPS的优化器优化的参数是根据网络不同的层数，设置不一样的优化器参数，有的层的weight_decay为0，参数初始化函数如下:

def group_weight(weight_group, module, norm_layer, lr):
    group_decay = []
    group_no_decay = []
    for m in module.sublayers():
        if isinstance(m, nn.Linear):
            group_decay.append(m.weight)
            if m.bias is not None:
                group_no_decay.append(m.bias)
        elif isinstance(m, (nn.Conv1D, nn.Conv2D, nn.Conv3D, nn.Conv2DTranspose, nn.Conv3DTranspose)):
            group_decay.append(m.weight)
            if m.bias is not None:
                group_no_decay.append(m.bias)
        elif isinstance(m, norm_layer) or isinstance(m, nn.BatchNorm1D) or isinstance(m, nn.BatchNorm2D) \
                or isinstance(m, nn.BatchNorm3D) or isinstance(m, nn.GroupNorm):
            if m.weight is not None:
                group_no_decay.append(m.weight)
            if m.bias is not None:
                group_no_decay.append(m.bias)
        elif isinstance(m, nn.Embedding):
            group_decay.append(m)

    weight_group.append(dict(params=group_decay))
    weight_group.append(dict(params=group_no_decay, weight_decay=.0))

    return weight_group

• 优化器定义如下：

params_list_l = []
params_list_l = group_weight(params_list_l, model.branch1.backbone,
                                BatchNorm2D, base_lr)
params_list_l = group_weight(params_list_l, model.branch1.head, BatchNorm2D,
                                base_lr)        # head lr * 10

optimizer_l = paddle.optimizer.Momentum(parameters=params_list_l,
                                learning_rate=base_lr,
                                momentum=config.momentum,
                                weight_decay=config.weight_decay)

2.3 训练过程

• 数据加载

minibatch = dataloader.next()
unsup_minibatch_0 = unsupervised_dataloader_0.next()
unsup_minibatch_1 = unsupervised_dataloader_1.next()

imgs = minibatch['data']
gts = minibatch['label']
unsup_imgs_0 = unsup_minibatch_0['data']
unsup_imgs_1 = unsup_minibatch_1['data']
mask_params = unsup_minibatch_0['mask_params']

• 两个子网络预测没有label的数据，并将预测结果分别作为标签

 batch_mix_masks = mask_params
unsup_imgs_mixed = unsup_imgs_0 * (1 - batch_mix_masks) + unsup_imgs_1 * batch_mix_masks
with paddle.no_grad():
    # Estimate the pseudo-label with branch#1 & supervise branch#2
    logits_u0_tea_1, _ = model(unsup_imgs_0)
    logits_u1_tea_1, _ = model(unsup_imgs_1)
    logits_u0_tea_1 = logits_u0_tea_1.detach()
    logits_u1_tea_1 = logits_u1_tea_1.detach()
    # Estimate the pseudo-label with branch#2 & supervise branch#1
    _, logits_u0_tea_2 = model(unsup_imgs_0)
    _, logits_u1_tea_2 = model(unsup_imgs_1)
    logits_u0_tea_2 = logits_u0_tea_2.detach()
    logits_u1_tea_2 = logits_u1_tea_2.detach()

# Mix teacher predictions using same mask
# It makes no difference whether we do this with logits or probabilities as
# the mask pixels are either 1 or 0
logits_cons_tea_1 = logits_u0_tea_1 * (1 - batch_mix_masks) + logits_u1_tea_1 * batch_mix_masks
ps_label_1 = paddle.argmax(logits_cons_tea_1, axis=1)
logits_cons_tea_2 = logits_u0_tea_2 * (1 - batch_mix_masks) + logits_u1_tea_2 * batch_mix_masks
ps_label_2 = paddle.argmax(logits_cons_tea_2, axis=1)

• 两个子网络各自的无label的图像的输出互为监督，计算cps损失

# Get student#1 and #2 prediction for mixed image
logits_cons_stu_1, logits_cons_stu_2 = model(unsup_imgs_mixed)

cps_loss = criterion_cps(logits_cons_stu_1, ps_label_2) + criterion_cps(logits_cons_stu_2, ps_label_1)
dist.all_reduce(cps_loss, dist.ReduceOp.SUM)
cps_loss = cps_loss / engine.world_size
cps_loss = cps_loss * config.cps_weight

• 两个预测有label的图片，计算交叉熵损失，再联合cps损失一同优化

# supervised loss on both models
sup_pred_l, sup_pred_r = model(imgs)

loss_sup = criterion(sup_pred_l, gts)
dist.all_reduce(loss_sup, dist.ReduceOp.SUM)
loss_sup = loss_sup / engine.world_size

loss_sup_r = criterion(sup_pred_r, gts)
dist.all_reduce(loss_sup_r, dist.ReduceOp.SUM)
loss_sup_r = loss_sup_r / engine.world_size

current_idx = epoch * config.niters_per_epoch + idx
lr = lr_policy.get_lr(current_idx)

optimizer_l.set_lr(lr)
optimizer_r.set_lr(lr)

loss = loss_sup + loss_sup_r + cps_loss
loss.backward()

optimizer_l.step()
optimizer_r.step()

3. 复现结果

• 对于Cityscapes数据集，论文使用8张V100进行CPS模型的训练，每张卡的batch size为1，total batch size为8。由于AI studio脚本任务只能开4卡，复现的每张卡batch size为2，total batch size也为8。复现结果如下表所示

CPS.resnet50.deeplabv3+(1/2 Cityscapes)	mIOU
pytorch	78.77%
paddle	78.28%

注：复现是没有完全按照论文所用的8卡训练，同时CPS的训练不太稳定，相同复现代码，我进行3次训练，不同次数之间的best mIOU的差距能到1% ，源代码的issue中也有人反馈使用源码训练达不到论文精度，本项目使用Paddle复现，mIOU差距在0.49%，算是波动范围内。复现的log文件在work文件夹下

• 部分可视化结果如下，左边为RGB图像，中间为预测图，右边为真值

4. 数据准备

• 使用CPS源代码所提供的Cityscapes数据集，已上传至AI Studio，使用以下命令将数据解压至PDSeg-SemiSeg/DATA文件夹下
• 同时，backbone的预训练权重已准备好在PDSeg-SemiSeg/DATA/paddle-weight文件夹下，准备好的数据组织如下所示

DATA/
|-- city
|-- paddle-weight
|   |-- resnet50_v1c.pdparams

!unzip -qo data/data177911/city.zip -d ./PDSeg-SemiSeg/DATA

%cd PDSeg-SemiSeg/

/home/aistudio/PDSeg-SemiSeg

5. 模型训练

• 配置文件为config.py，模型网络与优化器等参数已设置好，数据读取的.txt文件也已事先准备好，复现不需要进行修改
• 本项目复现的结果为1/2数据量的Cityscapes数据集，若训练其他数据量，只需修改两个变量即可

C.labeled_ratio = 2
C.nepochs = 240

• 数据量和训练epochs的关系如下表所示

Dataset	1/16	1/8	1/4	1/2
Cityscapes	128	137	160	240

• 复现的训练是使用单机四卡的脚本任务进行，notebook的单卡环境无法满足训练需求，此处只说明启动多卡训练的命令

python -m paddle.distributed.launch train.py

6. 模型测试

• 模型训练好之后，通过eval.py脚本对指定的权重在Cityscapes test数据集上进行测试，我已经把最优权重上传至AI Studio，执行以下命令进行测试并保存预测图像

!python eval.py -e ../data/data185470/epoch-last.pdparams --save_path ../work/test_results

• 测试的日志文件保存在PDSeg-SemiSeg/log文件夹下，可以看到mIOU为78.281%