PyTorch中如何高效实现避免特定输入对指定神经网络输出产生影响

我明白你现在的困扰——用三个独立LSTM分别对应输出确实能满足输入输出的隔离需求，但计算效率太低，相当于把模型跑了三遍，太浪费资源了。其实我们可以通过共享核心计算逻辑或者并行批处理的方式，在满足隔离要求的同时大幅提升效率，下面给你两个实用的方案：

方案一：合并输入批，共享LSTM参数（最推荐，高效易实现） #

这个思路的核心是把三个不同的输入子集（对应每个输出允许的输入）打包成一个大批次，只跑一次LSTM，再拆分结果给到各自的全连接层。这样既共享了LSTM的参数，又利用GPU的并行计算能力把三次LSTM的计算合并成一次，速度提升非常明显。

代码示例：

import torch
import torch.nn as nn

class EfficientLstmModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        # 只使用一个共享的LSTM层，输入维度是input_size-1（因为每个子集都排除一个输入）
        self.lstm = nn.LSTM(input_size - 1, hidden_size, batch_first=True)
        # 为每个输出配一个独立的全连接层
        self.fcs = nn.ModuleList([nn.Linear(hidden_size, output_size) for _ in range(3)])

    def forward(self, x):
        # x的形状：(batch_size, seq_length, input_size)，这里input_size=5
        batch_size, seq_len, _ = x.shape

        # 准备对应每个输出的输入子集：
        # 输出1排除第2个输入（索引1）
        input_for_out1 = x[:, :, [0, 2, 3, 4]]
        # 输出2排除第1个输入（索引0）
        input_for_out2 = x[:, :, [1, 2, 3, 4]]
        # 输出3排除第4个输入（索引3）
        input_for_out3 = x[:, :, [0, 1, 2, 4]]

        # 把三个输入子集合并成一个大批次，形状变为(3*batch_size, seq_len, 4)
        combined_input = torch.cat([input_for_out1, input_for_out2, input_for_out3], dim=0)

        # 只跑一次LSTM
        _, (hn, _) = self.lstm(combined_input)
        # 取最后一层的隐藏状态，形状是(3*batch_size, hidden_size)
        hn = hn[-1]

        # 把结果拆回三个批次，对应三个输出
        hn1, hn2, hn3 = hn.split(batch_size, dim=0)

        # 分别通过全连接层得到最终输出
        output1 = self.fcs[0](hn1)
        output2 = self.fcs[1](hn2)
        output3 = self.fcs[2](hn3)

        return output1, output2, output3

这个方案的优势：

• 参数总量只有原来的1/3（只用一个LSTM），减少了训练时的内存占用；
• 计算效率接近单次LSTM的耗时，GPU可以充分并行处理合并后的批次，比三次单独跑快很多；
• 完全满足你的隔离需求：每个输出对应的输入子集都排除了指定输入，所以输出不会受被排除输入的影响。

方案二：拆分输入维度，掩码隐藏状态（隔离更彻底） #

如果你需要更严格的输入输出隔离（比如确保被排除的输入信息完全不会进入对应输出的计算路径），可以把每个输入维度单独交给一个子LSTM处理，然后在隐藏状态层直接掩码掉对应输入的部分，再给到全连接层。

代码示例：

import torch
import torch.nn as nn

class SplitInputLSTM(nn.Module):
    def __init__(self, input_dims, hidden_size, output_size):
        super().__init__()
        # 为每个输入维度创建一个子LSTM，每个子LSTM的隐藏维度是总hidden_size的1/input_dims
        self.sub_lstms = nn.ModuleList([nn.LSTM(1, hidden_size//input_dims, batch_first=True) for _ in range(input_dims)])
        self.fcs = nn.ModuleList([nn.Linear(hidden_size, output_size) for _ in range(3)])
        self.hidden_size = hidden_size
        self.input_dims = input_dims

        # 预定义掩码，固定不可训练
        self.mask_out1 = torch.ones(hidden_size)
        self.mask_out1[hidden_size//input_dims : 2*hidden_size//input_dims] = 0.0  # 排除第2个输入的隐藏部分
        self.mask_out2 = torch.ones(hidden_size)
        self.mask_out2[:hidden_size//input_dims] = 0.0  # 排除第1个输入的隐藏部分
        self.mask_out3 = torch.ones(hidden_size)
        self.mask_out3[3*hidden_size//input_dims : 4*hidden_size//input_dims] = 0.0  # 排除第4个输入的隐藏部分

    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        hidden_states = []

        # 每个输入维度单独过子LSTM
        for i in range(self.input_dims):
            input_i = x[:, :, i:i+1]
            _, (hn_i, _) = self.sub_lstms[i](input_i)
            hidden_states.append(hn_i[-1])  # 形状：(batch_size, hidden_size//input_dims)
        
        # 合并所有子隐藏状态，得到完整隐藏状态：(batch_size, hidden_size)
        full_hidden = torch.cat(hidden_states, dim=1)

        # 把掩码移到和输入相同的设备上
        self.mask_out1 = self.mask_out1.to(x.device)
        self.mask_out2 = self.mask_out2.to(x.device)
        self.mask_out3 = self.mask_out3.to(x.device)

        # 应用掩码，阻断指定输入的影响
        hn1 = full_hidden * self.mask_out1
        hn2 = full_hidden * self.mask_out2
        hn3 = full_hidden * self.mask_out3

        # 通过全连接层得到输出
        output1 = self.fcs[0](hn1)
        output2 = self.fcs[1](hn2)
        output3 = self.fcs[2](hn3)

        return output1, output2, output3

这个方案的优势是从根源上阻断了指定输入到对应输出的路径，完全不会有信息泄露；缺点是子LSTM的数量等于输入维度，参数总量比方案一多一些，但计算效率依然远高于你的原始方案。

备注：内容来源于stack exchange，提问作者bird