pymoten icon indicating copy to clipboard operation
pymoten copied to clipboard

WIP+ENH add gpu capabilities

Open eickenberg opened this issue 4 years ago • 1 comments

Very much still work in progress. This is as yet untested, but I am putting it online because I cannot currently find the recommendations of how this should be integrated and need some input.

Should I leave the torch code as an independent core_torch file and keep it minimal?

Another option is to try to operate with backends, since the files do not differ that much, but that would involve a lot more work, and modifications to the remaining code.

Any input appreciated, since I currently don't have the big picture.

eickenberg avatar May 09 '20 02:05 eickenberg

Sweet! I think we can keep it simple for now. A separate core_torch would be fine for the time being. The pyramid class that will end up using this can be explicit and a bit ad-hoc for the time being (e.g. moten.pyramids.MotionEnergyPyramidGPU()).

Long term, I think you're right and operating with a backend would be better and will require a lot more work. Though if we replace the np.dot with the @ syntax, we'll be 80% there.

An intermediate option is to do something hacky :

xp = torch if backend=='cuda' else np
if backend == 'cuda': 
    xp.dot = xp.mm

But maybe let's not do that, yet =)

I think it's better as you suggest and keep things independent and minimal. We can then refactor for a backend. On that front, any pointers to implementing a numeric backend would be useful.

anwarnunez avatar May 09 '20 12:05 anwarnunez