语义分割丨PSPNet源码解析「网络训练」 (6)

PSPNet输入1x3x769x769，1为BS、3为RGB通道、769为cropsize。并有两个输出1x19x97x97和1x19x769x769，19为类别数，预测了每个位置属于各类的概率。(注意这里尚未softmax，概率之和不为1)。

criterion.py

语义分割的损失函数主要是交叉熵。由于采用了辅助损失，所以Loss应该包含两部分。

\(total\_loss=\alpha \cdot loss1+\beta \cdot loss2\)

此外，这里还定义了OHEM的损失计算，具体实现请看loss.py

import torch.nn as nn import math import torch.utils.model_zoo as model_zoo import torch import numpy as np from torch.nn import functional as F from torch.autograd import Variable from .loss import OhemCrossEntropy2d import scipy.ndimage as nd class CriterionDSN(nn.Module): ''' DSN : We need to consider two supervision for the model. 我们需要考虑两种损失 ''' def __init__(self, ignore_index=255, use_weight=True, reduce=True): super(CriterionDSN, self).__init__() self.ignore_index = ignore_index #交叉熵计算Loss，忽略了255类，并且对Loss取了平均 self.criterion = torch.nn.CrossEntropyLoss(ignore_index=ignore_index, reduce=reduce) if not reduce: print("disabled the reduce.") #criterion(preds, labels) def forward(self, preds, target): h, w = target.size(1), target.size(2) #769, 769 scale_pred = F.upsample(input=preds[0], size=(h, w), mode='bilinear', align_corners=True) loss1 = self.criterion(scale_pred, target) scale_pred = F.upsample(input=preds[1], size=(h, w), mode='bilinear', align_corners=True) loss2 = self.criterion(scale_pred, target) return loss1 + loss2*0.4 #采用难例挖掘 class CriterionOhemDSN(nn.Module): ''' DSN : We need to consider two supervision for the model. ''' def __init__(self, ignore_index=255, thresh=0.7, min_kept=100000, use_weight=True, reduce=True): super(CriterionOhemDSN, self).__init__() self.ignore_index = ignore_index self.criterion1 = OhemCrossEntropy2d(ignore_index, thresh, min_kept) #采用了新的计算方式 self.criterion2 = torch.nn.CrossEntropyLoss(ignore_index=ignore_index, reduce=reduce) def forward(self, preds, target): h, w = target.size(1), target.size(2) #769, 769 scale_pred = F.upsample(input=preds[0], size=(h, w), mode='bilinear', align_corners=True) loss1 = self.criterion1(scale_pred, target) scale_pred = F.upsample(input=preds[1], size=(h, w), mode='bilinear', align_corners=True) loss2 = self.criterion2(scale_pred, target) return loss1 + loss2*0.4 loss.py

OHEM目的是筛选出困难样本来训练模型，从而提升性能，其有两个超参数：\(\theta\)和\(K\)。

困难样本被定义为预测概率小于$\theta \(的像素，并且每个*minibatch*至少保证\)K$个困难样本。