How to use double sparse tensor with torch geometric?

[ray38]

Hi,

For my model, I have a node feature (called prime) that is stored as sparse tensor, which is multiplied by a double tensor (model derivative, named "gradient") during model "forward".

• feature preparation

data = Data(...)

# init, prime row, col, values, size are pre-computed
indices = np.vstack((prime_row, prime_col))
torch_indices = torch.LongTensor(indices)
torch_values = torch.tensor(prime_val, dtype=torch.DoubleTensor)

primes = torch.sparse.FloatTensor(torch_indices, torch_values, torch.Size(prime_size))
data.primes = primes

• Model Forward

def forward(self, batch):
    ....
    temp_val = torch.sparse.mm(batch.primes.t(), gradient .t())
    ....

I have two questions and would really appreciate some help:
1: currently, the model won't even train with the gpu version of pytorch, and game me an error (shown below), but it Does work with the cpu only version of pytorch. Seems it's realted to pin_memory, how could I get around this?

Traceback (most recent call last):
  File "train_ordernorm_model_custom_model_sqrt_lin_gelu_augmentation_precompute_force_test.py", line 233, in <module>
    trainer.train()
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/amptorch/trainer.py", line 341, in train
    self.net.fit(self.train_dataset, None)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/skorch/regressor.py", line 91, in fit
    return super(NeuralNetRegressor, self).fit(X, y, **fit_params)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/skorch/net.py", line 917, in fit
    self.partial_fit(X, y, **fit_params)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/skorch/net.py", line 876, in partial_fit
    self.fit_loop(X, y, **fit_params)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/skorch/net.py", line 790, in fit_loop
    step_fn=self.train_step, **fit_params)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/skorch/net.py", line 822, in run_single_epoch
    for data in self.get_iterator(dataset, training=training):
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 563, in _next_data
    data = _utils.pin_memory.pin_memory(data)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch/utils/data/_utils/pin_memory.py", line 58, in pin_memory
    return [pin_memory(sample) for sample in data]
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch/utils/data/_utils/pin_memory.py", line 58, in <listcomp>
    return [pin_memory(sample) for sample in data]
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch/utils/data/_utils/pin_memory.py", line 58, in pin_memory
    return [pin_memory(sample) for sample in data]
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch/utils/data/_utils/pin_memory.py", line 58, in <listcomp>
    return [pin_memory(sample) for sample in data]
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch/utils/data/_utils/pin_memory.py", line 60, in pin_memory
    return data.pin_memory()
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch_geometric/data/data.py", line 236, in pin_memory
    return self.apply(lambda x: x.pin_memory(), *args)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch_geometric/data/data.py", line 204, in apply
    store.apply(func, *args)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch_geometric/data/storage.py", line 146, in apply
    self[key] = recursive_apply(value, func)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch_geometric/data/storage.py", line 495, in recursive_apply
    return func(data)
  File "/home/xlei38/anaconda2/envs/CEMT3/lib/python3.6/site-packages/torch_geometric/data/data.py", line 236, in <lambda>
    return self.apply(lambda x: x.pin_memory(), *args)
NotImplementedError: Cannot access storage of SparseTensorImpl

2. My understanding is that the sparse tensor is of single precision. Is it possible to use Double precision instead?

Thanks a lot for your help! Please let me know if you need more details about my model

[rusty1s]

Where do you call pin_memory? Inside DataLoader? It looks like pin_memory() is not support for PyTorch sparse tensors.
I think double precision should be supported by just passing in a double tensor for values and creating the matrix via torch.sparse_coo_tensor(...). Best to ask directly in the PyTorch forums IMO.

[ray38]

@rusty1s Thanks a lot!
the only place I know that's related to pin_memory is when I initialize the neural net regressor

                net = NeuralNetRegressor(
                    module=self.model,
                    criterion=self.criterion,
                    ...
                    iterator_train__shuffle=False,
                    iterator_train__pin_memory=True,
                    iterator_valid__collate_fn=self.parallel_collater,
                    iterator_valid__shuffle=False,
                    iterator_valid__pin_memory=True,
                    device=self.device,
                    ...
                )

[rusty1s]

Sorry, I missed your reply. What happens if you set pin_memory args to False? Otherwise, I feel this issue seems more related to missing pin_memory support in torch.sparse? BTW, which PyG version are you using? I feel this should be already fixed in later versions.