CAN源码解析（HMER:Counting-Aware Network for HandwrittenMathematical Expression Recognition）

CAN源码解析（HMER:Counting-Aware Network for HandwrittenMathematical Expression Recognition）

paper: https://arxiv.org/abs/2207.11463

GitHub: https://github.com/LBH1024/CAN

0 整体架构

CAN是一个基于encoder-decoder架构的模型。整体由3个组成部分构成

• backbone (encoder)
• MSCM (辅助任务head，其输出也作为CCAD的输入)
• CCAD (decoder)

1 Encoder

1.1 输入

• Images ${\mathbb{R}}^{B\times C\times H\times W}$
• Image_masks ${\mathbb{R}}^{B\times 1\times H\times W}$ padding部分为0，图片区域为1
• labels ${\mathbb{R}}^{B\times L}$, latex字符在字典中的索引，（label后加停止符EOS，有些实现还会加起始符） (在decoder模块中使用到)
• labels mask. ${\mathbb{R}}^{B\times L}$ padding部分为0，图片区域为1 (在计算loss中使用到中使用到)

1.2 输出

• Cnn_features ${\mathbb{R}}^{B\times C^{\prime }\times \frac{H}{16}\times \frac{W}{16}}$
• Counting_mask ${\mathbb{R}}^{B\times C^{\prime }\times \frac{H}{16}\times \frac{W}{16}}$ 是输入的images mask间隔16采样得到 counting_mask = images_mask[:, :, ::self.ratio, ::self.ratio]

2 decoder

2.1 MSCM(Multi-Scale Counting Module)

    def forward(self, x, mask):
        """
        Args:
            x(tensor), cnn_feature, B, C',H//16, W//16 
            mask(tensor), counting_mask, B, 1, H//16, W//16
        """
        b, c, h, w = x.size()
        x = self.trans_layer(x)
        x = self.channel_att(x)
        x = self.pred_layer(x)
        if mask is not None:
            x = x * mask
        x = x.view(b, self.out_channel, -1)
        x1 = torch.sum(x, dim=-1)
        return x1, x.view(b, self.out_channel, h, w)

2.1.1 输入

• Cnn_features ${\mathbb{R}}^{B\times C^{\prime }\times \frac{H}{16}\times \frac{W}{16}}$
• Counting_mask ${\mathbb{R}}^{B\times C^{\prime }\times \frac{H}{16}\times \frac{W}{16}}$ 是输入的images mask间隔16采样得到 counting_mask = images_mask[:, :, ::self.ratio, ::self.ratio]

2.1.2 输出

• Counting_preds ${\mathbb{R}}^{B\times C_{token}}$, $C_{token}$是字典的长度

2.2 CCAD(Counting-Combined Attentional Decoder)

论文中的图片架构图如下

2.2.1 输入

• Cnn_features ${\mathbb{R}}^{B\times C^{\prime }\times \frac{H}{16}\times \frac{W}{16}}$
• Image_masks ${\mathbb{R}}^{B\times 1\times H\times W}$ padding部分为0，图片区域为1
• labels ${\mathbb{R}}^{B\times L}$, latex字符在字典中的索引，$L$是序列长度（label后加停止符EOS，有些实现还会加起始符）
• labels mask. ${\mathbb{R}}^{B\times L}$ padding部分为0，图片区域为1 （用于loss的计算）
• Counting_preds ${\mathbb{R}}^{B\times C_{token}}$, $C_{token}$是字典的长度（是上面MSCM的输出）

2.2.2 输出

• word_probs ${\mathbb{R}}^{B\times L\times C_{token}}$, $C_{token}$是字典的长度, $L$是序列长度
• (Optional) Word_alpha ${\mathbb{R}}^{B\times \frac{H}{16}\times \frac{W}{16}}$, attention map （用于可视化）

2.2.3 CCAD内部处理流程详细解析

内部的详细pipeline如下

3 其它

3.1 位置编码怎么做的

位置编码用了常规的正余弦编码

class PositionEmbeddingSine(nn.Module):
    """
    This is a more standard version of the position embedding, very similar to the one
    used by the Attention is all you need paper, generalized to work on images.
    """
    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
        super().__init__()
        self.num_pos_feats = num_pos_feats
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and normalize is False:
            raise ValueError("normalize should be True if scale is passed")
        if scale is None:
            scale = 2 * math.pi
        self.scale = scale

    def forward(self, x, mask):
        """
        Args:
            x: input features with shape [batch_size, channel, height, width]
            mask: mask with shape [batch_size, 1, height, width]
        """
        y_embed = mask.cumsum(1, dtype=torch.float32)  # [batch_size, 1, height, width]
        x_embed = mask.cumsum(2, dtype=torch.float32)  # [batch_size, 1, height, width]
        if self.normalize:
            eps = 1e-6
            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
        
        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)   
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)

        pos_x = x_embed[:, :, :, None] / dim_t  # pos_x.shape: [batch_size, 1, height, width, num_pos_feats]
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)  # pos_x.shape: [batch_size, 1, height, width, num_pos_feats]
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        return pos  # [batch_size, num_pos_feats*2, height, width]

3.2 损失函数

整体损失函数包括两个部分，其一是计数模块的smooth L1损失。其二是序列预测的交叉熵损失
$$\mathcal{L}={\mathcal{L}}_{cls}+{\mathcal{L}}_{counting},$$