torch_sparse.SparseTensor doesn't support torch_geometric.nn.batch_norm

[dxyzx0]

My code with dense tensor works well, but with both torch.sparse.Tensor and torch_sparse.SparseTensor, the torch_geometric.nn.Batch_norm won't work. How can I solve this problem?

[rusty1s]

Can you post a code snippet to reproduce? The BatchNorm does only take node features as argument. Are you trying to run BatchNorm with a sparse node feature matrix? This is not supported at the moment.

[dxyzx0]

Yes, my node feature matrix has many zeros.

[rusty1s]

I see. I suggest to convert your features into a dense matrix first as most GNNs expect dense features matrices (and all will produce dense output feature matrices).

[dxyzx0]

That's what I used to do. But I really need to increase the batch_size, so I want to make the node feature matrix or some other things to sparse tenor to reduce cuda memory consumption

[dxyzx0]

My node feature matrix is indeed a bipartite graph. The first n nodes(set A) and last m nodes(set B) are connected through the adjacency matrix, of which both own different features. Hence, my node feature matrix has many zeros, in the rows of A, features of B and rows of B, features of A. Is there a way to solve it?

[rusty1s]

Does it work when dropping the BatchNorm call on the input node features? If you desperately need normalization on sparse input features, you have to write a sparse BatchNorm variant by yourself, e.g.:

from torch_scatter import scatter_mean, scatter_std

row, col, value = input.coo()
mean_value = scatter_mean(value, col, dim=0)
std_value = scatter_std(value, col, dim=0)
value = (value - mean_value[col]) / (std_value[col] + 1e-6)
input = input.set_value(value, layout='coo')