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)
