CCNet（ICCV 2019）

paper：CCNet: Criss-Cross Attention for Semantic Segmentation

official implementation：https://github.com/speedinghzl/CCNet

third-party implementation：https://github.com/open-mmlab/mmcv/blob/main/mmcv/ops/cc_attention.py

存在的问题

传统的FCN网络受局部感受野的限制，无法捕获长距离的上下文信息。现有方法如扩张卷积和池化方法在聚合上下文信息时缺乏适应性，不能满足不同像素对不同上下文依赖的需求。基于图神经网络（GNN）的非局部神经网络虽然能够捕获全图像上下文信息，但计算复杂度高，占用大量GPU内存。

本文的创新点

本篇论文提出了一个新的语义分割网络CCNet，其中主要是提出了一个新的注意力模块Criss-Cross Attention Module（十字交叉注意力模块）。该模块通过在水平和垂直方向上聚合每个像素的上下文信息，通过递归操作使每个像素最终能够捕获全图像依赖关系。

Criss-Cross Attention

作者提出了一个交叉注意力模块，交叉注意力模块收集水平和垂直方向的上下文信息，以增强像素级表示能力。如图3所示

给定一个特征图 \(\mathbf{H}\in \mathbb{R}^{C\times W\times H}\)，首先通过两个1x1卷积得到两个特征图 \(\{\mathbf{Q},\mathbf{K}\}\in \mathbb{R}^{C'\times W\times H}\)，其中 \(C'\) 是通道数在降维时小于 \(C\)。

在得到 \(\mathbf{Q}\) 和 \(\mathbf{K}\) 后，我们通过Affinity operation得到一个attention map \(\mathbf{A}\in \mathbb{R}^{(H+W-1)\times (H\times W)}\)。对于 \(\mathbf{Q}\) 的每个空间位置 \(\mathbf{u}\)，我们可以得到一个向量 \(\mathbf{Q}_{\mathbf{u}}\in \mathbb{R}^{C'}\)。同时我们提取 \(\mathbf{K}\) 中同一位置 \(\mathbf{u}\) 所在行列的所有特征向量得到一个集和 \(\mathbf{\Omega}_{\mathbf{u}}\in \mathbb{R}^{(H+W-1)\times C'}\)。\(\mathbf{\Omega}_{i,\mathbf{u}}\in \mathbb{R}^{C'}\) 是 \(\mathbf{Q}_{\mathbf{u}}\) 的第 \(i\) 个元素。Affinity operation定义如下

其中 \(d_{i,\mathbf{u}}\in \mathbf{D}\) 是特征 \(\mathbf{Q}_{\mathbf{u}}\) 和 \(\mathbf{\Omega}_{i,\mathbf{u}},\ i=[1,...,H+W-1]\) 之间的相关程度，\(\mathbf{D}\in \mathbb{R}^{(H+W-1)\times (W\times H)}\)。然后我们沿 \(\mathbf{D}\) 的通道维度计算softmax就得到了attention map \(\mathbf{A}\)。

\(\mathbf{H}\) 经过另一个1x1卷积得到 \(\mathbf{V}\in \mathbb{R}^{C\times W\times H}\)。对于 \(\mathbf{V}\) 的每个空间位置 \(\mathbf{u}\)，我们可以得到一个向量 \(\mathbf{V}_{\mathbf{u}}\in \mathbb{R}^{C}\) 和 \(\mathbf{u}\) 所在行列所有特征的集和 \(\mathbf{\Phi}_{\mathbf{u}}\in \mathbb{R}^{(H+W-1)\times C}\)。然后按下式定义的Aggregation operation来聚合上下文信息

 

其中 \(\mathbf{H}'_{\mathbf{u}}\) 是 \(\mathbf{H}'\in \mathbb{R}^{C\times W\times H}\) 中位置 \(\mathbf{u}\) 处的特征向量，\(\mathbf{A}_{i,\mathbf{u}}\) 是 \(\mathbf{A}\) 中位置 \(\mathbf{u}\) 通道 \(i\) 处的标量值。将上下文信息添加到局部特征 \(\mathbf{H}\) 中，以增强像素级表示。因此，它具有广泛的上下文视角，并根据空间注意图选择性地聚合上下文。

代码

官方实现将十字交叉注意力分成横向、纵向两部分，其中将对角线的值设为INF是防止十字交叉的中心点计算两次注意力。 

import torch
import torch.nn as nn
from torch.nn import Softmax


def INF(B, H, W):  # 8,60,60
    # ()->(H)->(H,H)->(1,H,H)->(BW,H,H)
    return -torch.diag(torch.tensor(float("inf")).repeat(H), 0).unsqueeze(0).repeat(B * W, 1, 1)


class CrissCrossAttention(nn.Module):
    """Criss-Cross Attention Module.

    .. note::
        Before v1.3.13, we use a CUDA op. Since v1.3.13, we switch
        to a pure PyTorch and equivalent implementation. For more
        details, please refer to https://github.com/open-mmlab/mmcv/pull/1201.

        Speed comparison for one forward pass

        - Input size: [2,512,97,97]
        - Device: 1 NVIDIA GeForce RTX 2080 Ti

        +-----------------------+---------------+------------+---------------+
        |                       |PyTorch version|CUDA version|Relative speed |
        +=======================+===============+============+===============+
        |with torch.no_grad()   |0.00554402 s   |0.0299619 s |5.4x           |
        +-----------------------+---------------+------------+---------------+
        |no with torch.no_grad()|0.00562803 s   |0.0301349 s |5.4x           |
        +-----------------------+---------------+------------+---------------+

    Args:
        in_channels (int): Channels of the input feature map.
    """

    def __init__(self, in_channels: int) -> None:
        super().__init__()
        self.query_conv = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.key_conv = nn.Conv2d(in_channels, in_channels // 8, 1)
        self.value_conv = nn.Conv2d(in_channels, in_channels, 1)
        self.in_channels = in_channels  # 512

    def forward(self, x):
        self.softmax = Softmax(dim=3)
        m_batchsize, _, height, width = x.size()  # (8,512,60,60)
        proj_query = self.query_conv(x)  # (8,64,60,60)
        # (b,c,h,w)->(b,w,c,h)->(bw,c,h)->(bw,h,c)
        proj_query_H = proj_query.permute(0, 3, 1, 2).contiguous().view(m_batchsize*width, -1, height).permute(0, 2, 1)  # (bw,h,c),(480,60,64)
        # (b,c,h,w)->(b,h,c,w)->(bh,c,w)->(bh,w,c)
        proj_query_W = proj_query.permute(0, 2, 1, 3).contiguous().view(m_batchsize*height, -1, width).permute(0, 2, 1)  # (bh,w,c),(480,60,64)
        proj_key = self.key_conv(x)  # (8,64,60,60)
        # (b,c,h,w)->(b,w,c,h)->(bw,c,h)
        proj_key_H = proj_key.permute(0, 3, 1, 2).contiguous().view(m_batchsize*width, -1, height)  # (bw,c,h),(480,64,60)
        # (b,c,h,w)->(b,h,c,w)->(bh,c,w)
        proj_key_W = proj_key.permute(0, 2, 1, 3).contiguous().view(m_batchsize*height, -1, width)  # (bh,c,w),(480,64,60)
        proj_value = self.value_conv(x)  # (8,512,60,60)
        # (b,h,c,w)->(b,w,c,h)->(bw,c,h)
        proj_value_H = proj_value.permute(0, 3, 1, 2).contiguous().view(m_batchsize*width, -1, height)  # (bw,c,h),(480,512,60)
        # (b,c,h,w)->(b,h,c,w)->(bh,c,w)
        proj_value_W = proj_value.permute(0, 2, 1, 3).contiguous().view(m_batchsize*height, -1, width)  # (bh,c,w),(480,512,60)
        # (bw,h,c)@(bw,c,h)->(bw,h,h)=(480,60,60)
        # INF->(bw,h,h)=(480,60,60)
        # (bw,h,h)->(b,w,h,h)->(b,h,w,h)
        energy_H = (torch.bmm(proj_query_H, proj_key_H) + INF(m_batchsize, height, width)).view(m_batchsize, width, height, height).permute(0, 2, 1, 3)  # (8,60,60,60)
        # (bh,w,c)@(bh,c,w)->(bh,w,w)->(b,h,w,w)
        energy_W = torch.bmm(proj_query_W, proj_key_W).view(m_batchsize, height, width, width)  # (8,60,60,60)
        # [(b,h,w,h),(b,h,w,w)]->(b,h,w,h+w)
        # 实际应该是h+w-1，对角线设为-inf了，在softmax中不起作用，防止十字交叉中心点算两次
        concate = self.softmax(torch.cat([energy_H, energy_W], 3))

        att_H = concate[:, :, :, 0: height].permute(0, 2, 1, 3).contiguous().view(m_batchsize*width, height, height)
        # (b,h,w,h)->(b,w,h,h)->(bw,h,h)
        att_W = concate[:, :, :, height:height+width].contiguous().view(m_batchsize*height, width, width)
        # (b,h,w,w)->(bh,w,w)
        out_H = torch.bmm(proj_value_H, att_H.permute(0, 2, 1)).view(m_batchsize, width, -1, height).permute(0, 2, 3, 1)
        # (bw,c,h)@(bw,h,h)->(bw,c,h)->(b,w,c,h)->(b,c,h,w)
        out_W = torch.bmm(proj_value_W, att_W.permute(0, 2, 1)).view(m_batchsize, height, -1, width).permute(0, 2, 1, 3)
        # (bh,c,w)@(bh,w,w)->(bh,c,w)->(b,h,c,w)->(b,c,h,w)
        return self.gamma*(out_H + out_W) + x



    def __repr__(self) -> str:
        s = self.__class__.__name__
        s += f'(in_channels={self.in_channels})'
        return s


if __name__ == '__main__':
    model = CrissCrossAttention(512)
    x = torch.randn(8, 512, 60, 60)
    out = model(x)
    print(out.shape)