Skip to content

ManagedModule Class Reference

Bases: torch.nn.Module

A module that uses managed tensors. It can be used to convert a module to use managed tensors or can be used as a base class for a module.

Source code in managed/module.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class ManagedModule(torch.nn.Module):
    """
    A module that uses managed tensors. It can be used to convert a module
    to use managed tensors or can be used as a base class for a module.
    """
    @classmethod
    def from_module(cls, module: torch.nn.Module, tensor_cls=ManagedTensor):
        """
        Converts a module to use managed tensors.
        """
        if isinstance(module, cls):
            return module
        if isinstance(module, torch.nn.Module):
            module._apply(wrap_tensor(tensor_cls))
            for m in module.children():
                cls.from_module(m)
            module_class = module.__class__
            new_cls = type(
                f"{cls.__name__}.{module_class.__name__}",
                (module_class, cls),
                {}
            )
            module.__class__ = new_cls
        else:
            raise TypeError(f"Expected torch Module, got {type(module)}")
        return module

    def register_parameter(self, name: str, param: Optional[torch.nn.Parameter]) -> None:
        super().register_parameter(name, param)
        wrap_tensor(ManagedTensor)(param)

    def add_module(self, name: str, module: Optional[torch.nn.Module]) -> None:
        super().add_module(name, module)
        self.from_module(module)

    def cuda(self, *args, **kwargs):
        self._apply(lambda t: t.cuda(*args, **kwargs))
        return self

    def apply_(self, fn):
        self._apply(_inplace_func_wrap(fn))

    def pin(self):
        """
        Pins all the tensors in the module.
        """
        self.apply_(lambda t: t.pin())

    def unpin(self):
        """
        Unpins all the tensors in the module.
        """
        self.apply_(lambda t: t.unpin())

from_module(module, tensor_cls=ManagedTensor) classmethod

Converts a module to use managed tensors.

Source code in managed/module.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
@classmethod
def from_module(cls, module: torch.nn.Module, tensor_cls=ManagedTensor):
    """
    Converts a module to use managed tensors.
    """
    if isinstance(module, cls):
        return module
    if isinstance(module, torch.nn.Module):
        module._apply(wrap_tensor(tensor_cls))
        for m in module.children():
            cls.from_module(m)
        module_class = module.__class__
        new_cls = type(
            f"{cls.__name__}.{module_class.__name__}",
            (module_class, cls),
            {}
        )
        module.__class__ = new_cls
    else:
        raise TypeError(f"Expected torch Module, got {type(module)}")
    return module

pin()

Pins all the tensors in the module.

Source code in managed/module.py 
71
72
73
74
75
def pin(self):
    """
    Pins all the tensors in the module.
    """
    self.apply_(lambda t: t.pin())

unpin()

Unpins all the tensors in the module.

Source code in managed/module.py 
77
78
79
80
81
def unpin(self):
    """
    Unpins all the tensors in the module.
    """
    self.apply_(lambda t: t.unpin())