我需要创建一个神经网络，其中使用二进制门来将某些张量置零，这些张量是禁用电路的输出。
为了提高运行速度，我希望使用 torch.bool 二进制门来阻止网络中禁用电路的反向传播。然而，我使用官方 PyTorch 的 CIFAR-10 数据集示例进行了小型实验，无论 gate_A 和 gate_B 的值如何，运行速度都完全相同：（这意味着这个想法行不通）
class Net(nn.Module): def __init__(self): super().__init__() self.pool = nn.MaxPool2d(2, 2) self.conv1a = nn.Conv2d(3, 6, 5) self.conv2a = nn.Conv2d(6, 16, 5) self.conv1b = nn.Conv2d(3, 6, 5) self.conv2b = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(32 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): # 理论上应该随机启用一个门 # 然而，为了实验，我将值固定为 [1,0] 和 [1,1] choice = randint(0,1) gate_A = torch.tensor(choice ,dtype = torch.bool) gate_B = torch.tensor(1-choice ,dtype = torch.bool) a = self.pool(F.relu(self.conv1a(x))) a = self.pool(F.relu(self.conv2a(a))) b = self.pool(F.relu(self.conv1b(x))) b = self.pool(F.relu(self.conv2b(b))) a *= gate_A b *= gate_B x = torch.cat( [a,b], dim = 1 ) x = torch.flatten(x, 1) # 展平除批次外的所有维度 x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x如何定义 gate_A 和 gate_B，以便在它们为零时有效地停止反向传播？
附注：在运行时动态更改 concatenation 也会改变分配给每个模块的权重。（例如，关联到 a 的权重在另一轮中可能被分配给 b，这会扰乱网络的运行方式）。
回答：
您可以使用 torch.no_grad（下面的代码可能会更简洁）：
def forward(self, x): # 理论上应该随机启用一个门 # 然而，为了实验，我将值固定为 [1,0] 和 [1,1] choice = randint(0,1) gate_A = torch.tensor(choice ,dtype = torch.bool) gate_B = torch.tensor(1-choice ,dtype = torch.bool) if choice: a = self.pool(F.relu(self.conv1a(x))) a = self.pool(F.relu(self.conv2a(a))) a *= gate_A with torch.no_grad(): # 禁用梯度计算 b = self.pool(F.relu(self.conv1b(x))) b = self.pool(F.relu(self.conv2b(b))) b *= gate_B else: with torch.no_grad(): # 禁用梯度计算 a = self.pool(F.relu(self.conv1a(x))) a = self.pool(F.relu(self.conv2a(a))) a *= gate_A b = self.pool(F.relu(self.conv1b(x))) b = self.pool(F.relu(self.conv2b(b))) b *= gate_B x = torch.cat( [a,b], dim = 1 ) x = torch.flatten(x, 1) # 展平除批次外的所有维度 x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x再看一遍，我认为以下是对特定问题更简单的解决方案：
def forward(self, x): # 理论上应该随机启用一个门 # 然而，为了实验，我将值固定为 [1,0] 和 [1,1] choice = randint(0,1) if choice: a = self.pool(F.relu(self.conv1a(x))) a = self.pool(F.relu(self.conv2a(a))) b = torch.zeros(shape_of_conv_output) # 此处替换卷积输出的形状 else: b = self.pool(F.relu(self.conv1b(x))) b = self.pool(F.relu(self.conv2b(b))) a = torch.zeros(shape_of_conv_output) # 此处替换卷积输出的形状 x = torch.cat( [a,b], dim = 1 ) x = torch.flatten(x, 1) # 展平除批次外的所有维度 x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x