KeyError: 'conv4_2' while find the loss in style loss function and content loss function

Question

在为“使用 Pytorch（深度学习）进行神经风格迁移”进行编码时，我试图找到 Content Loss 和 Style Loss，但在执行代码期间我遇到了一个 KeyError: 'conv4_2' 我无法解决。

请看下面的代码：

# Get content, style features and create gram matrix 

def get_features(image,model):

  layers = { 
      '0' : 'conv1_1',
      '5' : 'conv2_1',
      '10' : 'conv3_1',
      '19' : 'conv4_1',
      '21' : 'conv4_2', # content_feature
      '28' : 'conv5_1'

  }

  x  = image 
  
  Features = {}

  for name, layer in model._modules.items():
    x = layer(x)
    if name in layers:
      Features[layers[name]] = x

    return Features

content_f = get_features(content_p,vgg)
style_f = get_features(style_p,vgg)

def gram_matrix(tensor):
  b,c,h,w = tensor.size()
  tensor = tensor.view(c,h*w)
  gram = torch.mm(tensor,tensor.t())
  return gram 

style_grams = {layer: gram_matrix(style_f[layer]) for layer in style_f}

# Creating Style and Content loss fucntion

def content_loss(target_conv4_2, content_conv4_2):
  loss = torch.mean((target_conv4_2-content_conv4_2)**2)
  return loss

style_weights = {
    'conv1_1' : 1.0,
    'conv2_1' : 0.75,
    'conv3_1' : 0.2,
    'conv4_1' : 0.2,
    'conv5_1' : 0.2

}

def style_loss(style_weights,target_features,style_grams):
  loss = 0
  for layer in style_weights:
    target_f = target_features[layer]
    target_gram = gram_matrix(target_f)
    style_gram = style_gram[layer]
    b,c,h,w = target_f.shape 
    layer_loss = style_weights[layers]*torch.mean((target_gram-style_gram)**2)
    loss += layer_loss/(c*h*w)
    return loss

target = content_p.clone().requires_grad_(True).to(device)
target_f = get_features(target,vgg)


print("Content Loss : ",content_loss(target_f['conv4_2'],content_f['conv4_2']))
print("Style Loss : ",style_loss(style_weights, target_f , style_grams))

最后2行代码的输出：

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-22-3b32a5406c6a> in <module>()
----> 1 print("Content Loss : ",content_loss(target_f['conv4_2'],content_f['conv4_2']))
      2 print("Style Loss : ",style_loss(style_weights, target_f , style_grams))

KeyError: 'conv4_2'

我将非常感谢您的快速回复！ 如果有人需要任何与代码相关的帮助来解决这个问题，请告诉我！

Answer 1

假设您使用的是标准的torchvision VGG 网络，类似于：

>>> vgg = torchvision.models.vgg16()

您希望在 get_features 函数中访问一些卷积层。但是，这些层包含在子模块中，而不是直接子模块（您使用 _modules 属性访问的子模块）。 顺便说一句，你不应该使用这个属性，下划线是为了指定 private 属性。

>>> dict(vgg.named_children()).keys()
dict_keys(['features', 'avgpool', 'classifier'])

所以在这里，您应该查看vgg.features，即 VGG 的 CNN 部分。

那么此时您的函数确实有意义，因为您是从模型卷积部分的层号索引：

def get_features(x, model):
   layers = {'0' : 'conv1_1',
             '5' : 'conv2_1',
             '10': 'conv3_1',
             '19': 'conv4_1',
             '21': 'conv4_2',
             '28': 'conv5_1'}

   features = {}
   for name, layer in model.named_children():
      x = layer(x)
      if name in layers:
         features[layers[name]] = x.detach()

   return features

那么features 将包含以下键：

>>> features.keys()
dict_keys(['conv1_1', 'conv2_1', 'conv3_1', 'conv4_1', 'conv4_2', 'conv5_1'])

---

或者，您可以使用forward hooks，我认为这是一个更好的解决方案。在this answer中我非常详细地解释了如何实现这样的方法。