PyTorch中如何获取深度神经网络各隐藏层神经元状态对应的张量

神经元的激活值属于模型前向传播的中间输出，不属于模型的可训练参数，因此无法直接用named_parameters()遍历获取，你可以用以下两种通用方案实现：

方案1：使用PyTorch前向钩子（无需修改模型结构，通用度最高） #

前向钩子可以在不修改模型forward逻辑的前提下，捕获每一层的输出结果，示例代码如下：

import torch

# 初始化存储各层激活值的字典
self.layer_activations = {}

# 定义钩子生成函数
def get_activation_hook(layer_name):
    def hook(module, input, output):
        # detach是为了把张量从计算图剥离，避免额外显存占用
        self.layer_activations[layer_name] = output.detach()
    return hook

# 遍历模型所有层，给目标隐藏层注册钩子
for name, layer in self.model.named_modules():
    # 可根据实际需求过滤层类型，比如只保留线性层、卷积层这类计算层，跳过激活、池化、Dropout层
    if isinstance(layer, (torch.nn.Linear, torch.nn.Conv2d, torch.nn.Conv3d)):
        layer.register_forward_hook(get_activation_hook(name))

每次模型完成前向传播后，就可以直接遍历存储的激活值写入TensorBoard：

for layer_name, activation in self.layer_activations.items():
    self.writer.add_histogram(f'Neuron activation/{layer_name}', activation, global_step=t)

方案2：修改模型forward直接返回激活值（自定义模型时更简单） #

如果是你自己开发的模型，直接在forward逻辑里收集各隐藏层的输出一起返回即可，不需要额外注册钩子，示例如下：

import torch.nn as nn
import torch.nn.functional as F

class CustomModel(nn.Module):
    def __init__(self, input_dim=10, num_classes=2):
        super().__init__()
        self.hidden1 = nn.Linear(input_dim, 64)
        self.hidden2 = nn.Linear(64, 32)
        self.output_layer = nn.Linear(32, num_classes)
    
    def forward(self, x):
        act1 = F.relu(self.hidden1(x))
        act2 = F.relu(self.hidden2(act1))
        out = self.output_layer(act2)
        # 同时返回最终预测结果和各层激活值字典
        return out, {"hidden_layer_1": act1, "hidden_layer_2": act2}

训练阶段拿到返回的激活值字典后，直接写入TensorBoard即可。

注意事项 #

• 如果只需要查看特定层的激活值，注册钩子或者收集返回值时可以直接过滤层名称，不需要全量收集减少性能消耗
• 激活值不需要参与反向传播，一定要加.detach()，否则会保留计算图占用大量显存

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