fix(backbones): 修复 shufflenet 无法构建,并让换骨架不再踩配置合并的坑
- shufflenet_v2_x1_0 列在支持清单里却直接 AttributeError:ShuffleNetV2 既不是 扁平 children 也没有 .features,需要按 conv1/maxpool/stage2..4/conv5 自己拼。 该 bug 同时影响 SimCLR 与检测模型 - 新增参数化测试覆盖清单里的每个 arch:能前向、声明的 out_channels 与实际一致 - 配置支持 `_replace_: true`:子字典整块替换而不与 base 合并。换骨架时 tiny_conv 的 widths/depths 会残留并传给新骨架,此前无论用 --set 还是 _base_ 都绕不开,只能改原文件 - Registry.build 提前校验字段,报错直接指出是哪些字段不被接受、该怎么办, 取代裸 TypeError - README 补充可选骨架表与三种改法(含不该用 --set 改 name 的说明) Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
e4bd794025
commit
e581d9b3d7
@@ -45,10 +45,7 @@ class TorchvisionBackbone(nn.Module):
|
||||
weights = "DEFAULT" if pretrained else None
|
||||
net = getattr(torchvision.models, arch)(weights=weights)
|
||||
|
||||
if arch.startswith("resnet"):
|
||||
self.features = nn.Sequential(*list(net.children())[:-2])
|
||||
else:
|
||||
self.features = net.features
|
||||
self.features = self._extract_features(net, arch)
|
||||
|
||||
if in_channels != 3:
|
||||
self._patch_first_conv(in_channels)
|
||||
@@ -60,6 +57,25 @@ class TorchvisionBackbone(nn.Module):
|
||||
for param in self.parameters():
|
||||
param.requires_grad_(False)
|
||||
|
||||
@staticmethod
|
||||
def _extract_features(net: nn.Module, arch: str) -> nn.Sequential:
|
||||
"""取到最后一个卷积 stage 为止,丢掉分类头。
|
||||
|
||||
不同系列的 torchvision 模型结构不一样:resnet 是扁平的 children,
|
||||
mobilenet 有现成的 `.features`,而 ShuffleNetV2 两者都不是——它把层
|
||||
散放在 conv1/maxpool/stage2..4/conv5 上,得按顺序自己拼。
|
||||
"""
|
||||
if arch.startswith("resnet"):
|
||||
return nn.Sequential(*list(net.children())[:-2])
|
||||
if arch.startswith("shufflenet"):
|
||||
return nn.Sequential(
|
||||
net.conv1, net.maxpool, net.stage2, net.stage3, net.stage4, net.conv5
|
||||
)
|
||||
features = getattr(net, "features", None)
|
||||
if features is None:
|
||||
raise ValueError(f"不知道如何从 {arch} 取特征,请在 _extract_features 里补一条分支")
|
||||
return features
|
||||
|
||||
def _patch_first_conv(self, in_channels: int) -> None:
|
||||
for module in self.features.modules():
|
||||
if isinstance(module, nn.Conv2d):
|
||||
|
||||
Reference in New Issue
Block a user