Torch scatter executes functional code upon import

[TimZaman]

To keep imports in Python snappy, it's wise not to execute any functional code;

import torch_scatter

Creates a cuda context on all visible gpus due to the call https://github.com/rusty1s/pytorch_scatter/blob/master/torch_scatter/__init__.py#L9

suffix = 'cuda' if torch.cuda.is_available() else 'cpu'

This statement torch.cuda.is_available() should be encapsulated and not be executed upon import.

I'm a bit busy to make a PR now [sorry], just putting this here for posterity.

[rusty1s]

Interesting. Do you know of another way to check for CUDA availability without creating a CUDA context?

[TimZaman]

Best is to simply do it lazily and save in a global variable.

…

On Fri, Apr 16, 2021, 00:05 Matthias Fey ***@***.***> wrote: Interesting. Do you know of another way to check for CUDA availability without creating a CUDA context? — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#209 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AB25ERAOJYKR4YK2GWXA44TTI7OVDANCNFSM43AG2FQQ> .

[github-actions]

This issue had no activity for 6 months. It will be closed in 2 weeks unless there is some new activity. Is this issue already resolved?

[rusty1s]

Still need to look into that. Thanks for reminding, and sorry for the delay.

[TimZaman]

$ time python -c "import torch_scatter"

real	0m3.838s
user	0m0.580s
sys	0m0.154s

This is a result of everything being thrown into the __init__.py, disabling any user to lazily import what they need

[rusty1s]

Interesting, it's way faster for me. Most of the time in import torch_scatter is spent on importing torch (1.62s), and only 0.05s are used to run the remaining code in __init__.py.

[TimZaman]

The reason it's fast for you is probably because you're on a 1 gpu system. The functional calls describred above will create a cuda context on each gpu at import-time. That's pretty bad, because that is before anyone would have the chance to eg call torch.cuda.set_device

[rusty1s]

Thanks :)

I think all we need to do for this is to replace the torch.cuda.is_available call here by torch.version.cuda. Can you confirm?