pytorch 实现预训练模型加载VGG16

1： 下载的预训练模型

# VGG-16
class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True, pretrained=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes  # 类别1
        self.bilinear = bilinear
        # 下载模型
        vgg16 = models.vgg16(pretrained)
        for param in vgg16.parameters():
            param.requires_grad_(False)
        self.vgg16 = vgg16

2： 使用本地下载的模型

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes  # 类别1
        self.bilinear = bilinear
        # 本地模型
        self.vgg16 = models.vgg16(pretrained=False)
        for param in self.vgg16.parameters():
            param.requires_grad_(False)
        self.vgg16.load_state_dict(torch.load('../vgg16-397923af.pth')) 

        # VGG 的预训练加载
        self.inc =  self.vgg16.features[:4]
        self.down1 = self.vgg16.features[4:9]
        self.down2 = self.vgg16.features[9:16]
        self.down3 = self.vgg16.features[16:23]
        self.down4 = self.vgg16.features[23:30]
        
 def forward(self, x):
        outputs = [] # 记录中间的特征图

        x1 = self.inc(x)     # [1, 64, 256, 256]  1/1
        x2 = self.down1(x1)  # [1, 128, 128, 128] 1/2
        x3 = self.down2(x2)  # [1, 256, 64, 64]   1/4
        x4 = self.down3(x3)  # [1, 512, 32, 32]   1/8
        x5 = self.down4(x4)  # [1, 512, 16, 16]   1/16
        return  x5
        
   if __name__ == '__main__':
    x = torch.rand(1,3,256,256)
    net = UNet(1,1)
    # print(net)
    output = net(x)
    print(output.shape)