每日Attention学习12——Exterior Contextual-Relation Module

模块出处

[ISBI 22] [link] [code] Duplex Contextual Relation Network for Polyp Segmentation

---

模块名称

Exterior Contextual-Relation Module (ECRM)

---

模块作用

内存型特征增强模块

---

模块结构

---

模块思想

原文表述：在临床环境中，不同样本之间存在息肉的同步视觉模式。基于这一关键观察，属于所有训练数据的同一语义类的区域特征应该具有上下文关系。因此，我们提出了一种新颖的跨不同样本的上下文关系探索模块。
具体做法则是，对于编码器最后一层得到的全局特征(图中红色方块)，进行两次增强：
第一次是直接将全局特征送入一个$1\times 1$卷积(图中浅紫色部分)以获取一个粗糙分割mask，该mask与全局特征相乘后便能得到过滤掉背景特征的增强特征(图中enqueue左边的部分)。
第二次增强则是基于网络存储的源自其他训练样本的历史上下文信息(图中的Cross-Batch Memory)。即，当前特征与Memory内特征进行Cross Attention操作，从而利用历史经验对当前状态进行补全。

---

模块代码

代码实现有几个额外要注意的地方：

• 模块返回的aux_out要进行side supervision监督，以保证准确性；
• Memory负责维护网络的历史信息，为防止被破坏，这部分信息并不参与梯度更新过程；
• 在测试阶段，Memory不再更新，直接使用训练所存储的历史信息，这一思想与BatchNorm类似。

import torch
from torch import nn

def conv2d(in_channel, out_channel, kernel_size):
    layers = [
        nn.Conv2d(
            in_channel, out_channel, kernel_size, padding=kernel_size // 2, bias=False
        ),
        nn.BatchNorm2d(out_channel),
        nn.ReLU(),
    ]
    return nn.Sequential(*layers)


def conv1d(in_channel, out_channel):
    layers = [
        nn.Conv1d(in_channel, out_channel, 1, bias=False),
        nn.BatchNorm1d(out_channel),
        nn.ReLU(),
    ]
    return nn.Sequential(*layers)


class ECRM(nn.Module):
    def __init__(self, bank_size=20, feat_channels=512, num_classes=1):
        super(ECRM, self).__init__()  
        # BANK CONFIG
        self.bank_size = bank_size
        self.register_buffer("bank_ptr", torch.zeros(1, dtype=torch.long))  # memory bank pointer
        self.register_buffer("bank", torch.zeros(self.bank_size, feat_channels, num_classes))  # memory bank
        self.bank_full = False

        # ATTENTION CONFIG
        self.feat_channels = feat_channels
        self.L = nn.Conv2d(feat_channels, num_classes, 1)
        self.X = conv2d(feat_channels, 512, 3)
        self.phi = conv1d(512, 256)
        self.psi = conv1d(512, 256)
        self.delta = conv1d(512, 256)
        self.rho = conv1d(256, 512)
        self.g = conv2d(512 + 512, 512, 1)

    def init(self):
        self.bank_ptr[0] = 0
        self.bank_full = False

    @torch.no_grad()
    def update_bank(self, x):
        ptr = int(self.bank_ptr)
        batch_size = x.shape[0]
        vacancy = self.bank_size - ptr
        if batch_size >= vacancy:
            self.bank_full = True
        pos = min(batch_size, vacancy)
        self.bank[ptr:ptr+pos] = x[0:pos].clone()
        # update pointer
        ptr = (ptr + pos) % self.bank_size
        self.bank_ptr[0] = ptr

    def enhance_by_memory(self, bank, X_flat, X):
        batch, n_class, height, width = X.shape
        # query = S * C
        query = self.phi(bank).squeeze(dim=2)
        # key: = B * C * HW
        key = self.psi(X_flat)
        # logit = HW * S * B (cross image relation)
        logit = torch.matmul(query, key).transpose(0,2)
        # attn = HW * S * B
        attn = torch.softmax(logit, 2)
        # delta = S * C
        delta = self.delta(bank).squeeze(dim=2)
        # attn_sum = B * C * HW
        attn_sum = torch.matmul(attn.transpose(1,2), delta).transpose(1,2)
        # x_obj = B * C * H * W
        X_obj = self.rho(attn_sum).view(batch, -1, height, width)
        concat = torch.cat([X, X_obj], 1)
        out = self.g(concat)
        return out
    
    def get_prototype(self, input):
        L = self.L(input)
        aux_out = L
        batch, n_class, _, _ = L.shape
        l_flat = L.view(batch, n_class, -1)
        M = torch.softmax(l_flat, -1)
        X = self.X(input)
        channel = X.shape[1]
        X_flat = X.view(batch, channel, -1)
        f_k = (M @ X_flat.transpose(1, 2)).transpose(1, 2)
        return aux_out, f_k, X_flat, X

    def forward(self, x, flag='train'):
        # x [3, 512, 11, 11]
        # patch [3, 512, 1]
        aux_out, patch, feats_flat, feats = self.get_prototype(x)
        if flag == 'train':
            self.update_bank(patch)
            ptr = int(self.bank_ptr)
            if self.bank_full == True:
                out = self.enhance_by_memory(self.bank, feats_flat, feats)
            else:
                out = self.enhance_by_memory(self.bank[0:ptr], feats_flat, feats)
        elif flag == 'test':
            out = self.enhance_by_memory(patch, feats_flat, feats)
        return out, aux_out
    
if __name__ == '__main__':
    x = torch.randn([3, 512, 11, 11])
    ecrm = ECRM()
    out = ecrm(x)
    print(out[0].shape)  # 3, 512, 11, 11
    print(out[1].shape)  # 3, 1, 11, 11

---