iRMB反向残差移动块代码和解析

import math  # 导入数学库
from functools import partial  # 从functools模块导入partial，用于固定函数参数值
from einops import rearrange  # 从einops库导入rearrange函数，用于重新排列张量
from timm.models.layers.activations import *  # 从timm库中导入所有激活层
from timm.models.layers import DropPath  # 从timm库中导入DropPath，用于实现随机路径丢弃
from timm.models.efficientnet_builder import _parse_ksize  # 导入解析核大小的函数
from timm.models.efficientnet_blocks import num_groups, SqueezeExcite as SE  # 导入组数计算和压缩激励类
 
# ========== 通用层定义区 ==========
class LayerNorm2d(nn.Module):  # 定义二维层归一化类
    def __init__(self, normalized_shape, eps=1e-6, elementwise_affine=True):  # 初始化函数
        super().__init__()  # 调用父类初始化
        self.norm = nn.LayerNorm(normalized_shape, eps, elementwise_affine)  # 创建层归一化实例
 
    def forward(self, x):  # 前向传播函数
        x = rearrange(x, 'b c h w -> b h w c').contiguous()  # 重排输入张量的维度，使之适应层归一化的输入需求
        x = self.norm(x)  # 应用层归一化
        x = rearrange(x, 'b h w c -> b c h w').contiguous()  # 恢复张量的原始维度排列
        return x  # 返回处理后的张量
 
def get_norm(norm_layer='in_1d'):  # 定义函数以获取不同类型的归一化层
    eps = 1e-6  # 设置一个小的epsilon值，防止除以零
    norm_dict = {  # 定义归一化层的字典
        # 各种归一化层和对应的构造函数，使用partial固定一些参数值
        'none': nn.Identity,  # 不进行任何操作的层
        'in_1d': partial(nn.InstanceNorm1d, eps=eps),
        'in_2d': partial(nn.InstanceNorm2d, eps=eps),
        'in_3d': partial(nn.InstanceNorm3d, eps=eps),
        'bn_1d': partial(nn.BatchNorm1d, eps=eps),
        'bn_2d': partial(nn.BatchNorm2d, eps=eps),
        'bn_3d': partial(nn.BatchNorm3d, eps=eps),
        'gn': partial(nn.GroupNorm, eps=eps),
        'ln_1d': partial(nn.LayerNorm, eps=eps),
        'ln_2d': partial(LayerNorm2d, eps=eps),
    }
    return