企业官网建设流程全解析

【Bug已解决】PyTorch: Extract learned weights correctly 解决方案

问题描述

在 PyTorch 深度学习开发中，训练完成后的模型权重提取是一个非常常见的需求。无论是用于模型部署、权重可视化、模型压缩、知识蒸馏还是简单的调试，我们都需要正确地从模型中提取学习到的权重参数。然而，许多开发者在使用state_dict()、named_parameters()、parameters()等方法时，经常会遇到各种问题，例如：

• 提取的权重张量与预期形状不一致
• 提取后修改权重导致原始模型参数也被修改（浅拷贝问题）
• 混淆了state_dict和parameters的使用场景
• 在多 GPU 训练（DataParallel/DistributedDataParallel）后提取权重时出现module.前缀问题
• 提取 BatchNorm 的running_mean、running_var等非参数缓冲区时遗漏

本文将全面讲解 PyTorch 中正确提取模型权重的方法，涵盖从基础到高级的各种场景，并提供完整的可运行代码示例。

错误复现

以下是一些常见的错误用法及其导致的问题：

错误 1：浅拷贝导致权重被意外修改

import torch import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(10, 20) self.fc2 = nn.Linear(20, 5) def forward(self, x): x = torch.relu(self.fc1(x)) return self.fc2(x) model = SimpleNet() # 错误：直接引用参数对象，修改会影响原模型 weights = dict(model.named_parameters()) weights['fc1.weight'] *= 0 # 这会直接修改模型权重！ print("修改后模型权重:", model.fc1.weight.data) # 全是0，模型被破坏了

错误 2：遗漏非参数缓冲区

# 错误：named_parameters() 不包含 BatchNorm 的 running_mean/running_var class BNNet(nn.Module): def __init__(self): super().__init__() self.bn = nn.BatchNorm1d(10) self.fc = nn.Linear(10, 5) model = BNNet() x = torch.randn(32, 10) model(x) # 前向传播以更新 BN 统计量 # 只用 named_parameters() 会遗漏 running_mean, running_var, num_batches_tracked params = dict(model.named_parameters()) print("named_parameters 包含的键:", list(params.keys())) # 输出: ['bn.weight', 'bn.bias', 'fc.weight', 'fc.bias'] # 缺少: bn.running_mean, bn.running_var, bn.num_batches_tracked

错误 3：DataParallel 模型的module.前缀

# 在多 GPU 环境下 model = SimpleNet() model = nn.DataParallel(model) # 此时 state_dict 的键都带有 'module.' 前缀 sd = model.state_dict() print(list(sd.keys())) # ['module.fc1.weight', 'module.fc1.bias', 'module.fc2.weight', 'module.fc2.bias'] # 直接加载到非 DataParallel 模型会报错

根因分析

1. PyTorch 参数与缓冲区的区别

PyTorch 的nn.Module中有三种类型的张量：

BatchNorm层中的running_mean、running_var和num_batches_tracked都是缓冲区，不是参数。它们在前向传播中被更新但不参与反向传播。因此，named_parameters()不会包含它们，而state_dict()会包含。

2. 浅拷贝 vs 深拷贝

dict(model.named_parameters())返回的字典中，值是对原始参数张量的引用。当你修改这个引用指向的张量时，原始模型的参数也会被修改。这是因为 PyTorch 张量是可变对象。

3. DataParallel 的模块嵌套

nn.DataParallel会将你的模型包装在一个Module容器中，原始模型成为其module属性。因此所有参数名都会自动加上module.前缀。DistributedDataParallel也有类似行为。

解决方案

方案一：使用state_dict()提取完整模型状态（推荐）

state_dict()是提取模型权权的标准方法，它包含所有参数和缓冲区：

import torch import torch.nn as nn import copy class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(10, 20) self.bn = nn.BatchNorm1d(20) self.fc2 = nn.Linear(20, 5) def forward(self, x): x = torch.relu(self.fc1(x)) x = self.bn(x) return self.fc2(x) model = SimpleNet() model.train() x = torch.randn(32, 10) output = model(x) # 前向传播更新 BN 统计量 # 提取完整 state_dict（包含参数和缓冲区） state_dict = model.state_dict() print("state_dict 包含的键:") for key, value in state_dict.items(): print(f" {key}: {value.shape}")

输出：

state_dict 包含的键: fc1.weight: torch.Size([20, 10]) fc1.bias: torch.Size([20]) bn.weight: torch.Size([20]) bn.bias: torch.Size([20]) bn.running_mean: torch.Size([20]) bn.running_var: torch.Size([20]) bn.num_batches_tracked: torch.Size([]) fc2.weight: torch.Size([5, 20]) fc2.bias: torch.Size([5])

方案二：使用named_parameters()提取可学习参数

当你只需要可学习参数（不需要 BN 的统计量）时：

# 提取可学习参数的独立副本 learnable_params = {} for name, param in model.named_parameters(): learnable_params[name] = param.data.clone() # .clone() 确保是独立副本 # 或者使用 detach() 创建不跟踪梯度的副本 learnable_params_v2 = {} for name, param in model.named_parameters(): learnable_params_v2[name] = param.detach().clone() print("可学习参数:") for name, param in learnable_params.items(): print(f" {name}: {param.shape}")

方案三：提取特定层的权重

# 方法 1：通过层属性直接访问 fc1_weight = model.fc1.weight.data # torch.Tensor, shape: [20, 10] fc1_bias = model.fc1.bias.data # torch.Tensor, shape: [20] # 方法 2：通过 named_parameters 过滤 for name, param in model.named_parameters(): if 'fc1' in name: print(f"{name}: {param.data.shape}") # 方法 3：通过 children() 遍历子模块 for layer in model.children(): if isinstance(layer, nn.Linear): print(f"Linear层权重: {layer.weight.data.shape}")

方案四：处理 DataParallel 模型

# 方法 1：通过 .module 属性访问原始模型 if isinstance(model, nn.DataParallel): original_model = model.module state_dict = original_model.state_dict() else: state_dict = model.state_dict() # 方法 2：去除 state_dict 中的 'module.' 前缀 state_dict = model.state_dict() clean_state_dict = {} for key, value in state_dict.items(): new_key = key.replace('module.', '') if key.startswith('module.') else key clean_state_dict[new_key] = value # 方法 3：DistributedDataParallel if isinstance(model, torch.nn.parallel.DistributedDataParallel): state_dict = model.module.state_dict()

完整修复代码

以下是一个完整的示例，展示如何正确提取、保存和加载模型权重：

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset import copy # ==================== 定义模型 ==================== class ClassificationNet(nn.Module): def __init__(self, input_dim=784, hidden_dim=256, num_classes=10): super().__init__() self.fc1 = nn.Linear(input_dim, hidden_dim) self.bn1 = nn.BatchNorm1d(hidden_dim) self.fc2 = nn.Linear(hidden_dim, hidden_dim) self.bn2 = nn.BatchNorm1d(hidden_dim) self.fc3 = nn.Linear(hidden_dim, num_classes) self.dropout = nn.Dropout(0.3) def forward(self, x): x = torch.relu(self.fc1(x)) x = self.bn1(x) x = self.dropout(x) x = torch.relu(self.fc2(x)) x = self.bn2(x) x = self.dropout(x) return self.fc3(x) def train_model(model, dataloader, epochs=5): """训练模型""" criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) model.train() for epoch in range(epochs): total_loss = 0 for batch_x, batch_y in dataloader: optimizer.zero_grad() outputs = model(batch_x) loss = criterion(outputs, batch_y) loss.backward() optimizer.step() total_loss += loss.item() print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(dataloader):.4f}") def extract_weights_comprehensive(model): """ 全面提取模型权重的工具函数 返回包含参数、缓冲区和完整 state_dict 的字典 """ result = { 'state_dict': copy.deepcopy(model.state_dict()), 'named_parameters': {}, 'named_buffers': {}, 'layer_weights': {} } # 提取可学习参数 for name, param in model.named_parameters(): result['named_parameters'][name] = { 'data': param.data.clone(), 'requires_grad': param.requires_grad, 'dtype': param.dtype, 'shape': tuple(param.shape) } # 提取缓冲区（如 BN 的 running_mean/var） for name, buffer in model.named_buffers(): result['named_buffers'][name] = { 'data': buffer.data.clone(), 'dtype': buffer.dtype, 'shape': tuple(buffer.shape) } # 按层提取权重 for name, module in model.named_modules(): if name == '': continue if hasattr(module, 'weight') and module.weight is not None: result['layer_weights'][name] = { 'weight': module.weight.data.clone(), 'bias': module.bias.data.clone() if module.bias is not None else None } return result def verify_weights(model, extracted): """验证提取的权重是否与模型一致""" current_sd = model.state_dict() extracted_sd = extracted['state_dict'] all_match = True for key in current_sd: if key in extracted_sd: if not torch.equal(current_sd[key], extracted_sd[key]): print(f" [不匹配] {key}") all_match = False else: print(f" [缺失] {key}") all_match = False if all_match: print(" 所有权重验证通过！") return all_match def load_weights_safely(model, state_dict, strict=True): """ 安全加载权重，处理键名不匹配的情况 """ # 去除 DataParallel 的 module. 前缀 model_keys = set(model.state_dict().keys()) sd_keys = set(state_dict.keys()) # 检查是否有 module. 前缀 has_module_prefix = any(k.startswith('module.') for k in sd_keys) if has_module_prefix: new_state_dict = {} for k, v in state_dict.items(): new_k = k[7:] if k.startswith('module.') else k new_state_dict[new_k] = v state_dict = new_state_dict # 加载权重 missing, unexpected = model.load_state_dict(state_dict, strict=strict) if missing: print(f"缺失的键: {missing}") if unexpected: print(f"多余的键: {unexpected}") return model # ==================== 主程序 ==================== def main(): # 创建模拟数据 torch.manual_seed(42) X = torch.randn(1000, 784) y = torch.randint(0, 10, (1000,)) dataset = TensorDataset(X, y) dataloader = DataLoader(dataset, batch_size=32, shuffle=True) # 创建并训练模型 model = ClassificationNet() print("=" * 60) print("开始训练模型...") train_model(model, dataloader, epochs=3) # 提取权重 print("\n" + "=" * 60) print("提取模型权重...") extracted = extract_weights_comprehensive(model) # 显示提取结果 print(f"\n可学习参数数量: {len(extracted['named_parameters'])}") for name, info in extracted['named_parameters'].items(): print(f" {name}: shape={info['shape']}, dtype={info['dtype']}, " f"requires_grad={info['requires_grad']}") print(f"\n缓冲区数量: {len(extracted['named_buffers'])}") for name, info in extracted['named_buffers'].items(): print(f" {name}: shape={info['shape']}, dtype={info['dtype']}") # 验证权重 print("\n" + "=" * 60) print("验证提取的权重...") verify_weights(model, extracted) # 保存和加载 print("\n" + "=" * 60) print("保存和加载模型...") # 方法 1：保存完整 state_dict torch.save(model.state_dict(), 'model_weights.pth') print("已保存 model_weights.pth") # 加载到新模型 new_model = ClassificationNet() loaded_sd = torch.load('model_weights.pth', weights_only=True) load_weights_safely(new_model, loaded_sd) # 验证两个模型输出一致 test_input = torch.randn(5, 784) model.eval() new_model.eval() with torch.no_grad(): out1 = model(test_input) out2 = new_model(test_input) print(f"输出差异: {(out1 - out2).abs().max().item():.2e}") print("模型加载验证通过！" if torch.allclose(out1, out2, atol=1e-6) else "模型加载有误！") # 方法 2：提取特定层权重用于分析 print("\n" + "=" * 60) print("提取特定层权重进行分析...") fc1_weights = model.fc1.weight.data print(f"FC1 权重统计:") print(f" 均值: {fc1_weights.mean().item():.6f}") print(f" 标准差: {fc1_weights.std().item():.6f}") print(f" 最大值: {fc1_weights.max().item():.6f}") print(f" 最小值: {fc1_weights.min().item():.6f}") print(f" L2范数: {fc1_weights.norm().item():.6f}") # 方法 3：提取权重用于可视化 print("\n提取权重用于可视化...") all_weights = [] for name, param in model.named_parameters(): if 'weight' in name: all_weights.append(param.data.flatten()) all_weights_cat = torch.cat(all_weights) print(f"所有权重合并后: {all_weights_cat.shape[0]} 个值") print(f" 全局均值: {all_weights_cat.mean().item():.6f}") print(f" 全局标准差: {all_weights_cat.std().item():.6f}") # 绘制权重分布直方图 import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt fig, axes = plt.subplots(1, 3, figsize=(15, 4)) for idx, layer_name in enumerate(['fc1', 'fc2', 'fc3']): layer = getattr(model, layer_name) weights = layer.weight.data.flatten().numpy() axes[idx].hist(weights, bins=50, alpha=0.7, color='steelblue') axes[idx].set_title(f'{layer_name} Weight Distribution') axes[idx].set_xlabel('Weight Value') axes[idx].set_ylabel('Frequency') plt.tight_layout() plt.savefig('weight_distributions.png', dpi=150) print("权重分布图已保存为 weight_distributions.png") # 方法 4：提取权重用于模型压缩（剪枝） print("\n" + "=" * 60) print("权重剪枝示例...") pruning_ratio = 0.5 for name, module in model.named_modules(): if isinstance(module, nn.Linear): weight = module.weight.data threshold = torch.quantile(torch.abs(weight).flatten(), pruning_ratio) mask = torch.abs(weight) > threshold print(f" {name}: 剪枝比例 {pruning_ratio*100:.0f}%, " f"阈值 {threshold.item():.6f}, " f"保留 {mask.sum().item()}/{mask.numel()} 个权重") print("\n所有权重提取操作完成！") if __name__ == '__main__': main()

常见陷阱与注意事项

1..datavs.detach()的区别

# .data：返回的张量与原张量共享存储，但不跟踪梯度历史 # 修改 .data 的内容不会触发 autograd 检查，可能导致梯度计算错误 weight_copy = model.fc1.weight.data.clone() # 安全：先 clone # .detach()：返回的张量与原张量共享存储，但从计算图中分离 # 更安全，推荐使用 weight_copy_safe = model.fc1.weight.detach().clone()

2.state_dict()返回的是引用还是副本

# state_dict() 返回的字典是新字典，但值是引用 sd = model.state_dict() sd['fc1.weight'][0, 0] = 999 # 这会修改模型权重！ # 正确做法：深拷贝 import copy sd_safe = copy.deepcopy(model.state_dict()) # 或者逐个 clone sd_safe2 = {k: v.clone() for k, v in model.state_dict().items()}

3. 保存整个模型 vs 只保存 state_dict

# 推荐：只保存 state_dict（与模型结构解耦） torch.save(model.state_dict(), 'model.pth') # 不推荐：保存整个模型（依赖 Python pickle，路径耦合） # torch.save(model, 'model_full.pth') # 加载时需要模型类定义在同一个文件路径

4. 多 GPU 模型的权重提取

# DataParallel model_dp = nn.DataParallel(model) # 提取时需要去除 module. 前缀 state_dict = {k.replace('module.', ''): v for k, v in model_dp.state_dict().items()} # DistributedDataParallel model_ddp = DDP(model, device_ids=[0]) # 使用 .module 获取原始模型 state_dict = model_ddp.module.state_dict()

5. 自定义层权重的提取

class CustomLayer(nn.Module): def __init__(self, in_features, out_features): super().__init__() # 自定义参数 self.weight = nn.Parameter(torch.randn(out_features, in_features)) self.bias = nn.Parameter(torch.zeros(out_features)) # 自定义缓冲区 self.register_buffer('running_stat', torch.zeros(out_features)) def forward(self, x): return x @ self.weight.T + self.bias layer = CustomLayer(10, 5) # state_dict 会包含 weight, bias 和 running_stat print(layer.state_dict().keys()) # odict_keys(['weight', 'bias', 'running_stat'])

总结

在 PyTorch 中正确提取模型权重的核心要点：

1. 使用state_dict()获取完整的模型状态（参数 + 缓冲区），这是最全面、最标准的方法。
2. 使用named_parameters()当只需要可学习参数时，但注意它不包含 BN 的running_mean/var。
3. 使用.clone()或copy.deepcopy()确保提取的权重是独立副本，避免意外修改原模型。
4. 处理多 GPU 前缀：DataParallel/DDP 模型提取时需要去除module.前缀或使用.module属性。
5. 区分参数和缓冲区：参数参与梯度计算，缓冲区不参与但仍是模型状态的一部分。
6. 保存时只保存state_dict，避免保存整个模型对象导致的路径耦合问题。

通过掌握这些方法，你可以灵活地在各种场景下正确提取和使用 PyTorch 模型的学习权重，无论是用于部署、分析还是调试。