PyTorch模型如何为全连接层神经元添加字符串或数值型名称标识

该方案完全可实现，且可以做到不影响原有模型的前向计算、反向传播逻辑，以下是两种常见场景的实现方式：

场景1：仅调试打印时需要附带神经元ID（推荐，侵入性最小） #

这种场景下不需要修改原有模型的计算逻辑，只需要在打印前额外生成ID张量、拼接为调试用结构即可，原有计算流完全不受影响：

import torch
import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(1, 5)
        self.fc2 = nn.Linear(5, 10)
        self.fc3 = nn.Linear(10, 1)

    def forward(self, x):
        x = self.fc1(x)
        # === 新增调试代码开始 ===
        # 生成神经元数值ID，1对应第一个神经元，2对应第二个，以此类推
        neuron_ids = torch.arange(1, x.size(-1)+1, dtype=torch.float32, device=x.device)
        # 沿最后一维拼接值和ID，仅用于打印查看
        debug_tensor = torch.stack([x, neuron_ids.expand_as(x)], dim=-1)
        print("fc1层输出（值+神经元ID）：\n", debug_tensor)
        # 提取数值部分直接用原有变量即可，后续计算完全不受影响
        # 如果需要单独取数值/ID：value = debug_tensor[..., 0], id = debug_tensor[..., 1].int()
        # === 新增调试代码结束 ===
        x = torch.relu(x)        
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Model()
opt = optim.Adam(net.parameters())
features = torch.rand((3,1)) 
net(features)

输出示例：

fc1层输出（值+神经元ID）：
 tensor([[[ 0.6703, 1.0000],
         [ 0.4484, 2.0000],
         [-0.8529, 3.0000],
         [ 1.3119, 4.0000],
         [ 0.6741, 5.0000]],

        [[ 0.9112, 1.0000],
         [ 0.6496, 2.0000],
         [-1.2960, 3.0000],
         [ 1.8264, 4.0000],
         [ 0.4547, 5.0000]],

        [[ 0.7483, 1.0000],
         [ 0.5135, 2.0000],
         [-0.9963, 3.0000],
         [ 1.4785, 4.0000],
         [ 0.6031, 5.0000]]], grad_fn=<StackBackward0>)

场景2：全流程需要保留神经元ID信息 #

如果需要在前向传播过程中一直关联神经元ID，可以将值和ID分开存储为元组，使用时取第一个元素即可，避免维度错误。
你也可以封装自定义全连接层，复用ID生成逻辑，适合多层都需要加神经元标识的场景：

class LinearWithNeuronID(nn.Module):
    def __init__(self, in_features, out_features, **kwargs):
        super().__init__()
        self.core = nn.Linear(in_features, out_features, **kwargs)
        # 提前存储神经元ID，不需要每次生成
        self.register_buffer('neuron_ids', torch.arange(1, out_features+1))

    def forward(self, x, return_id=False):
        val = self.core(x)
        if return_id:
            return val, self.neuron_ids.to(val.device).expand_as(val)
        return val

# 修改模型初始化替换原有Linear层即可
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = LinearWithNeuronID(1, 5)
        self.fc2 = LinearWithNeuronID(5, 10)
        self.fc3 = LinearWithNeuronID(10, 1)

    def forward(self, x):
        # 同时拿到输出值和对应神经元ID
        x_val, x_id = self.fc1(x, return_id=True)
        # 你需要的[值, ID]对应结构，分开存储不会影响计算
        neuron_tensor = (x_val, x_id)
        print("神经元值：\n", neuron_tensor[0])
        print("神经元ID：\n", neuron_tensor[1])
        # 后续计算直接取值部分即可
        x = torch.relu(x_val)
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

注意：不要将ID作为新维度合并到参与计算的张量中，否则会导致后续全连接层输入维度不匹配，出现运行错误。

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