PyTorch自定义类U-Net模型summary统计参数量显示为0的原因
问题根因
你的模型检测不到参数的核心原因是所有网络层都没有被注册为模型的子模块。
PyTorch的nn.Module只会统计两类参数:一是在__init__初始化方法中通过self.xxx = 层实例方式绑定为类属性的层的参数,二是被加入到nn.ModuleList/nn.Sequential等容器中的层的参数。
你现在的写法是在conv2d_block方法和forward前向传播过程中临时创建Conv2d、BatchNorm2d、ConvTranspose2d等层实例,这些层是前向传播时的临时变量,用完就会被销毁,根本没有被注册到模型结构里,所以参数统计工具读取不到任何可训练参数。
这种写法除了统计不到参数,还会导致模型完全无法训练:每次前向传播时所有卷积、BN层的权重都是随机初始化的新值,反向传播更新的梯度根本不会被保留。
修正方案
- 把所有层(包括无参数的池化、激活、dropout层)统一定义到
__init__方法中,绑定为self的属性。带训练参数的卷积、BN、转置卷积层必须放在__init__里,否则参数不会被优化器捕获。 - 调整
conv2d_block和forward的逻辑,前向传播时直接调用__init__中已经初始化好的层实例,不要在传播过程中新建层。
修正示例
你可以用nn.Sequential封装重复的卷积块结构,减少冗余代码,示例如下:
import torch import torch.nn as nn import pytorch_model_summary class unet_like(nn.Module): def __init__(self): super(unet_like, self).__init__() # 封装卷积块为Sequential def conv2d_block(in_ch, out_ch): return nn.Sequential( nn.Conv2d(in_ch, out_ch, kernel_size=3, padding="same"), nn.BatchNorm2d(out_ch), nn.ReLU(), nn.Conv2d(out_ch, out_ch, kernel_size=3, padding="same"), nn.BatchNorm2d(out_ch), nn.ReLU() ) # 下采样路径 self.c1 = conv2d_block(3, 16) self.p1 = nn.MaxPool2d(kernel_size=2) self.d1 = nn.Dropout2d(0.1) self.c2 = conv2d_block(16, 32) self.p2 = nn.MaxPool2d(kernel_size=2) self.d2 = nn.Dropout2d(0.1) self.c3 = conv2d_block(32, 64) self.p3 = nn.MaxPool2d(kernel_size=2) self.d3 = nn.Dropout2d(0.1) self.c4 = conv2d_block(64, 128) self.p4 = nn.MaxPool2d(kernel_size=2) self.d4 = nn.Dropout2d(0.1) # 瓶颈层 self.c5 = conv2d_block(128, 256) # 上采样路径 self.up6 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2, output_padding=(0,1)) self.d6 = nn.Dropout2d(0.1) self.c6 = conv2d_block(256, 128) self.up7 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2, output_padding=(1,0)) self.d7 = nn.Dropout2d(0.1) self.c7 = conv2d_block(128, 64) self.up8 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2, output_padding=(0,1)) self.d8 = nn.Dropout2d(0.1) self.c8 = conv2d_block(64, 32) self.up9 = nn.ConvTranspose2d(32, 16, kernel_size=2, stride=2, output_padding=(0,1)) self.d9 = nn.Dropout2d(0.1) self.c9 = conv2d_block(32, 16) # 输出层 self.out_conv = nn.Conv2d(16, 1, kernel_size=3, padding=1) self.sigmoid = nn.Sigmoid() def forward(self, x): c1 = self.c1(x) p1 = self.d1(self.p1(c1)) c2 = self.c2(p1) p2 = self.d2(self.p2(c2)) c3 = self.c3(p2) p3 = self.d3(self.p3(c3)) c4 = self.c4(p3) p4 = self.d4(self.p4(c4)) c5 = self.c5(p4) u6 = self.up6(c5) u6 = torch.cat([u6, c4], dim=1) u6 = self.d6(u6) c6 = self.c6(u6) u7 = self.up7(c6) u7 = torch.cat([u7, c3], dim=1) u7 = self.d7(u7) c7 = self.c7(u7) u8 = self.up8(c7) u8 = torch.cat([u8, c2], dim=1) u8 = self.d8(u8) c8 = self.c8(u8) u9 = self.up9(c8) u9 = torch.cat([u9, c1], dim=1) u9 = self.d9(u9) c9 = self.c9(u9) outputs = self.sigmoid(self.out_conv(c9)) return outputs # 测试模型参数统计 model = unet_like().to("cpu") # 生成和你输入尺寸一致的测试张量 test_input = torch.randn(1, 3, 100, 75) print(pytorch_model_summary.summary(model, test_input, show_input=True))
运行上述修正后的代码,就能正常统计到所有卷积、BN层的参数,模型也可以正常被优化器训练。
注意:你原代码里下采样阶段混用了
nn.Dropout2d和普通nn.Dropout,处理2D图像特征图时建议统一用nn.Dropout2d,普通Dropout是给全连接层设计的,在特征图上使用效果不符合预期。
内容的提问来源于stack exchange,提问作者Hyunjin Lee
相关产品推荐
相关产品推荐

