search

teamtomo / libtilt

Image processing for cryo-electron microscopy in PyTorch
Other
14 stars 5 forks source link

unittesting device agnostic code #72

Closed McHaillet closed 2 weeks ago

McHaillet commented 3 months ago

Some of the code is giving issues when running on cuda devices and I was thinking whether we could have device testing integrated in unittests.

To highligth some things, backproject_fourier is not running with cuda tensors and I noticed @rsanchezgarc also pointed out some issues in PR #54.

Running all the tests with both tensors on cpu and a cuda device would of course solve it. However, if there would be automated testing for PR's in the future, the default GitHub instances do not support this. Secondly, nice about pytorch is that it also allows development on a CPU only system and still have portability for GPUs.

I was looking online, cause I imagined more people need this, and found that torch supports a 'meta' device: https://pytorch.org/torchdistx/latest/fake_tensor.html . Also see the discussion here (https://github.com/pytorch/pytorch/issues/61654). I don't know whether its fully supported throughout pytorch, but I could play around with it and see how applicable it is.

For example though, the following produces appropriate errors:

>>> import torch
>>> a = torch.zeros((10,10))  # default will initialize on 'cpu'
>>> b = torch.zeros((10,10), device='meta')
>>> c = a * b
RuntimeError: Tensor on device meta is not on the expected device cpu!
alisterburt commented 3 months ago

Totally agree this is a current limitation - it looks like the meta device is an incomplete solution and we can't actually run any compute on that device, is that correct? Otherwise I'm all for anything that makes automated testing of this kind of thing possible in GitHub actions

McHaillet commented 3 months ago

it looks like the meta device is an incomplete solution and we can't actually run any compute on that device, is that correct?

Yep, that's correct! So we would need to run each test with both the 'cpu' and 'meta' device.

alisterburt commented 3 months ago

Nice! Seems like we could relatively easily write a decorator for tests that would make sure things work with the device on both inputs