torchvison.models
rgb_encoder = models.resnet101(pretrained=True)
self.conv0_rgb = nn.Sequential(rgb_encoder.conv1, rgb_encoder.bn1, nn.PReLU())
self.conv1_rgb = nn.Sequential(rgb_encoder.maxpool, *rgb_encoder.layer1)
self.conv2_rgb = rgb_encoder.layer2
self.conv3_rgb = rgb_encoder.layer3
self.conv4_rgb = rgb_encoder.layer4
nn.ModuleList
详解PyTorch中的ModuleList和Sequential:
链接: 知乎.
class BasicResNetEncoder(nn.Module):
def __init__(self, in_channel, model):
super(BasicResNetEncoder, self).__init__()
self.encoders = nn.ModuleList(model(in_channel))
def forward(self, x):
outs = []
for en in self.encoders:
x = en(x)
outs.append(x)
return outs
getattr & setattr
self.base_model = getattr(torchvision.models, base_model)(pretrained=True)
feature_dim = getattr(self.base_model, 'fc').in_features
setattr(self.base_model, 'conv1',
nn.Conv2d(2 * self.new_length, 64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False))