python - 具有自定义层的 PyTorch 网络在 CPU 上运行良好，但在移动到 GPU 时会出现 cudaErrorIllegalAddress

Question

我正在尝试实现我自己的 Graph-Attention-network 版本

自定义的GAT层如下

class GATLayer(nn.Module):
    def __init__(self, input_dim: int, output_dim: int, adj: torch.tensor):
        super().__init__()
        self.W = nn.Parameter(torch.zeros(size=(output_dim, input_dim)))
        self.a = nn.Parameter(torch.zeros(size=(2 * output_dim,)))
        self.adj = adj
        self.n_points = adj.shape[0]
        #print(f"input dim:{input_dim}")

    def forward(self, h: torch.Tensor):
        B, T, N, F = h.size()
        hh = functional.linear(h, self.W)
        output = torch.zeros_like(hh)
        for i in range(self.n_points):
          #  print(i)
            hhj = hh[:, :, self.adj[i], :]
            hhi = torch.cat([hh[:, :, i:i + 1, :]] * hhj.size(2), 2)
            hhij = torch.cat([hhi, hhj], 3)
            e = torch.mm(hhij.reshape(B * T * hhj.size(2), -1), self.a.reshape(self.a.size(0), 1)).reshape(B, T, -1)
            alpha = functional.softmax(e, dim=2)
            output[:, :, i, :] = torch.sum(hhj * torch.cat([torch.unsqueeze(alpha, 3)] * hhj.size(3), 3), dim=2)
        return output

整个网络定义为:

class AQIP(nn.Module):
    def __init__(self, adj: torch.tensor, seq_len: int, with_aqi: bool = True):
        super().__init__()
        self.hid_size = 128
        self.seq_len = seq_len
        self.gat_layers = [
            GATLayer(input_dim=16 + int(with_aqi), output_dim=128, adj=adj),
            GATLayer(input_dim=128, output_dim=128, adj=adj),
        ]
        self.rnns = [
            nn.LSTM(input_size=128, hidden_size=128, num_layers=4, bias=True, batch_first=True),
        ]
        self.linear = nn.Linear(in_features=128 * 4, out_features=1, bias=True)

    def forward(self, x: torch.Tensor, site_idx: int):
        h = torch.zeros(size=(4, x.size(0), 128))
        c = torch.zeros(size=(4, x.size(0), 128))
        for gat in self.gat_layers:
            x = gat(x)
        for rnn in self.rnns:
            x[:, :, site_idx, :], (h, c) = rnn(x[:, :, site_idx, :], (h, c))
        h = h.permute(1, 0, 2)
        h = h.reshape(h.size(0), -1)
        return self.linear(h).squeeze()

当我用下面的代码独立测试自定义的 GAT 层时，事实证明 GAT 层即使在 GPU 上也能正常工作

model = GATLayer(3, 1024, torch.tensor(np.array([[1, 0, 1], [0, 0, 1], [1, 0, 1]], dtype='bool')))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
print(model(torch.randn(5, 5, 3, 3)).shape)

输出 torch.Size([5, 5, 3, 1024])

当我用 CPU 和下面的代码测试整个网络时，它也工作正常

#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
adj = torch.tensor(np.array([[1, 0, 0], [0, 1, 1], [1, 1, 1]], dtype="bool"))
exp = torch.randn(3, 8, 3, 17)
gpus = [0]
model = AQIP(adj, seq_len=8)
#model = model.to(device, non_blocking=True)
print(model(exp, 1))

输出 tensor([-0.0320, -0.0320, -0.0320], grad_fn=<SqueezeBackward0>)

但是只要我尝试将模型移动到 GPU 并取消对设备和线路的注释，我就会收到以下错误并回溯到与我的代码无关的某个 Formatter 类:

RuntimeError: copy_if failed to synchronize: cudaErrorIllegalAddress: an illegal memory access was encountered

使用 CUDA_LAUNCH_BLOCKING=1 时运行代码，我得到:

RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc)`

根本没有帮助我定位错误

我也在同一台机器上运行官方示例并尝试将它们移动到 GPU，结果它们都运行良好。所以我想这与 CUDA 和 Cudnn 或 GPU 驱动程序版本不兼容无关。但我也无法在我的代码中找到问题所在。请帮忙!如果您能帮我解决这个问题，我将不胜感激。

Answer 1

经过无数次的努力，终于找到了问题所在。事实证明，如果你将图层放入列表中，例如

self.gat_layers = [
           GATLayer(input_dim=16 + int(with_aqi), output_dim=128, adj=adj).cuda(),
           GATLayer(input_dim=128, output_dim=128, adj=adj).cuda(),
         ]

然后 PyTorch 不会自动识别这些层，以便在调用 .to(device) 时，它们的参数不会传输到 GPU。因此解决方案之一是逐层声明。

虽然更好的解决方案是使用nn.ModuleList 来包含您想要的所有层，因此代码可以更改为

self.gat_layers = nn.ModuleList([
            GATLayer(input_dim=16 + int(with_aqi), output_dim=128, adj=adj).cuda(),
            GATLayer(input_dim=128, output_dim=128, adj=adj).cuda(),
        ])