Cannot use GCNConv with pytorch.cuda.amp

[dxyzx0]

Cannot use pytorch.cuda.amp since some tensor in GCNConv doesn't change dtype to half float
PS: sgcn_conv.py is just a rewrite of gcn_conv.py with only code modification of gcn_norm, I think it won't affect adj_t

I find that SparseTensor will be autocast by torch.cuda.amp, that's where the problem resides.

[rusty1s]

Does the same code work for you when using the edge_index representation instead of SparseTensor? And do you have a small script to reproduce? torch-sparse supports half-precision by now, and conversion of adj_t to torch.half should work as well.

[rusty1s]

I can look into it. If you have a minimal example to share, please do so. That would help me out tremendously.

[dxyzx0]

import torch
import torch_geometric
from torch_geometric.data import Data
from torch_sparse import SparseTensor


class Net(torch.nn.Module):
    def __init__(self, device):
        super(Net, self).__init__()
        self.device = device
        self.gcn = torch_geometric.nn.GCNConv(2, 3)

    def forward(self, data: SparseTensor):
        X, Adj, edge_weight = data.x, data.edge_index, data.edge_attr
        return self.gcn(X, Adj, edge_weight)


if __name__ == "__main__":
    edge_index = torch.LongTensor([[0, 1], [1, 0]])
    edge_weight = torch.Tensor([0.1, 0.2])
    x = torch.Tensor([[0, 1], [1, 1]])
    y = torch.Tensor([0, 1])
    data = Data(x=x, edge_index=edge_index, edge_weight=edge_weight, y=y)

    model = Net(device="cuda")
    model.to(model.device)
    with torch.cuda.amp.autocast():
        data.edge_index = SparseTensor(row=data.edge_index[0], col=data.edge_index[1], value=data.edge_attr,
                                       sparse_sizes=(data.x.size(0), data.x.size(0)))
        data.edge_attr = None
        data = data.to(model.device)

        out = model(data)

And what I get:

[dxyzx0]

My env:

[rusty1s]

Thanks for reporting. I fixed it in the latest torch-sparse commit, see here.

[rusty1s]

It seems that x.pt already contains inf values. One additional think that might cause this is the eps value of LayerNorm, which you may want to set to a higher value in a half-precision scenario.