如何在PyTorch中实现节点非全连接的自定义神经网络结构？

你要的非全连接结构本质是对线性变换的权重做固定位置的屏蔽，核心用掩码矩阵实现即可，这个方案可以直接嵌入到现有RNN、LSTM的结构中替换原有的全连接层：

• 首先根据连接规则生成固定的掩码矩阵：输入层3个节点（a/b/c）、输出层3个节点（d/e/f），需要连通的位置设为1，不需要连通的位置设为0，对应规则的掩码如下（行对应输出节点，列对应输入节点）：

  [[1, 0, 0], # d仅接收a的输入
   [1, 1, 0], # e接收a、b的输入
   [1, 1, 1]] # f接收所有输入
  

• 自定义线性层时，将可学习的权重参数和掩码矩阵做元素乘法，即可屏蔽掉不需要的连接，反向传播时被屏蔽位置的权重梯度会为0，不会被更新。
• 掩码矩阵需要注册为模型的buffer，不会被优化器更新，同时会自动和模型同步设备（CPU/GPU）。

自定义非全连接线性层 #

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

class MaskedLinear(nn.Module):
    def __init__(self, in_features=3, out_features=3, bias=True):
        super().__init__()
        # 定义可学习的权重和偏置
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.randn(out_features)) if bias else None
        # 注册固定掩码为buffer，不参与训练
        mask = torch.tensor([
            [1, 0, 0],
            [1, 1, 0],
            [1, 1, 1]
        ], dtype=torch.float32)
        self.register_buffer('mask', mask)
    
    def forward(self, x):
        # 前向传播时先对权重做掩码，屏蔽不需要的连接
        masked_weight = self.weight * self.mask
        return F.linear(x, masked_weight, self.bias)

嵌入到RNN单元使用 #

如果要替换内置RNN的全连接层，自定义RNN单元即可：

class CustomRNNCell(nn.Module):
    def __init__(self, input_size, hidden_size=3):
        super().__init__()
        self.hidden_size = hidden_size
        # 输入到隐状态的变换可按需选择是否用掩码层
        self.input2hidden = nn.Linear(input_size, hidden_size)
        # 隐状态到隐状态的变换使用自定义非全连接层
        self.hidden2hidden = MaskedLinear(hidden_size, hidden_size)
    
    def forward(self, x, hidden_prev):
        hidden_next = torch.tanh(self.input2hidden(x) + self.hidden2hidden(hidden_prev))
        return hidden_next

可以通过单节点激活测试验证连接规则是否正确：

# 构造测试输入：分别单独激活a、b、c三个节点
test_x = torch.tensor([
    [1, 0, 0], # 仅a激活
    [0, 1, 0], # 仅b激活
    [0, 0, 1]  # 仅c激活
], dtype=torch.float32)

layer = MaskedLinear()
output = layer(test_x)
print(output)
# 输出符合预期：
# 第一行（仅a激活）：d、e、f三个节点均有非零值
# 第二行（仅b激活）：d节点为0，e、f有非零值
# 第三行（仅c激活）：d、e节点为0，仅f有非零值

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