pytorch
a688b297 - Support custom Python classes in CUDAFuture (#56516)

Commit
3 years ago
Support custom Python classes in CUDAFuture (#56516) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/56516 One problem with CUDAFuture's extraction of DataPtrs from IValues is that it only supported Python objects that could be converted to "regular" IValues (e.g., lists/dicts/tuples of ints/strings/tensors/...). One notable exception are custom Python classes, which are in fact a very common data type transferred over RPC. The only solution we found for those is to use the Python pickler to extract the tensors contained in them. We can't insert a Python dependency directly into CUDAFuture, so instead I'm proposing to use the same indirection technique used to support `getSubValues` on Python objects: define some methods on the abstract class `PyObjectHolder` (which can be used by CUDAFuture) but only implement them in the concrete subclass `ConcretePyObjectHolder` (which is only built when Python support is enabled). I am a bit worried about the performance toll of this (pickling isn't exactly known to be cheap) but I think we should start by providing a functionally complete API. We already have ideas on how to make this faster if needed, for example by having users provide a custom DataPtr extractor tailored to their class via a decorator. (Or just use TorchScript). ghstack-source-id: 127295014 Test Plan: Added a test later in the stack Reviewed By: mrshenli Differential Revision: D27887189 fbshipit-source-id: 9d27e4e62390b836e5bb4f06f401cc002f0cf95b
Author
lw lw
Parents
Loading