如何为PyTorch预训练VGG16添加权重归一化（Weight Normalization）

你当前的实现无法正确生效，存在两个核心问题：

• 问题1：子模块替换逻辑错误
  你在for k,m in self.encoder.encoder._modules.items()循环中，直接将wn(m)赋值给临时变量m，这一操作不会修改self.encoder.encoder中实际存储的子模块，权重归一化层并没有真正插入到VGG16的特征提取网络中，相当于你的添加操作完全不生效。
• 问题2：类名语义不匹配（可选优化）
  你定义的ResnetEncoder类内部实际加载的是VGG16模型，命名和实际功能不符，会降低代码可读性，建议调整类名为VGGEncoder。

正确实现方案 #

修正核心逻辑：将应用权重归一化后的层重新赋值回模型的子模块字典中，代码如下：

import torch.nn as nn
from torch.nn.utils import weight_norm as wn
import torchvision.models as models

# 优化类名，和实际使用的模型匹配
class VGGEncoder(nn.Module):
    def __init__(self):
        super(VGGEncoder, self).__init__()
        # 此处省略你原有的mean、std、num_ch_enc等参数初始化逻辑
        self.encoder = models.vgg16(pretrained=True).features
        # 直接在Encoder初始化时添加权重归一化，逻辑更内聚
        for k, m in self.encoder._modules.items():
            if isinstance(m, nn.Conv2d):
                self.encoder._modules[k] = wn(m)

    def forward(self, input_image):
        self.features = []
        x = (input_image - self.mean) / self.std
        
        self.features.append(self.encoder(x))
        # 省略你原有的多尺度特征保存逻辑

        return self.features

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.encoder = VGGEncoder()
        self.decoder = DepthDecoder(self.encoder.num_ch_enc)
        # 如果你不想在Encoder内部加，也可以在这里实现替换，逻辑如下：
        # for k,m in self.encoder.encoder._modules.items():
        #     if isinstance(m,nn.Conv2d):
        #         self.encoder.encoder._modules[k] = wn(m)

    def forward(self,x):
        return self.decoder(self.encoder(x))

vgg_backbone_model = Net()
vgg_backbone_model.train()

额外注意事项 #

• 预训练权重适配：PyTorch的weight_norm会自动从原有Conv2d层的预训练权重中计算出weight_g和weight_v两个参数，不需要额外做初始化适配，原有预训练特征不会被破坏。
• 训练阶段注意：权重归一化在训练阶段会对权重做归一化计算，推理阶段可以调用nn.utils.remove_weight_norm移除归一化层加快推理速度，同时保证输出结果和训练阶段一致。

内容的提问来源于stack exchange，提问作者Mohit Lamba