【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错误 3DataParallel 模型的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中有三种类型的张量类型方法是否参与梯度计算是否出现在state_dict是否出现在parameters参数 (Parameter)nn.Parameter是是是缓冲区 (Buffer)register_buffer否是否普通属性self.x tensor否否否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(fLinear层权重: {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 # 方法 3DistributedDataParallel 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_dim784, hidden_dim256, num_classes10): 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, epochs5): 训练模型 criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.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(fEpoch {epoch1}/{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, strictTrue): 安全加载权重处理键名不匹配的情况 # 去除 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, strictstrict) 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_size32, shuffleTrue) # 创建并训练模型 model ClassificationNet() print( * 60) print(开始训练模型...) train_model(model, dataloader, epochs3) # 提取权重 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]}, frequires_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_onlyTrue) 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, atol1e-6) else 模型加载有误) # 方法 2提取特定层权重用于分析 print(\n * 60) print(提取特定层权重进行分析...) fc1_weights model.fc1.weight.data print(fFC1 权重统计:) 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, bins50, alpha0.7, colorsteelblue) 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, dpi150) 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 中正确提取模型权重的核心要点使用state_dict()获取完整的模型状态参数 缓冲区这是最全面、最标准的方法。使用named_parameters()当只需要可学习参数时但注意它不包含 BN 的running_mean/var。使用.clone()或copy.deepcopy()确保提取的权重是独立副本避免意外修改原模型。处理多 GPU 前缀DataParallel/DDP 模型提取时需要去除module.前缀或使用.module属性。区分参数和缓冲区参数参与梯度计算缓冲区不参与但仍是模型状态的一部分。保存时只保存state_dict避免保存整个模型对象导致的路径耦合问题。通过掌握这些方法你可以灵活地在各种场景下正确提取和使用 PyTorch 模型的学习权重无论是用于部署、分析还是调试。