#打算取VGG19的第二个全连接层的输出,那么就需要构建一个类,这个类要包含VGG的全部卷积层,
#以及到第二个全连接层的全部网络还有他们对应的参数
class Classification_att(nn.Module):
def __init__(self, rgb_range):
super(Classification_att, self).__init__()
self.vgg19 =models.vgg19(pretrained=True)
vgg = models.vgg19(pretrained=True).features
conv_modules = [m for m in vgg]
self.vgg_conv = nn.Sequential(*conv_modules[:37])
classfi = models.vgg19(pretrained=True).classifier
classif_modules = [n for n in classfi]
self.vgg_class = nn.Sequential(*classif_modules[:4])
vgg_mean = (0.485, 0.456, 0.406)
vgg_std = (0.229 * rgb_range, 0.224 * rgb_range, 0.225 * rgb_range)
self.sub_mean = common.MeanShift(rgb_range, vgg_mean, vgg_std)
for p in self.vgg_conv.parameters():
p.requires_grad = False
for p in self.vgg_class.parameters():
p.requires_grad = False
self.classifi = nn.Sequential(
nn.Linear(4096, 1024),
nn.ReLU(True),
nn.Linear(1024, 256),
nn.ReLU(True),
nn.Linear(256, 64),
)
def forward(self, x):
x = F.interpolate(x, size=[224, 224], scale_factor=None, mode='bilinear',
align_corners=False)
x = self.sub_mean(x)
x = self.vgg_conv(x)
x = self.vgg_class(x) #执行这部报错,说张量不匹配
原因是因为卷积层的输出不能直接连接全连接层,即使输出的张量的总的大小是一致的
查看vgg的pytorch源码发现是
x = self.features(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
#自己的代码没有torch.flatten(x, 1)这步
所以自己的少了一步
x = torch.flatten(x, 1)
补上就好了