pytorch模型加载的一些方式

模型加载

  • torchvison.models
  • nn.ModuleList
  • getattr & setattr

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))

你可能感兴趣的:(VOS)