【implementation】SimAM

【ICML-2021】SimAM

架构

能量函数

最小化能量函数的方程

最终的输出结果

$$\begin{aligned} \frac{1}{e} &= \frac{(t-\hat \mu)^2}{4(\hat \sigma^2 + \lambda)} + \frac{2(\hat \sigma^2 + \lambda)}{4(\hat \sigma^2 + \lambda)} &= \frac{(t-\hat \mu)^2}{4(\hat \sigma^2 + \lambda)} + \frac {1}{2} \end{aligned}$$

实现

import torch
from torch import nn
import torch.nn.functional as F

class SimAM(nn.Module):
    def __init__(self,Lambda=0.1):
        """
        :param Lambda: 能量函数里面的超参数
        """
        super(SimAM, self).__init__()
        self.Lambda = Lambda

    def forward(self,x):
        n,c,h,w = x.shape
        M = h * w - 1
        t_minus_mu_square = x - x.mean(dim=[2,3],keepdim = True).pow(2)
        sigma_square = torch.sum(t_minus_mu_square) / M
        one_divide_e = t_minus_mu_square / (4 * sigma_square + self.Lambda) + 0.5
        out = F.sigmoid(one_divide_e) * x
        return out

if __name__ == '__main__':
    x = torch.randn(2,8,32,32)
    print("iutput shape:{}".format(x.shape))
    model = SimAM()
    print("output shape:{}".format(model(x).shape))

FAQ