如何存储PyTorch MLP隐藏层变量用于可视化并转换为NumPy格式

你之前保存失败的核心原因是PyTorch张量默认绑定计算梯度的计算图、或是存放在GPU显存中时，无法直接转换为NumPy数组，按照以下步骤修改即可：

步骤1：修改MLP的forward方法，返回需要的中间张量 #

import torch
import torch.nn as nn
import numpy as np

class MLP(nn.Module):
  def __init__(self):
    super().__init__()
    self.in_dim = 28 * 28
    self.out_dim = 10
    self.fc1 = nn.Linear(self.in_dim,512)
    self.fc2 = nn.Linear(512, 256)
    self.fc3 = nn.Linear(256, 128)
    self.fc4 = nn.Linear(128, 64)
    self.fc5 = nn.Linear(64, self.out_dim)
    self.relu = nn.ReLU()

  def forward(self, x):
      # 新增变量保存fc1的原始输出
      fc1_raw = self.fc1(x.view(-1,self.in_dim))
      a1 = self.relu(fc1_raw)    

      a2 = self.relu(self.fc2(a1))
      a3 = self.relu(self.fc3(a2))
      a4 = self.relu(self.fc4(a3))
      logit = self.fc5(a4)
      # 返回需要的中间结果
      return logit, fc1_raw, a1

步骤2：前向传播后转换张量并保存 #

推理阶段保存（无梯度计算场景） #

推理阶段直接关闭梯度计算即可避免计算图相关报错，写法如下：

# 实例化模型、加载权重、准备输入x的常规步骤省略
model = MLP()
model.eval()

with torch.no_grad():
    logit, fc1_raw, a1 = model(x)
    # 先将张量移到CPU，再转NumPy，兼容CPU/GPU运行的所有场景
    fc1_raw_np = fc1_raw.cpu().numpy()
    a1_np = a1.cpu().numpy()

# 保存为npy文件
np.save("fc1_raw_output.npy", fc1_raw_np)
np.save("a1_relu_output.npy", a1_np)

训练阶段保存（需保留梯度计算场景） #

如果是训练过程中需要同步保存中间结果，不需要关闭梯度，转换时添加detach操作剥离计算图即可：

# 训练过程中前向传播代码
logit, fc1_raw, a1 = model(x)
# 损失计算、反向传播、优化器迭代等常规训练步骤省略

# 转换保存代码
fc1_raw_np = fc1_raw.detach().cpu().numpy()
a1_np = a1.detach().cpu().numpy()
np.save("fc1_raw_train.npy", fc1_raw_np)
np.save("a1_relu_train.npy", a1_np)

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