本文是论文 RobustNet: Improving Domain Generalization in Urban-Scene Segmentation via Instance Selective Whitening (ISW) 的解析。ISW 的贡献是提出了一个能够有选择性的白化协方差矩阵部分区域的损失函数，论文行文脉络十分清晰。

Instance Whitening Loss

首先，论文指出了 deep whitening transformation (DWT) （通过设计损失函数使得各特征图之间组成的协方差矩阵主对角线元素为 1，其它元素为 0）不能同时优化对角线元素和其它元素的缺陷，使用 IN 先将协方差矩阵的对角线元素归一，这样之后只需要优化非对角线元素即可，如下图 (a) 所示。

Instance Whitening 最基本的思路就是将协方差矩阵非对角线元素 -> 0，通过 $L I W L_{IW}$ LIW 损失函数优化即可，如下图 (c) (d) 所示。

源码中每过一次 IW 就做一次 IN 操作，并额外返回经过 IN 处理的特征图

python 复制代码

class InstanceWhitening(nn.Module):

    def __init__(self, dim):
        super(InstanceWhitening, self).__init__()
        self.instance_standardization = nn.InstanceNorm2d(dim, affine=False)

    def forward(self, x):

        x = self.instance_standardization(x)
        w = x

        return x, w

在模型中保留这些特征图，用于计算 loss

python 复制代码

for module in i_block:
    if isinstance(module, InstanceWhitening):
        x, w = module(x)
        w_arr.append(w)
...
data_dict['w_arr'] = w_arr

计算 loss，传入的 mask 是一个 C x C 大小的主对角线元素为 0，其它元素为 1 的矩阵

python 复制代码

w_arr = output_dict['w_arr']
...
wt_loss = torch.FloatTensor([0]).cuda()
for index, f_map in enumerate(w_arr):
    B, C, H, W = f_map.shape
    M_ones = torch.ones(C,C).cuda()
    diag = torch.diag(M_ones)
    diag = torch.diag_embed(diag)
    M_ones = M_ones - diag
    loss = instance_whitening_loss(f_map, None, M_ones, 0, 10000)
    wt_loss = wt_loss + loss
wt_loss = wt_loss / len(w_arr)
total_loss += wt_loss.item()

$L I W L_{IW}$ LIW 源码，协方差矩阵通过矩阵乘积得到，与掩码矩阵逐位相乘，得到需要优化的协方差矩阵，这里其实应该是协方差矩阵的上三角，传入的掩码应该下三角为 0

python 复制代码

def instance_whitening_loss(f_map, eye, mask_matrix, margin, num_remove_cov):
    f_cor, B = get_covariance_matrix(f_map, eye=eye)
    f_cor_masked = f_cor * mask_matrix

    off_diag_sum = torch.sum(torch.abs(f_cor_masked), dim=(1,2), keepdim=True) - margin # B X 1 X 1
    loss = torch.clamp(torch.div(off_diag_sum, num_remove_cov), min=0) # B X 1 X 1
    loss = torch.sum(loss) / B

    return loss


def get_covariance_matrix(f_map, eye=None):
    eps = 1e-5
    B, C, H, W = f_map.shape  # i-th feature size (B X C X H X W)
    HW = H * W
    if eye is None:
        eye = torch.eye(C).cuda()
    f_map = f_map.contiguous().view(B, C, -1)  # B X C X H X W > B X C X (H X W)
    f_cor = torch.bmm(f_map, f_map.transpose(1, 2)).div(HW-1) + (eps * eye)  # B X C X C / HW

    return f_cor, B

Margin-based relaxation of whitening loss

作者认为将协方差矩阵的非对角线元素全部优化为 0，会影响模型的鉴别能力，因此设计了一个 margin 参数，在上面的代码中已有体现：off_diag_sum = torch.sum(torch.abs(f_cor_masked), dim=(1,2), keepdim=True) - margin

Separating Covariance Elements

又到了特征解耦的时候了，本文的出发点是对原始数据引入一个光照变换，比较原图和数据增强后图像对应特征图的协方差矩阵，差异较小的部分认为是 domain-invariant 部分，其它部分为 domain-specific 部分。通过这个操作来得到一个 Selective 的掩码矩阵，只对原始协方差矩阵的这些部分做优化。

论文解析 | RobustNet / ISW

Instance Whitening Loss

Margin-based relaxation of whitening loss

Separating Covariance Elements