如何在PyTorch中创建自定义连接/非全连接神经网络层

首先明确连接规则对应的约束：隐藏层每2个神经元为一组连接1个输出神经元，因此首先需要保证hidden_size为偶数，最终输出层神经元数量固定为hidden_size // 2。

你可以用两种方式实现这个自定义连接层：

方法1：自定义无冗余参数的线性层（推荐） #

直接按照连接规则实现前向逻辑，不存在多余的无效参数，计算效率最高，且兼容任意维度的输入格式：

import math
import torch
from torch import nn

class PairConnectLinear(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        if hidden_size % 2 != 0:
            raise ValueError("隐藏层神经元数量必须为偶数，才能按2个一组划分连接")
        self.out_features = hidden_size // 2
        # 每个输出神经元仅对应2个可学习权重，无冗余参数
        self.weight = nn.Parameter(torch.empty(self.out_features, 2))
        self.bias = nn.Parameter(torch.empty(self.out_features))
        # 采用和nn.Linear一致的默认初始化逻辑
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
        nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, x):
        # 保留输入除最后一维（隐藏层维度）外的所有维度形状
        pre_shape = x.shape[:-1]
        # 将隐藏层输出按每2个神经元拆分为独立分组
        x = x.reshape(*pre_shape, self.out_features, 2)
        # 逐组计算加权和+偏置得到输出
        return (x * self.weight).sum(dim=-1) + self.bias

# 直接实例化即可
layer2 = PairConnectLinear(hidden_size)

方法2：普通Linear+掩码（快速验证用） #

如果你不想自定义层，可以给普通全连接层加固定连接掩码，训练时始终把非连接位置的权重清零即可，好处是修改连接规则时只需要调整掩码：

from torch import nn

out_size = hidden_size // 2
layer2 = nn.Linear(hidden_size, out_size)

# 构造连接掩码：仅保留每组对应位置的权重连接
connect_mask = torch.zeros_like(layer2.weight)
for out_idx in range(out_size):
    connect_mask[out_idx, 2*out_idx] = 1.0
    connect_mask[out_idx, 2*out_idx + 1] = 1.0
# 将掩码注册为buffer，自动跟随模型设备切换，不参与梯度更新
layer2.register_buffer("connect_mask", connect_mask)
# 初始化后先清空非连接位置的权重
layer2.weight.data *= connect_mask

# 注册反向传播钩子，每次参数更新后自动清空非连接位置的权重
def mask_zero_hook(module, grad_input, grad_output):
    module.weight.data *= module.connect_mask
layer2.register_full_backward_hook(mask_zero_hook)

两种方式的前向计算结果完全等价，第一种训练速度更快、显存占用更低，第二种改连接规则更方便，按需选择即可。

内容的提问来源于stack exchange，提问作者Swain Subrat Kumar