Shaun ZhouandClaude Opus 5
6e543660ce
resnet101 此前不在 TorchvisionBackbone 的支持列表里，先补进 _SUPPORTED （out_channels 2048），随后被既有的参数化测试自动覆盖：能前向、声明维度与实际一致。 configs/cluster/film_simclr_resnet101.yaml 仿照 resnet18/50 两份，只改骨架与 batch_size，其余超参一致以便横向对比： - 42.5M 参数（resnet18 11.2M / resnet50 23.5M），layer2+3 仍是 1536 维， 与 resnet50 相同，差别在深度（layer3 有 23 个 block 而非 6 个） - batch_size 取 128：MPS 上 resnet50 在 1024 已 OOM，resnet101 约重 1.8 倍故再降一档； 注释里说明该数字来自 MPS，16 GB 独显上可往上调 - pretrained 首次使用会下载约 171 MB 权重 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
"""torchvision 卷积骨架封装（resnet / mobilenet 等纯卷积模型）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import torch
|
||
from torch import nn
|
||
|
||
from models.registry import BACKBONES
|
||
|
||
# 仅允许纯卷积骨架；ViT/Swin 等不在部署范围内。
|
||
_SUPPORTED = {
|
||
"resnet18": 512,
|
||
"resnet34": 512,
|
||
"resnet50": 2048,
|
||
"resnet101": 2048,
|
||
"mobilenet_v3_small": 576,
|
||
"mobilenet_v3_large": 960,
|
||
"shufflenet_v2_x1_0": 1024,
|
||
}
|
||
|
||
|
||
@BACKBONES.register("torchvision")
|
||
class TorchvisionBackbone(nn.Module):
|
||
"""输出 stride=32 的单尺度特征。
|
||
|
||
Args:
|
||
arch: `_SUPPORTED` 中的模型名。
|
||
pretrained: 是否加载 ImageNet 权重（需要联网或本地缓存）。
|
||
in_channels: 输入通道，非 3 时会把第一层卷积权重按通道求均值复制。
|
||
frozen: 是否冻结骨架参数。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
arch: str = "resnet18",
|
||
pretrained: bool = True,
|
||
in_channels: int = 3,
|
||
frozen: bool = False,
|
||
) -> None:
|
||
super().__init__()
|
||
if arch not in _SUPPORTED:
|
||
raise ValueError(f"不支持的骨架 {arch}，可选 {sorted(_SUPPORTED)}")
|
||
|
||
import torchvision
|
||
|
||
weights = "DEFAULT" if pretrained else None
|
||
net = getattr(torchvision.models, arch)(weights=weights)
|
||
|
||
self.features = self._extract_features(net, arch)
|
||
|
||
if in_channels != 3:
|
||
self._patch_first_conv(in_channels)
|
||
|
||
self.out_channels = _SUPPORTED[arch]
|
||
self.out_stride = 32
|
||
|
||
if frozen:
|
||
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):
|
||
weight = module.weight.data.mean(dim=1, keepdim=True).repeat(1, in_channels, 1, 1)
|
||
module.in_channels = in_channels
|
||
module.weight = nn.Parameter(weight * (3.0 / in_channels))
|
||
return
|
||
raise RuntimeError("骨架中未找到 Conv2d，无法调整输入通道")
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
return self.features(x)
|