PyTorch实现神经风格迁移计算损失时出现KeyError: 'conv4_2'如何解决

问题原因&修复方案 #

1 直接触发KeyError的核心问题 #

get_features 函数的return语句缩进错误，被放在了for循环内部：

for name, layer in model._modules.items():
    x = layer(x)
    if name in layers:
      Features[layers[name]] = x
    return Features # 这里缩进错误，第一次循环就直接返回

程序遍历第一层（name='0'）后就直接返回了结果，此时Features字典里只有conv1_1一个键，自然找不到后续的conv4_2。
修复方法：把return语句缩进到和for循环同级，等遍历完所有层后再返回：

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 # 修正缩进位置

2 其他隐藏的可运行错误 #

你的style_loss函数还有三处代码错误，修复完第一个问题后会触发：

• 变量名拼写错误：参数传入的是style_grams，但代码里写了style_gram[layer]
• 权重取值错误：style_weights[layers]应该改为style_weights[layer]，取当前遍历层的权重
• return缩进错误：同样放在了for循环内部，只会计算第一层风格损失就返回

修复后的style_loss代码：

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_grams[layer] # 修正变量名
    b,c,h,w = target_f.shape 
    layer_loss = style_weights[layer]*torch.mean((target_gram-style_gram)**2) # 修正权重取值
    loss += layer_loss/(c*h*w)
  return loss # 修正缩进位置

内容的提问来源于stack exchange，提问作者Rachit S Garg