电脑基础 · 2023年4月3日

深监督,辅助损失,auxiliary loss

今天看到了在PSPNet文章中使用了深监督即辅助损失,特此来记录一下。

辅助损失由Deeply-Supervised Nets提出的,经过Training Deeper Convolutional Networks with Deep Supervision改良

在PSPnet中,在Resnet的stage3最后一层的卷积处使用了辅助损失,是为了解决反向传播不能传递到浅层,加了辅助损失,两个损失都可以传递到之前的层,辅助损失优化了学习过程,主分支损失承担了主要的责任,且对辅助损失加了一个权重。

在训练过程中我们使用辅助损失,在测试阶段我们不使用辅助损失,使用经过辅助损失优化好的网络。

深监督,辅助损失,auxiliary loss

 那么辅助损失是如何实现的呢?pytorch代码实现

在PSPnet网络定义的类中定义了辅助损失:

        if self.training:
            self.aux = nn.Sequential(
                nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(256),
                nn.ReLU(inplace=True),
                nn.Dropout2d(p=dropout),
                nn.Conv2d(256, classes, kernel_size=1)

根据代码我们可以看到,辅助损失由两个卷积构成,第一个卷积的输入就是resnet第三层卷积的输入为1024,然后辅助损失再经过一个输出为class的卷积。也就是说和主干网络的输出一样。

深监督,辅助损失,auxiliary loss

 看完定义后看如何使用:

在网络的forward函数中,对经过第四层的卷积输出,输入进辅助损失函数中,得到输出再经过上采样到原始图片大小,接着进行两个损失计算,输出x与原始标签计算主损失,辅助损失输出与原始标签计算辅助损失。最后我们PSPnet得到输出x,主损失,辅助损失。

        if self.training:
            aux = self.aux(x_tmp)
            if self.zoom_factor != 1:
                aux = F.interpolate(aux, size=(h, w), mode='bilinear', align_corners=True)
            main_loss = self.criterion(x, y)
            aux_loss = self.criterion(aux, y)
            return x.max(1)[1], main_loss, aux_loss
        else:
            return x

在train.py中,我们实例化model=PSPnet,将输入和标签输送进model中,总损失就是主损失加辅助损失乘以权重。然后进行梯度反向传播。

        output, main_loss, aux_loss = model(input, target)
        if not args.multiprocessing_distributed:
            main_loss, aux_loss = torch.mean(main_loss), torch.mean(aux_loss)
        loss = main_loss + args.aux_weight * aux_loss
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在验证中不使用辅助损失:

    model.eval()
    end = time.time()
    for i, (input, target) in enumerate(val_loader):
        data_time.update(time.time() - end)
        input = input.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)
        output = model(input)
        if args.zoom_factor != 8:
            output = F.interpolate(output, size=target.size()[1:], mode='bilinear', align_corners=True)
        loss = criterion(output, target)

辅助损失的效果:

辅助损失要乘以一个权重,那么这个权重为多少合适,对辅助损失的消融实验:

在0-1之间尝试,发现0.4效果最好。

深监督,辅助损失,auxiliary loss