Remove torch_scatter and torch_cluster from CI pipeline's dependencies #233
Conversation
dhpitt
changed the title
|
Small empirical tests show that my neighbor search function runs in about 1/3 the time of |
There was a problem hiding this comment.
Thanks for the changes @dhpitt -- should we have a small unit-test for the fallbacks?
neuralop/layers/neighbor_search.py
Outdated
| device=neighbors_count.device)), | ||
| dim=0) | ||
| if not self.use_torch_cluster: | ||
| return_dict = self.search_fn(data, queries, radius) |
There was a problem hiding this comment.
@kovachki Should we just keep the manual implementation and just choose between open3d (fast) and fallback (slow)?
There was a problem hiding this comment.
Yes, I think this is best. Let's just remove cluster altogether.
| if torch.backends.cuda.is_built(): | ||
| """only import torch_scatter when cuda is available""" | ||
| import torch_scatter.segment_csr as scatter_segment_csr | ||
| return scatter_segment_csr(src, indptr, reduce) |
There was a problem hiding this comment.
Hmm cuda being built doesn't necessarily mean torch_scatter will be available? Shouldn't we check for both?
There was a problem hiding this comment.
For now I thought we would keep torch_scatter as a dependency, but only for GPU users. In that case is it reasonable to assume it's available whenever CUDA is?
There was a problem hiding this comment.
@dhpitt could you please check that your segment_csr works with backwards. You're doing some inplace operations so I'm not sure.
|
@kovachki suggested also checking whether this works with backwards -- can you do add a simple sum loss and check for backward in the unit-test? |
|
Awesome, thanks @dhpitt ! |
We currently use the module
torch_scatteron main for one function calledsegment_csr, which aggregates features across each neighborhood for the GINO model inneuralop.layers.integral_transform.py. This implements a simple python version of the same function to remove our dependence ontorch_scatterand fix that annoying OSError in our CI.Second note: the same error occurs with the CPU version of
torch_cluster, so I wrote a neighborhood search for CPU runners only.