Unable to run the code with error importing 'spmm'

[PengfeiHePower]

Hi, I followed the instruction to reproduce results but had a problem with module 'spmm'. My torch version is 1.8.2, torchdrug is 0.1.2. Any ideas how to fix it?

12:53:15 Epoch 0 begin
Traceback (most recent call last):
File "script/run.py", line 78, in
File "script/run.py", line 30, in train_and_validate
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\core\engine.py", line 143, in train
loss, metric = model(batch)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\tasks\reasoning.py", line 85, in forward
pred = self.predict(batch, all_loss, metric)
File "C:\Users\Pengfei\Documents\cse research\NBFNet-master\nbfnet\task.py", line 288, in predict
pred = self.model(graph, h_index, t_index, r_index, all_loss=all_loss, metric=metric)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "C:\Users\Pengfei\Documents\cse research\NBFNet-master\nbfnet\model.py", line 149, in forward
output = self.bellmanford(graph, h_index[:, 0], r_index[:, 0])
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\decorator.py", line 232, in fun
return caller(func, *(extras + args), **kw)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\utils\decorator.py", line 56, in wrapper
return forward(self, *args, **kwargs)
File "C:\Users\Pengfei\Documents\cse research\NBFNet-master\nbfnet\model.py", line 115, in bellmanford
hidden = layer(step_graph, layer_input)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torch\nn\modules\module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\layers\conv.py", line 91, in forward
update = self.message_and_aggregate(graph, input)
File "C:\Users\Pengfei\Documents\cse research\NBFNet-master\nbfnet\layer.py", line 140, in message_and_aggregate
sum = functional.generalized_rspmm(adjacency, relation_input, input, sum="add", mul=mul)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\layers\functional\spmm.py", line 378, in generalized_rspmm
return Function.apply(sparse.coalesce(), relation, input)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\layers\functional\spmm.py", line 172, in forward
forward = spmm.rspmm_add_mul_forward_cuda
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\utils\torch.py", line 27, in getattr
return getattr(self.module, key)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\utils\decorator.py", line 21, in get
result = self.func(obj)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torchdrug\utils\torch.py", line 31, in module
return cpp_extension.load(self.name, self.sources, self.extra_cflags, self.extra_cuda_cflags,
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torch\utils\cpp_extension.py", line 1079, in load
return _jit_compile(
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torch\utils\cpp_extension.py", line 1317, in _jit_compile
return _import_module_from_library(name, build_directory, is_python_module)
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\site-packages\torch\utils\cpp_extension.py", line 1700, in _import_module_from_library
file, path, description = imp.find_module(module_name, [path])
File "C:\Users\Pengfei\anaconda3\envs\py38\lib\imp.py", line 296, in find_module
raise ImportError(_ERR_MSG.format(name), name=name)
ImportError: No module named 'spmm'

[KiddoZhu]

Hi! spmm is compiled by TorchDrug on the fly when you run the code. So you need to prepare the correct compilation tools for that. For windows, please see the installation instructions of TorchDrug for details.

[jnaeg]

Hello dear authors of NBFNet,

trying to reproduce NBFNet, and directly running into compiling issues of TorchDrug's spmm and its other cpp_extensions,
I realised that the conda environment needed codatoolkit-dev installed in the env, not simple cudatoolkit.
Maybe adjust your installation instructions accordingly? (As otherwise nvcc is missing)

Kind regards, and thank you for sharing this lovely piece of efficient code with us!
jnaeg