如何将MNIST数据集DNN隐藏层输出转换为NumPy格式并保存

第一步：修改forward函数返回所需变量 #

首先调整MLP类的前向传播逻辑，把你需要存储的输入、第一层全连接输出、第一层激活值和logit一同返回，同时修正原示例代码中重复调用fc2的笔误：

import torch.nn as nn
import torch
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):
        # 计算需要保存的中间变量
        x_flat = x.view(-1, self.in_dim)
        z1 = self.fc1(x_flat) # 第一层全连接输出z[1]
        a1 = self.relu(z1) # 第一层激活值a[1]
        a2 = self.relu(self.fc2(a1))
        a3 = self.relu(self.fc3(a2))
        a4 = self.relu(self.fc4(a3))
        logit = self.fc5(a4)
        # 返回所有需要保存的变量
        return x, z1, a1, logit

第二步：张量转NumPy格式 #

推理阶段拿到返回的张量后，需要先做格式转换，注意如果张量在GPU上要先移到CPU，带梯度的张量要先断开计算图：

# 初始化模型、加载权重后先设置为评估模式
model = MLP()
model.eval()
# 推理阶段关闭梯度计算，可省去后续detach操作
with torch.no_grad():
    # input_tensor为输入的MNIST批次数据
    x, z1, a1, logit = model(input_tensor)

# 张量转NumPy数组
x_np = x.cpu().numpy()
z1_np = z1.cpu().numpy()
a1_np = a1.cpu().numpy()
logit_np = logit.cpu().numpy()

第三步：NumPy数组存储 #

可根据需求选择两种存储方式：

• 单数组存为独立npy文件，使用np.save

# 存储
np.save("x.npy", x_np)
np.save("z1.npy", z1_np)
np.save("a1.npy", a1_np)
np.save("logit.npy", logit_np)

# 读取示例
x_load = np.load("x.npy")

• 多数组打包存为单个npz文件，方便统一管理，使用np.savez

# 存储
np.savez(
    "mnist_pcl_result.npz",
    x=x_np,
    z1=z1_np,
    a1=a1_np,
    logit=logit_np
)

# 读取示例
data = np.load("mnist_pcl_result.npz")
x_load = data["x"]
z1_load = data["z1"]
a1_load = data["a1"]
logit_load = data["logit"]

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