PyTorch GCNConv层前向传播输出不一致问题求助

我每次执行以下PyTorch代码时，输出结果都会不同。由于这是全连接图且各节点输入特征一致，同一节点的输出每次是相同的，但我期望每次执行脚本后整体输出完全一致。

我需要保证输入矩阵相同时，前向传播输出固定，以此测试我用其他语言实现的GCN层。正如代码所示，我已手动设置权重矩阵和偏置向量，也尝试过设置随机种子，且未使用Dropout，但问题仍存在，请问可能的原因是什么？

import torch
from torch_geometric.nn import GCNConv

x = torch.tensor([[1.0], [1.0], [1.0]], dtype=torch.float)
edge_index = torch.tensor([[0,0,1,1,2,2], [1,2,0,2,0,1]], dtype=torch.long)

conv_layer = GCNConv(in_channels=1, out_channels=1)

new_weight_values = torch.tensor([[1.0]])
new_bias_values = torch.tensor([[0.0]])

conv_layer.weight = torch.nn.Parameter(new_weight_values)
conv_layer.bias = torch.nn.Parameter(new_bias_values)

output = conv_layer.forward(x, edge_index)

print("Input Features:")
print(x)
print("Output Features:")
print(output)

原因分析与解决方案 #

输出不一致的核心原因是浮点数计算的非确定性，结合PyTorch和PyTorch Geometric的底层实现细节，具体原因和解决办法如下：

1. CPU多线程并行导致求和顺序变化 #

PyTorch默认启用CPU多线程加速，GCNConv内部使用的scatter_add聚合操作在多线程下的元素求和顺序不固定。由于浮点数加法不满足结合律，不同求和顺序会产生细微精度差异（比如1.0变成0.99999994或1.0000001），看起来输出不一致。

解决办法：强制PyTorch使用单线程运行，固定求和顺序：

torch.set_num_threads(1)
torch.set_num_interop_threads(1)

2. 随机种子未完全固定 #

即使手动设置了模型参数，仍需确保所有可能引入随机性的组件都被固定，包括Python原生随机库、NumPy和PyTorch的全局种子，以及CUDA相关的确定性设置（如果使用GPU）。

解决办法：在脚本最开头添加以下代码：

import random
import numpy as np

# 固定Python全局随机种子
random.seed(42)
# 固定NumPy随机种子
np.random.seed(42)
# 固定PyTorch CPU随机种子
torch.manual_seed(42)
# 如果使用GPU，固定CUDA种子与确定性设置
if torch.cuda.is_available():
    torch.cuda.manual_seed(42)
    torch.cuda.manual_seed_all(42)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

3. 验证设备一致性 #

确保输入张量、边索引和GCN层在同一设备上（均为CPU或均为GPU），跨设备的隐式数据传输可能引入意外差异。可以通过.device属性检查：

print(x.device, edge_index.device, conv_layer.weight.device)

如果不一致，使用.to()方法统一设备：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = x.to(device)
edge_index = edge_index.to(device)
conv_layer = conv_layer.to(device)

最终固定输出的完整代码 #

整合上述修改后的完整代码如下：

import random
import numpy as np
import torch
from torch_geometric.nn import GCNConv

# 固定所有随机种子与计算确定性
random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
torch.set_num_threads(1)
torch.set_num_interop_threads(1)
if torch.cuda.is_available():
    torch.cuda.manual_seed(42)
    torch.cuda.manual_seed_all(42)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

x = torch.tensor([[1.0], [1.0], [1.0]], dtype=torch.float)
edge_index = torch.tensor([[0,0,1,1,2,2], [1,2,0,2,0,1]], dtype=torch.long)

conv_layer = GCNConv(in_channels=1, out_channels=1)

new_weight_values = torch.tensor([[1.0]])
new_bias_values = torch.tensor([[0.0]])

conv_layer.weight = torch.nn.Parameter(new_weight_values)
conv_layer.bias = torch.nn.Parameter(new_bias_values)

output = conv_layer.forward(x, edge_index)

print("Input Features:")
print(x)
print("Output Features:")
print(output)

修改后每次执行脚本的输出都会完全一致，可满足跨语言GCN层的测试需求。

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