模型定义和训练的一些注意事项

定义了一个动态跟随输入张量的形状调整输入输出通道数的模块，这种做法是有问题的。

错误做法如下：

class ARM(nn.Module):
    def __init__(self, *args, **kwargs):
        super(ARM, self).__init__()
        # 先初始化为None，在第一次forward时初始化
        self.in_chan = None
        self.out_chan = None
        self.conv = None
        self.conv_atten = None
        self.bn_atten = None
        self.sigmoid_atten = nn.Sigmoid()  # 这个不需要通道数，可以先初始化

    def _init_layers(self, x):
        # 根据输入x的通道数动态初始化各层
        self.in_chan = x.size(1)  # 获取输入通道数
        self.out_chan = x.size(1)  # 输出通道数与输入相同

        # 初始化各层
        self.conv = ConvBNReLU(self.in_chan, self.out_chan, ks=3, stride=1, padding=1)
        self.conv_atten = nn.Conv2d(self.out_chan, self.out_chan, kernel_size=1, bias=False)
        self.bn_atten = nn.BatchNorm2d(self.out_chan)

        # 初始化权重
        self.init_weight()

    def init_weight(self):
        if self.conv is not None:
            for ly in [self.conv, self.conv_atten]:
                if hasattr(ly, 'weight'):
                    nn.init.kaiming_normal_(ly.weight, a=1)
            if hasattr(self.bn_atten, 'weight'):
                nn.init.constant_(self.bn_atten.weight, 0)
                nn.init.constant_(self.bn_atten.bias, 0)

    def forward(self, x):
        if self.conv is None:  # 第一次forward时初始化各层
            self._init_layers(x)

        feat = self.conv(x)
        atten = F.avg_pool2d(feat, feat.size()[2:])
        atten = self.conv_atten(atten)
        atten = self.bn_atten(atten)
        atten = self.sigmoid_atten(atten)
        out = torch.mul(feat, atten)
        return out
class YourNet():
	def init_(self):
	    self.module = ARM()
	def forward(x):
	    x = self.module(x)

这样的话，训练可以跑通，但是self.module不会得到保存，也就是训练后的模型权重保存不全，正确做法是一定要初始化你想要保存到权重文件中的模块

正确的一种可选方式：（但不建议，有些画蛇添足，不如老老实实再定义模块的初始化函数中确定好输入输出维度）

class YourModel(nn.Module):
    def __init__(self):
        super(YourModel, self).__init__()
        self.relaxARM1 = ARM()
        self.relaxARM2 = ARM()
        self.relaxARM3 = ARM()
        self.relaxARM4 = ARM()
        self.relaxARM5 = ARM()
        
        # 初始化所有 ARM（假设输入通道数为 3）
        self._initialize_arms(input_channels=3)

    def _initialize_arms(self, input_channels):
        dummy_input = torch.randn(1, input_channels, 32, 32)  # 假设输入尺寸为 32x32
        for name, module in self.named_children():
            if name.startswith("relaxARM"):
                module(dummy_input)  # 触发初始化

    def forward(self, x):
        # 正常调用各 ARM
        x = self.relaxARM1(x)
        x = self.relaxARM2(x)
        # ... 其他调用
        return x