_classes.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. import types
  2. import torch._C
  3. class _ClassNamespace(types.ModuleType):
  4. def __init__(self, name):
  5. super().__init__("torch.classes" + name)
  6. self.name = name
  7. def __getattr__(self, attr):
  8. proxy = torch._C._get_custom_class_python_wrapper(self.name, attr)
  9. if proxy is None:
  10. raise RuntimeError(f"Class {self.name}.{attr} not registered!")
  11. return proxy
  12. class _Classes(types.ModuleType):
  13. __file__ = "_classes.py"
  14. def __init__(self):
  15. super().__init__("torch.classes")
  16. def __getattr__(self, name):
  17. namespace = _ClassNamespace(name)
  18. setattr(self, name, namespace)
  19. return namespace
  20. @property
  21. def loaded_libraries(self):
  22. return torch.ops.loaded_libraries
  23. def load_library(self, path):
  24. """
  25. Loads a shared library from the given path into the current process.
  26. The library being loaded may run global initialization code to register
  27. custom classes with the PyTorch JIT runtime. This allows dynamically
  28. loading custom classes. For this, you should compile your class
  29. and the static registration code into a shared library object, and then
  30. call ``torch.classes.load_library('path/to/libcustom.so')`` to load the
  31. shared object.
  32. After the library is loaded, it is added to the
  33. ``torch.classes.loaded_libraries`` attribute, a set that may be inspected
  34. for the paths of all libraries loaded using this function.
  35. Args:
  36. path (str): A path to a shared library to load.
  37. """
  38. torch.ops.load_library(path)
  39. # The classes "namespace"
  40. classes = _Classes()