Skip to content

zamba.pytorch.transforms

imagenet_normalization_values

Classes

ConvertTCHWtoCTHW (Module)

Convert tensor from (T, C, H, W) to (C, T, H, W)

Source code in zamba/pytorch/transforms.py
class ConvertTCHWtoCTHW(torch.nn.Module):
    """Convert tensor from (T, C, H, W) to (C, T, H, W)"""

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.permute(1, 0, 2, 3)

Attributes

T_destination inherited
dump_patches: bool inherited

This allows better BC support for :meth:load_state_dict. In :meth:state_dict, the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.

If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module's _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.

Methods

__init__(self) -> None inherited special
Source code in zamba/pytorch/transforms.py
def __init__(self) -> None:
    """
    Initializes internal Module state, shared by both nn.Module and ScriptModule.
    """
    torch._C._log_api_usage_once("python.nn_module")

    self.training = True
    self._parameters: Dict[str, Optional[Parameter]] = OrderedDict()
    self._buffers: Dict[str, Optional[Tensor]] = OrderedDict()
    self._non_persistent_buffers_set: Set[str] = set()
    self._backward_hooks: Dict[int, Callable] = OrderedDict()
    self._is_full_backward_hook = None
    self._forward_hooks: Dict[int, Callable] = OrderedDict()
    self._forward_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._state_dict_hooks: Dict[int, Callable] = OrderedDict()
    self._load_state_dict_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._modules: Dict[str, Optional['Module']] = OrderedDict()
add_module(self, name: str, module: Optional[Module]) -> None inherited

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

Name Type Description Default
name string

name of the child module. The child module can be accessed from this module using the given name

required
module Module

child module to be added to the module.

required
Source code in zamba/pytorch/transforms.py
def add_module(self, name: str, module: Optional['Module']) -> None:
    r"""Adds a child module to the current module.

    The module can be accessed as an attribute using the given name.

    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\", got: {}".format(name))
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module
apply(self: ~T, fn: Callable[[Module], NoneType]) -> ~T inherited

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also :ref:nn-init-doc).

Parameters:

Name Type Description Default
fn

class:Module -> None): function to be applied to each submodule

required

Returns:

Type Description
Module

self

Example::

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Source code in zamba/pytorch/transforms.py
def apply(self: T, fn: Callable[['Module'], None]) -> T:
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`nn-init-doc`).

    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule

    Returns:
        Module: self

    Example::

        >>> @torch.no_grad()
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self
bfloat16(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to bfloat16 datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def bfloat16(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
buffers(self, recurse: bool = True) -> Iterator[torch.Tensor] inherited

Returns an iterator over module buffers.

Parameters:

Name Type Description Default
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
torch.Tensor

module buffer

Example::

>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
    r"""Returns an iterator over module buffers.

    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        torch.Tensor: module buffer

    Example::

        >>> for buf in model.buffers():
        >>>     print(type(buf), buf.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for _, buf in self.named_buffers(recurse=recurse):
        yield buf
children(self) -> Iterator[Module] inherited

Returns an iterator over immediate children modules.

Yields:

Type Description
Module

a child module

Source code in zamba/pytorch/transforms.py
def children(self) -> Iterator['Module']:
    r"""Returns an iterator over immediate children modules.

    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module
cpu(self: ~T) -> ~T inherited

Moves all model parameters and buffers to the CPU.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())
cuda(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))
double(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to double datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def double(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``double`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)
eval(self: ~T) -> ~T inherited

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

This is equivalent with :meth:self.train(False) <torch.nn.Module.train>.

See :ref:locally-disable-grad-doc for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def eval(self: T) -> T:
    r"""Sets the module in evaluation mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.eval()` and several similar mechanisms that may be confused with it.

    Returns:
        Module: self
    """
    return self.train(False)
extra_repr(self) -> str inherited

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Source code in zamba/pytorch/transforms.py
def extra_repr(self) -> str:
    r"""Set the extra representation of the module

    To print customized extra information, you should re-implement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''
float(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to float datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def float(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``float`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)
forward(self, vid: Tensor) -> Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

.. note:: Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Source code in zamba/pytorch/transforms.py
def forward(self, vid: torch.Tensor) -> torch.Tensor:
    return vid.permute(1, 0, 2, 3)
get_buffer(self, target: str) -> Tensor inherited

Returns the buffer given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.Tensor

The buffer referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not a buffer

Source code in zamba/pytorch/transforms.py
def get_buffer(self, target: str) -> "Tensor":
    """
    Returns the buffer given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the buffer
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.Tensor: The buffer referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not a
            buffer
    """
    module_path, _, buffer_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, buffer_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + buffer_name + "`")

    buffer: torch.Tensor = getattr(mod, buffer_name)

    if buffer_name not in mod._buffers:
        raise AttributeError("`" + buffer_name + "` is not a buffer")

    return buffer
get_extra_state(self) -> Any inherited

Returns any extra state to include in the module's state_dict. Implement this and a corresponding :func:set_extra_state for your module if you need to store extra state. This function is called when building the module's state_dict().

Note that extra state should be pickleable to ensure working serialization of the state_dict. We only provide provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Type Description
object

Any extra state to store in the module's state_dict

Source code in zamba/pytorch/transforms.py
def get_extra_state(self) -> Any:
    """
    Returns any extra state to include in the module's state_dict.
    Implement this and a corresponding :func:`set_extra_state` for your module
    if you need to store extra state. This function is called when building the
    module's `state_dict()`.

    Note that extra state should be pickleable to ensure working serialization
    of the state_dict. We only provide provide backwards compatibility guarantees
    for serializing Tensors; other objects may break backwards compatibility if
    their serialized pickled form changes.

    Returns:
        object: Any extra state to store in the module's state_dict
    """
    raise RuntimeError(
        "Reached a code path in Module.get_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
get_parameter(self, target: str) -> Parameter inherited

Returns the parameter given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Parameter

The Parameter referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Parameter

Source code in zamba/pytorch/transforms.py
def get_parameter(self, target: str) -> "Parameter":
    """
    Returns the parameter given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the Parameter
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Parameter: The Parameter referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Parameter``
    """
    module_path, _, param_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, param_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + param_name + "`")

    param: torch.nn.Parameter = getattr(mod, param_name)

    if not isinstance(param, torch.nn.Parameter):
        raise AttributeError("`" + param_name + "` is not an "
                             "nn.Parameter")

    return param
get_submodule(self, target: str) -> Module inherited

Returns the submodule given by target if it exists, otherwise throws an error.

For example, let's say you have an nn.Module A that looks like this:

.. code-block::text

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Module

The submodule referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Module

Source code in zamba/pytorch/transforms.py
def get_submodule(self, target: str) -> "Module":
    """
    Returns the submodule given by ``target`` if it exists,
    otherwise throws an error.

    For example, let's say you have an ``nn.Module`` ``A`` that
    looks like this:

    .. code-block::text

        A(
            (net_b): Module(
                (net_c): Module(
                    (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
                )
                (linear): Linear(in_features=100, out_features=200, bias=True)
            )
        )

    (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
    submodule ``net_b``, which itself has two submodules ``net_c``
    and ``linear``. ``net_c`` then has a submodule ``conv``.)

    To check whether or not we have the ``linear`` submodule, we
    would call ``get_submodule("net_b.linear")``. To check whether
    we have the ``conv`` submodule, we would call
    ``get_submodule("net_b.net_c.conv")``.

    The runtime of ``get_submodule`` is bounded by the degree
    of module nesting in ``target``. A query against
    ``named_modules`` achieves the same result, but it is O(N) in
    the number of transitive modules. So, for a simple check to see
    if some submodule exists, ``get_submodule`` should always be
    used.

    Args:
        target: The fully-qualified string name of the submodule
            to look for. (See above example for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Module: The submodule referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Module``
    """
    if target == "":
        return self

    atoms: List[str] = target.split(".")
    mod: torch.nn.Module = self

    for item in atoms:

        if not hasattr(mod, item):
            raise AttributeError(mod._get_name() + " has no "
                                 "attribute `" + item + "`")

        mod = getattr(mod, item)

        if not isinstance(mod, torch.nn.Module):
            raise AttributeError("`" + item + "` is not "
                                 "an nn.Module")

    return mod
half(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to half datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def half(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``half`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)
load_state_dict(self, state_dict: OrderedDict[str, Tensor], strict: bool = True) inherited

Copies parameters and buffers from :attr:state_dict into this module and its descendants. If :attr:strict is True, then the keys of :attr:state_dict must exactly match the keys returned by this module's :meth:~torch.nn.Module.state_dict function.

Parameters:

Name Type Description Default
state_dict dict

a dict containing parameters and persistent buffers.

required
strict bool

whether to strictly enforce that the keys in :attr:state_dict match the keys returned by this module's :meth:~torch.nn.Module.state_dict function. Default: True

True

Returns:

Type Description
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields
  • missing_keys is a list of str containing the missing keys
    • unexpected_keys is a list of str containing the unexpected keys

!!! note If a parameter or buffer is registered as None and its corresponding key exists in :attr:state_dict, :meth:load_state_dict will raise a RuntimeError.

Source code in zamba/pytorch/transforms.py
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                    strict: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``

    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Note:
        If a parameter or buffer is registered as ``None`` and its corresponding key
        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
        ``RuntimeError``.
    """
    missing_keys: List[str] = []
    unexpected_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        # mypy isn't aware that "_metadata" exists in state_dict
        state_dict._metadata = metadata  # type: ignore[attr-defined]

    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')

    load(self)
    del load

    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))

    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)
modules(self) -> Iterator[Module] inherited

Returns an iterator over all modules in the network.

Yields:

Type Description
Module

a module in the network

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
Source code in zamba/pytorch/transforms.py
def modules(self) -> Iterator['Module']:
    r"""Returns an iterator over all modules in the network.

    Yields:
        Module: a module in the network

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)

        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)

    """
    for _, module in self.named_modules():
        yield module
named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]] inherited

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all buffer names.

''
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
(string, torch.Tensor)

Tuple containing the name and buffer

Example::

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
Source code in zamba/pytorch/transforms.py
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.

    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer

    Example::

        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())

    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
named_children(self) -> Iterator[Tuple[str, Module]] inherited

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

Type Description
(string, Module)

Tuple containing a name and child module

Example::

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Source code in zamba/pytorch/transforms.py
def named_children(self) -> Iterator[Tuple[str, 'Module']]:
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.

    Yields:
        (string, Module): Tuple containing a name and child module

    Example::

        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)

    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module
named_modules(self, memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True) inherited

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

Name Type Description Default
memo Optional[Set[Module]]

a memo to store the set of modules already added to the result

None
prefix str

a prefix that will be added to the name of the module

''
remove_duplicate bool

whether to remove the duplicated module instances in the result

True

Yields:

Type Description
(string, Module)

Tuple of name and module

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
Source code in zamba/pytorch/transforms.py
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = '', remove_duplicate: bool = True):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.

    Args:
        memo: a memo to store the set of modules already added to the result
        prefix: a prefix that will be added to the name of the module
        remove_duplicate: whether to remove the duplicated module instances in the result
        or not

    Yields:
        (string, Module): Tuple of name and module

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)

        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

    """

    if memo is None:
        memo = set()
    if self not in memo:
        if remove_duplicate:
            memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix, remove_duplicate):
                yield m
named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]] inherited

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all parameter names.

''
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
(string, Parameter)

Tuple containing the name and parameter

Example::

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
Source code in zamba/pytorch/transforms.py
def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.

    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        (string, Parameter): Tuple containing the name and parameter

    Example::

        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())

    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter] inherited

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

Name Type Description Default
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
Parameter

module parameter

Example::

>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters.

    This is typically passed to an optimizer.

    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        Parameter: module parameter

    Example::

        >>> for param in model.parameters():
        >>>     print(type(param), param.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param
register_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

This function is deprecated in favor of :meth:~torch.nn.Module.register_full_backward_hook and the behavior of this function will change in future versions.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
    the behavior of this function will change in future versions.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is True:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = False

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None inherited

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

Name Type Description Default
name string

name of the buffer. The buffer can be accessed from this module using the given name

required
tensor Tensor or None

buffer to be registered. If None, then operations that run on buffers, such as :attr:cuda, are ignored. If None, the buffer is not included in the module's :attr:state_dict.

required
persistent bool

whether the buffer is part of this module's :attr:state_dict.

True

Example::

>>> self.register_buffer('running_mean', torch.zeros(num_features))
Source code in zamba/pytorch/transforms.py
def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    r"""Adds a buffer to the module.

    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the module's state. Buffers, by
    default, are persistent and will be saved alongside parameters. This
    behavior can be changed by setting :attr:`persistent` to ``False``. The
    only difference between a persistent buffer and a non-persistent buffer
    is that the latter will not be a part of this module's
    :attr:`state_dict`.

    Buffers can be accessed as attributes using given names.

    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor or None): buffer to be registered. If ``None``, then operations
            that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
            the buffer is **not** included in the module's :attr:`state_dict`.
        persistent (bool): whether the buffer is part of this module's
            :attr:`state_dict`.

    Example::

        >>> self.register_buffer('running_mean', torch.zeros(num_features))

    """
    if persistent is False and isinstance(self, torch.jit.ScriptModule):
        raise RuntimeError("ScriptModule does not support non-persistent buffers")

    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor
        if persistent:
            self._non_persistent_buffers_set.discard(name)
        else:
            self._non_persistent_buffers_set.add(name)
register_forward_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward hook on the module.

The hook will be called every time after :func:forward has computed an output. It should have the following signature::

hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after :func:forward is called.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward hook on the module.

    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::

        hook(module, input, output) -> None or modified output

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the output. It can modify the input inplace but
    it will not have effect on forward since this is called after
    :func:`forward` is called.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle
register_forward_pre_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward pre-hook on the module.

The hook will be called every time before :func:forward is invoked. It should have the following signature::

hook(module, input) -> None or modified input

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward pre-hook on the module.

    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::

        hook(module, input) -> None or modified input

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the input. User can either return a tuple or a
    single modified value in the hook. We will wrap the value into a tuple
    if a single value is returned(unless that value is already a tuple).

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle
register_full_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature::

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The :attr:grad_input and :attr:grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of :attr:grad_input in subsequent computations. :attr:grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in :attr:grad_input and :attr:grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module's forward function.

.. warning :: Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_full_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::

        hook(module, grad_input, grad_output) -> tuple(Tensor) or None

    The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
    with respect to the inputs and outputs respectively. The hook should
    not modify its arguments, but it can optionally return a new gradient with
    respect to the input that will be used in place of :attr:`grad_input` in
    subsequent computations. :attr:`grad_input` will only correspond to the inputs given
    as positional arguments and all kwarg arguments are ignored. Entries
    in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
    arguments.

    For technical reasons, when this hook is applied to a Module, its forward function will
    receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
    of each Tensor returned by the Module's forward function.

    .. warning ::
        Modifying inputs or outputs inplace is not allowed when using backward hooks and
        will raise an error.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is False:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = True

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None inherited

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

Name Type Description Default
name string

name of the parameter. The parameter can be accessed from this module using the given name

required
param Parameter or None

parameter to be added to the module. If None, then operations that run on parameters, such as :attr:cuda, are ignored. If None, the parameter is not included in the module's :attr:state_dict.

required
Source code in zamba/pytorch/transforms.py
def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    r"""Adds a parameter to the module.

    The parameter can be accessed as an attribute using given name.

    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter or None): parameter to be added to the module. If
            ``None``, then operations that run on parameters, such as :attr:`cuda`,
            are ignored. If ``None``, the parameter is **not** included in the
            module's :attr:`state_dict`.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")

    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))

    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param
requires_grad_(self: ~T, requires_grad: bool = True) -> ~T inherited

Change if autograd should record operations on parameters in this module.

This method sets the parameters' :attr:requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See :ref:locally-disable-grad-doc for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

Name Type Description Default
requires_grad bool

whether autograd should record operations on parameters in this module. Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def requires_grad_(self: T, requires_grad: bool = True) -> T:
    r"""Change if autograd should record operations on parameters in this
    module.

    This method sets the parameters' :attr:`requires_grad` attributes
    in-place.

    This method is helpful for freezing part of the module for finetuning
    or training parts of a model individually (e.g., GAN training).

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.requires_grad_()` and several similar mechanisms that may be confused with it.

    Args:
        requires_grad (bool): whether autograd should record operations on
                              parameters in this module. Default: ``True``.

    Returns:
        Module: self
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self
set_extra_state(self, state: Any) inherited

This function is called from :func:load_state_dict to handle any extra state found within the state_dict. Implement this function and a corresponding :func:get_extra_state for your module if you need to store extra state within its state_dict.

Parameters:

Name Type Description Default
state dict

Extra state from the state_dict

required
Source code in zamba/pytorch/transforms.py
def set_extra_state(self, state: Any):
    """
    This function is called from :func:`load_state_dict` to handle any extra state
    found within the `state_dict`. Implement this function and a corresponding
    :func:`get_extra_state` for your module if you need to store extra state within its
    `state_dict`.

    Args:
        state (dict): Extra state from the `state_dict`
    """
    raise RuntimeError(
        "Reached a code path in Module.set_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
share_memory(self: ~T) -> ~T inherited

See :meth:torch.Tensor.share_memory_

Source code in zamba/pytorch/transforms.py
def share_memory(self: T) -> T:
    r"""See :meth:`torch.Tensor.share_memory_`"""
    return self._apply(lambda t: t.share_memory_())
state_dict(self, destination = None, prefix = '', keep_vars = False) inherited

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Returns:

Type Description
dict

a dictionary containing a whole state of the module

Example::

>>> module.state_dict().keys()
['bias', 'weight']
Source code in zamba/pytorch/transforms.py
def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.

    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Parameters and buffers set to ``None`` are not included.

    Returns:
        dict:
            a dictionary containing a whole state of the module

    Example::

        >>> module.state_dict().keys()
        ['bias', 'weight']

    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    self._save_to_state_dict(destination, prefix, keep_vars)
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination
to(self, *args, **kwargs) inherited

Moves and/or casts the parameters and buffers.

This can be called as

.. function:: to(device=None, dtype=None, non_blocking=False) :noindex:

.. function:: to(dtype, non_blocking=False) :noindex:

.. function:: to(tensor, non_blocking=False) :noindex:

.. function:: to(memory_format=torch.channels_last) :noindex:

Its signature is similar to :meth:torch.Tensor.to, but only accepts floating point or complex :attr:dtype\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:dtype (if given). The integral parameters and buffers will be moved :attr:device, if that is given, but with dtypes unchanged. When :attr:non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device

class:torch.device): the desired device of the parameters and buffers in this module

required
dtype

class:torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module

required
tensor torch.Tensor

Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

required
memory_format

class:torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword only argument)

required

Returns:

Type Description
Module

self

Examples::

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
Source code in zamba/pytorch/transforms.py
def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.

    This can be called as

    .. function:: to(device=None, dtype=None, non_blocking=False)
       :noindex:

    .. function:: to(dtype, non_blocking=False)
       :noindex:

    .. function:: to(tensor, non_blocking=False)
       :noindex:

    .. function:: to(memory_format=torch.channels_last)
       :noindex:

    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`\ s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.

    See below for examples.

    .. note::
        This method modifies the module in-place.

    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)

    Returns:
        Module: self

    Examples::

        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)

        >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.3741+0.j,  0.2382+0.j],
                [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
        >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
        tensor([[0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

    """

    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)

    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError('nn.Module.to only accepts floating point or complex '
                            'dtypes, but got desired dtype={}'.format(dtype))
        if dtype.is_complex:
            warnings.warn(
                "Complex modules are a new feature under active development whose design may change, "
                "and some modules might not work as expected when using complex tensors as parameters or buffers. "
                "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
                "if a complex module does not work as expected.")

    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

    return self._apply(convert)
to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T inherited

Moves the parameters and buffers to the specified device without copying storage.

Parameters:

Name Type Description Default
device

class:torch.device): The desired device of the parameters and buffers in this module.

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def to_empty(self: T, *, device: Union[str, device]) -> T:
    r"""Moves the parameters and buffers to the specified device without copying storage.

    Args:
        device (:class:`torch.device`): The desired device of the parameters
            and buffers in this module.

    Returns:
        Module: self
    """
    return self._apply(lambda t: torch.empty_like(t, device=device))
train(self: ~T, mode: bool = True) -> ~T inherited

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Parameters:

Name Type Description Default
mode bool

whether to set training mode (True) or evaluation mode (False). Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def train(self: T, mode: bool = True) -> T:
    r"""Sets the module in training mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    Args:
        mode (bool): whether to set training mode (``True``) or evaluation
                     mode (``False``). Default: ``True``.

    Returns:
        Module: self
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self
type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T inherited

Casts all parameters and buffers to :attr:dst_type.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
dst_type type or string

the desired type

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def type(self: T, dst_type: Union[dtype, str]) -> T:
    r"""Casts all parameters and buffers to :attr:`dst_type`.

    .. note::
        This method modifies the module in-place.

    Args:
        dst_type (type or string): the desired type

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))
xpu(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the XPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on XPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.xpu(device))
zero_grad(self, set_to_none: bool = False) -> None inherited

Sets gradients of all model parameters to zero. See similar function under :class:torch.optim.Optimizer for more context.

Parameters:

Name Type Description Default
set_to_none bool

instead of setting to zero, set the grads to None. See :meth:torch.optim.Optimizer.zero_grad for details.

False
Source code in zamba/pytorch/transforms.py
def zero_grad(self, set_to_none: bool = False) -> None:
    r"""Sets gradients of all model parameters to zero. See similar function
    under :class:`torch.optim.Optimizer` for more context.

    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            See :meth:`torch.optim.Optimizer.zero_grad` for details.
    """
    if getattr(self, '_is_replica', False):
        warnings.warn(
            "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
            "The parameters are copied (in a differentiable manner) from the original module. "
            "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
            "If you need gradients in your forward method, consider using autograd.grad instead.")

    for p in self.parameters():
        if p.grad is not None:
            if set_to_none:
                p.grad = None
            else:
                if p.grad.grad_fn is not None:
                    p.grad.detach_()
                else:
                    p.grad.requires_grad_(False)
                p.grad.zero_()

ConvertTHWCtoCTHW (Module)

Convert tensor from (0:T, 1:H, 2:W, 3:C) to (3:C, 0:T, 1:H, 2:W)

Source code in zamba/pytorch/transforms.py
class ConvertTHWCtoCTHW(torch.nn.Module):
    """Convert tensor from (0:T, 1:H, 2:W, 3:C) to (3:C, 0:T, 1:H, 2:W)"""

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.permute(3, 0, 1, 2)

Attributes

T_destination inherited
dump_patches: bool inherited

This allows better BC support for :meth:load_state_dict. In :meth:state_dict, the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.

If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module's _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.

Methods

__init__(self) -> None inherited special
Source code in zamba/pytorch/transforms.py
def __init__(self) -> None:
    """
    Initializes internal Module state, shared by both nn.Module and ScriptModule.
    """
    torch._C._log_api_usage_once("python.nn_module")

    self.training = True
    self._parameters: Dict[str, Optional[Parameter]] = OrderedDict()
    self._buffers: Dict[str, Optional[Tensor]] = OrderedDict()
    self._non_persistent_buffers_set: Set[str] = set()
    self._backward_hooks: Dict[int, Callable] = OrderedDict()
    self._is_full_backward_hook = None
    self._forward_hooks: Dict[int, Callable] = OrderedDict()
    self._forward_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._state_dict_hooks: Dict[int, Callable] = OrderedDict()
    self._load_state_dict_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._modules: Dict[str, Optional['Module']] = OrderedDict()
add_module(self, name: str, module: Optional[Module]) -> None inherited

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

Name Type Description Default
name string

name of the child module. The child module can be accessed from this module using the given name

required
module Module

child module to be added to the module.

required
Source code in zamba/pytorch/transforms.py
def add_module(self, name: str, module: Optional['Module']) -> None:
    r"""Adds a child module to the current module.

    The module can be accessed as an attribute using the given name.

    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\", got: {}".format(name))
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module
apply(self: ~T, fn: Callable[[Module], NoneType]) -> ~T inherited

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also :ref:nn-init-doc).

Parameters:

Name Type Description Default
fn

class:Module -> None): function to be applied to each submodule

required

Returns:

Type Description
Module

self

Example::

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Source code in zamba/pytorch/transforms.py
def apply(self: T, fn: Callable[['Module'], None]) -> T:
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`nn-init-doc`).

    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule

    Returns:
        Module: self

    Example::

        >>> @torch.no_grad()
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self
bfloat16(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to bfloat16 datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def bfloat16(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
buffers(self, recurse: bool = True) -> Iterator[torch.Tensor] inherited

Returns an iterator over module buffers.

Parameters:

Name Type Description Default
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
torch.Tensor

module buffer

Example::

>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
    r"""Returns an iterator over module buffers.

    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        torch.Tensor: module buffer

    Example::

        >>> for buf in model.buffers():
        >>>     print(type(buf), buf.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for _, buf in self.named_buffers(recurse=recurse):
        yield buf
children(self) -> Iterator[Module] inherited

Returns an iterator over immediate children modules.

Yields:

Type Description
Module

a child module

Source code in zamba/pytorch/transforms.py
def children(self) -> Iterator['Module']:
    r"""Returns an iterator over immediate children modules.

    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module
cpu(self: ~T) -> ~T inherited

Moves all model parameters and buffers to the CPU.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())
cuda(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))
double(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to double datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def double(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``double`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)
eval(self: ~T) -> ~T inherited

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

This is equivalent with :meth:self.train(False) <torch.nn.Module.train>.

See :ref:locally-disable-grad-doc for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def eval(self: T) -> T:
    r"""Sets the module in evaluation mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.eval()` and several similar mechanisms that may be confused with it.

    Returns:
        Module: self
    """
    return self.train(False)
extra_repr(self) -> str inherited

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Source code in zamba/pytorch/transforms.py
def extra_repr(self) -> str:
    r"""Set the extra representation of the module

    To print customized extra information, you should re-implement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''
float(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to float datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def float(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``float`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)
forward(self, vid: Tensor) -> Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

.. note:: Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Source code in zamba/pytorch/transforms.py
def forward(self, vid: torch.Tensor) -> torch.Tensor:
    return vid.permute(3, 0, 1, 2)
get_buffer(self, target: str) -> Tensor inherited

Returns the buffer given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.Tensor

The buffer referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not a buffer

Source code in zamba/pytorch/transforms.py
def get_buffer(self, target: str) -> "Tensor":
    """
    Returns the buffer given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the buffer
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.Tensor: The buffer referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not a
            buffer
    """
    module_path, _, buffer_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, buffer_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + buffer_name + "`")

    buffer: torch.Tensor = getattr(mod, buffer_name)

    if buffer_name not in mod._buffers:
        raise AttributeError("`" + buffer_name + "` is not a buffer")

    return buffer
get_extra_state(self) -> Any inherited

Returns any extra state to include in the module's state_dict. Implement this and a corresponding :func:set_extra_state for your module if you need to store extra state. This function is called when building the module's state_dict().

Note that extra state should be pickleable to ensure working serialization of the state_dict. We only provide provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Type Description
object

Any extra state to store in the module's state_dict

Source code in zamba/pytorch/transforms.py
def get_extra_state(self) -> Any:
    """
    Returns any extra state to include in the module's state_dict.
    Implement this and a corresponding :func:`set_extra_state` for your module
    if you need to store extra state. This function is called when building the
    module's `state_dict()`.

    Note that extra state should be pickleable to ensure working serialization
    of the state_dict. We only provide provide backwards compatibility guarantees
    for serializing Tensors; other objects may break backwards compatibility if
    their serialized pickled form changes.

    Returns:
        object: Any extra state to store in the module's state_dict
    """
    raise RuntimeError(
        "Reached a code path in Module.get_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
get_parameter(self, target: str) -> Parameter inherited

Returns the parameter given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Parameter

The Parameter referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Parameter

Source code in zamba/pytorch/transforms.py
def get_parameter(self, target: str) -> "Parameter":
    """
    Returns the parameter given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the Parameter
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Parameter: The Parameter referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Parameter``
    """
    module_path, _, param_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, param_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + param_name + "`")

    param: torch.nn.Parameter = getattr(mod, param_name)

    if not isinstance(param, torch.nn.Parameter):
        raise AttributeError("`" + param_name + "` is not an "
                             "nn.Parameter")

    return param
get_submodule(self, target: str) -> Module inherited

Returns the submodule given by target if it exists, otherwise throws an error.

For example, let's say you have an nn.Module A that looks like this:

.. code-block::text

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Module

The submodule referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Module

Source code in zamba/pytorch/transforms.py
def get_submodule(self, target: str) -> "Module":
    """
    Returns the submodule given by ``target`` if it exists,
    otherwise throws an error.

    For example, let's say you have an ``nn.Module`` ``A`` that
    looks like this:

    .. code-block::text

        A(
            (net_b): Module(
                (net_c): Module(
                    (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
                )
                (linear): Linear(in_features=100, out_features=200, bias=True)
            )
        )

    (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
    submodule ``net_b``, which itself has two submodules ``net_c``
    and ``linear``. ``net_c`` then has a submodule ``conv``.)

    To check whether or not we have the ``linear`` submodule, we
    would call ``get_submodule("net_b.linear")``. To check whether
    we have the ``conv`` submodule, we would call
    ``get_submodule("net_b.net_c.conv")``.

    The runtime of ``get_submodule`` is bounded by the degree
    of module nesting in ``target``. A query against
    ``named_modules`` achieves the same result, but it is O(N) in
    the number of transitive modules. So, for a simple check to see
    if some submodule exists, ``get_submodule`` should always be
    used.

    Args:
        target: The fully-qualified string name of the submodule
            to look for. (See above example for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Module: The submodule referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Module``
    """
    if target == "":
        return self

    atoms: List[str] = target.split(".")
    mod: torch.nn.Module = self

    for item in atoms:

        if not hasattr(mod, item):
            raise AttributeError(mod._get_name() + " has no "
                                 "attribute `" + item + "`")

        mod = getattr(mod, item)

        if not isinstance(mod, torch.nn.Module):
            raise AttributeError("`" + item + "` is not "
                                 "an nn.Module")

    return mod
half(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to half datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def half(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``half`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)
load_state_dict(self, state_dict: OrderedDict[str, Tensor], strict: bool = True) inherited

Copies parameters and buffers from :attr:state_dict into this module and its descendants. If :attr:strict is True, then the keys of :attr:state_dict must exactly match the keys returned by this module's :meth:~torch.nn.Module.state_dict function.

Parameters:

Name Type Description Default
state_dict dict

a dict containing parameters and persistent buffers.

required
strict bool

whether to strictly enforce that the keys in :attr:state_dict match the keys returned by this module's :meth:~torch.nn.Module.state_dict function. Default: True

True

Returns:

Type Description
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields
  • missing_keys is a list of str containing the missing keys
    • unexpected_keys is a list of str containing the unexpected keys

!!! note If a parameter or buffer is registered as None and its corresponding key exists in :attr:state_dict, :meth:load_state_dict will raise a RuntimeError.

Source code in zamba/pytorch/transforms.py
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                    strict: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``

    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Note:
        If a parameter or buffer is registered as ``None`` and its corresponding key
        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
        ``RuntimeError``.
    """
    missing_keys: List[str] = []
    unexpected_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        # mypy isn't aware that "_metadata" exists in state_dict
        state_dict._metadata = metadata  # type: ignore[attr-defined]

    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')

    load(self)
    del load

    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))

    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)
modules(self) -> Iterator[Module] inherited

Returns an iterator over all modules in the network.

Yields:

Type Description
Module

a module in the network

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
Source code in zamba/pytorch/transforms.py
def modules(self) -> Iterator['Module']:
    r"""Returns an iterator over all modules in the network.

    Yields:
        Module: a module in the network

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)

        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)

    """
    for _, module in self.named_modules():
        yield module
named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]] inherited

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all buffer names.

''
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
(string, torch.Tensor)

Tuple containing the name and buffer

Example::

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
Source code in zamba/pytorch/transforms.py
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.

    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer

    Example::

        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())

    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
named_children(self) -> Iterator[Tuple[str, Module]] inherited

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

Type Description
(string, Module)

Tuple containing a name and child module

Example::

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Source code in zamba/pytorch/transforms.py
def named_children(self) -> Iterator[Tuple[str, 'Module']]:
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.

    Yields:
        (string, Module): Tuple containing a name and child module

    Example::

        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)

    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module
named_modules(self, memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True) inherited

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

Name Type Description Default
memo Optional[Set[Module]]

a memo to store the set of modules already added to the result

None
prefix str

a prefix that will be added to the name of the module

''
remove_duplicate bool

whether to remove the duplicated module instances in the result

True

Yields:

Type Description
(string, Module)

Tuple of name and module

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
Source code in zamba/pytorch/transforms.py
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = '', remove_duplicate: bool = True):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.

    Args:
        memo: a memo to store the set of modules already added to the result
        prefix: a prefix that will be added to the name of the module
        remove_duplicate: whether to remove the duplicated module instances in the result
        or not

    Yields:
        (string, Module): Tuple of name and module

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)

        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

    """

    if memo is None:
        memo = set()
    if self not in memo:
        if remove_duplicate:
            memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix, remove_duplicate):
                yield m
named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]] inherited

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all parameter names.

''
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
(string, Parameter)

Tuple containing the name and parameter

Example::

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
Source code in zamba/pytorch/transforms.py
def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.

    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        (string, Parameter): Tuple containing the name and parameter

    Example::

        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())

    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter] inherited

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

Name Type Description Default
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
Parameter

module parameter

Example::

>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters.

    This is typically passed to an optimizer.

    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        Parameter: module parameter

    Example::

        >>> for param in model.parameters():
        >>>     print(type(param), param.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param
register_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

This function is deprecated in favor of :meth:~torch.nn.Module.register_full_backward_hook and the behavior of this function will change in future versions.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
    the behavior of this function will change in future versions.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is True:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = False

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None inherited

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

Name Type Description Default
name string

name of the buffer. The buffer can be accessed from this module using the given name

required
tensor Tensor or None

buffer to be registered. If None, then operations that run on buffers, such as :attr:cuda, are ignored. If None, the buffer is not included in the module's :attr:state_dict.

required
persistent bool

whether the buffer is part of this module's :attr:state_dict.

True

Example::

>>> self.register_buffer('running_mean', torch.zeros(num_features))
Source code in zamba/pytorch/transforms.py
def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    r"""Adds a buffer to the module.

    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the module's state. Buffers, by
    default, are persistent and will be saved alongside parameters. This
    behavior can be changed by setting :attr:`persistent` to ``False``. The
    only difference between a persistent buffer and a non-persistent buffer
    is that the latter will not be a part of this module's
    :attr:`state_dict`.

    Buffers can be accessed as attributes using given names.

    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor or None): buffer to be registered. If ``None``, then operations
            that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
            the buffer is **not** included in the module's :attr:`state_dict`.
        persistent (bool): whether the buffer is part of this module's
            :attr:`state_dict`.

    Example::

        >>> self.register_buffer('running_mean', torch.zeros(num_features))

    """
    if persistent is False and isinstance(self, torch.jit.ScriptModule):
        raise RuntimeError("ScriptModule does not support non-persistent buffers")

    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor
        if persistent:
            self._non_persistent_buffers_set.discard(name)
        else:
            self._non_persistent_buffers_set.add(name)
register_forward_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward hook on the module.

The hook will be called every time after :func:forward has computed an output. It should have the following signature::

hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after :func:forward is called.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward hook on the module.

    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::

        hook(module, input, output) -> None or modified output

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the output. It can modify the input inplace but
    it will not have effect on forward since this is called after
    :func:`forward` is called.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle
register_forward_pre_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward pre-hook on the module.

The hook will be called every time before :func:forward is invoked. It should have the following signature::

hook(module, input) -> None or modified input

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward pre-hook on the module.

    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::

        hook(module, input) -> None or modified input

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the input. User can either return a tuple or a
    single modified value in the hook. We will wrap the value into a tuple
    if a single value is returned(unless that value is already a tuple).

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle
register_full_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature::

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The :attr:grad_input and :attr:grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of :attr:grad_input in subsequent computations. :attr:grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in :attr:grad_input and :attr:grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module's forward function.

.. warning :: Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_full_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::

        hook(module, grad_input, grad_output) -> tuple(Tensor) or None

    The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
    with respect to the inputs and outputs respectively. The hook should
    not modify its arguments, but it can optionally return a new gradient with
    respect to the input that will be used in place of :attr:`grad_input` in
    subsequent computations. :attr:`grad_input` will only correspond to the inputs given
    as positional arguments and all kwarg arguments are ignored. Entries
    in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
    arguments.

    For technical reasons, when this hook is applied to a Module, its forward function will
    receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
    of each Tensor returned by the Module's forward function.

    .. warning ::
        Modifying inputs or outputs inplace is not allowed when using backward hooks and
        will raise an error.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is False:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = True

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None inherited

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

Name Type Description Default
name string

name of the parameter. The parameter can be accessed from this module using the given name

required
param Parameter or None

parameter to be added to the module. If None, then operations that run on parameters, such as :attr:cuda, are ignored. If None, the parameter is not included in the module's :attr:state_dict.

required
Source code in zamba/pytorch/transforms.py
def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    r"""Adds a parameter to the module.

    The parameter can be accessed as an attribute using given name.

    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter or None): parameter to be added to the module. If
            ``None``, then operations that run on parameters, such as :attr:`cuda`,
            are ignored. If ``None``, the parameter is **not** included in the
            module's :attr:`state_dict`.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")

    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))

    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param
requires_grad_(self: ~T, requires_grad: bool = True) -> ~T inherited

Change if autograd should record operations on parameters in this module.

This method sets the parameters' :attr:requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See :ref:locally-disable-grad-doc for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

Name Type Description Default
requires_grad bool

whether autograd should record operations on parameters in this module. Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def requires_grad_(self: T, requires_grad: bool = True) -> T:
    r"""Change if autograd should record operations on parameters in this
    module.

    This method sets the parameters' :attr:`requires_grad` attributes
    in-place.

    This method is helpful for freezing part of the module for finetuning
    or training parts of a model individually (e.g., GAN training).

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.requires_grad_()` and several similar mechanisms that may be confused with it.

    Args:
        requires_grad (bool): whether autograd should record operations on
                              parameters in this module. Default: ``True``.

    Returns:
        Module: self
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self
set_extra_state(self, state: Any) inherited

This function is called from :func:load_state_dict to handle any extra state found within the state_dict. Implement this function and a corresponding :func:get_extra_state for your module if you need to store extra state within its state_dict.

Parameters:

Name Type Description Default
state dict

Extra state from the state_dict

required
Source code in zamba/pytorch/transforms.py
def set_extra_state(self, state: Any):
    """
    This function is called from :func:`load_state_dict` to handle any extra state
    found within the `state_dict`. Implement this function and a corresponding
    :func:`get_extra_state` for your module if you need to store extra state within its
    `state_dict`.

    Args:
        state (dict): Extra state from the `state_dict`
    """
    raise RuntimeError(
        "Reached a code path in Module.set_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
share_memory(self: ~T) -> ~T inherited

See :meth:torch.Tensor.share_memory_

Source code in zamba/pytorch/transforms.py
def share_memory(self: T) -> T:
    r"""See :meth:`torch.Tensor.share_memory_`"""
    return self._apply(lambda t: t.share_memory_())
state_dict(self, destination = None, prefix = '', keep_vars = False) inherited

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Returns:

Type Description
dict

a dictionary containing a whole state of the module

Example::

>>> module.state_dict().keys()
['bias', 'weight']
Source code in zamba/pytorch/transforms.py
def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.

    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Parameters and buffers set to ``None`` are not included.

    Returns:
        dict:
            a dictionary containing a whole state of the module

    Example::

        >>> module.state_dict().keys()
        ['bias', 'weight']

    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    self._save_to_state_dict(destination, prefix, keep_vars)
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination
to(self, *args, **kwargs) inherited

Moves and/or casts the parameters and buffers.

This can be called as

.. function:: to(device=None, dtype=None, non_blocking=False) :noindex:

.. function:: to(dtype, non_blocking=False) :noindex:

.. function:: to(tensor, non_blocking=False) :noindex:

.. function:: to(memory_format=torch.channels_last) :noindex:

Its signature is similar to :meth:torch.Tensor.to, but only accepts floating point or complex :attr:dtype\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:dtype (if given). The integral parameters and buffers will be moved :attr:device, if that is given, but with dtypes unchanged. When :attr:non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device

class:torch.device): the desired device of the parameters and buffers in this module

required
dtype

class:torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module

required
tensor torch.Tensor

Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

required
memory_format

class:torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword only argument)

required

Returns:

Type Description
Module

self

Examples::

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
Source code in zamba/pytorch/transforms.py
def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.

    This can be called as

    .. function:: to(device=None, dtype=None, non_blocking=False)
       :noindex:

    .. function:: to(dtype, non_blocking=False)
       :noindex:

    .. function:: to(tensor, non_blocking=False)
       :noindex:

    .. function:: to(memory_format=torch.channels_last)
       :noindex:

    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`\ s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.

    See below for examples.

    .. note::
        This method modifies the module in-place.

    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)

    Returns:
        Module: self

    Examples::

        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)

        >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.3741+0.j,  0.2382+0.j],
                [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
        >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
        tensor([[0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

    """

    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)

    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError('nn.Module.to only accepts floating point or complex '
                            'dtypes, but got desired dtype={}'.format(dtype))
        if dtype.is_complex:
            warnings.warn(
                "Complex modules are a new feature under active development whose design may change, "
                "and some modules might not work as expected when using complex tensors as parameters or buffers. "
                "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
                "if a complex module does not work as expected.")

    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

    return self._apply(convert)
to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T inherited

Moves the parameters and buffers to the specified device without copying storage.

Parameters:

Name Type Description Default
device

class:torch.device): The desired device of the parameters and buffers in this module.

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def to_empty(self: T, *, device: Union[str, device]) -> T:
    r"""Moves the parameters and buffers to the specified device without copying storage.

    Args:
        device (:class:`torch.device`): The desired device of the parameters
            and buffers in this module.

    Returns:
        Module: self
    """
    return self._apply(lambda t: torch.empty_like(t, device=device))
train(self: ~T, mode: bool = True) -> ~T inherited

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Parameters:

Name Type Description Default
mode bool

whether to set training mode (True) or evaluation mode (False). Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def train(self: T, mode: bool = True) -> T:
    r"""Sets the module in training mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    Args:
        mode (bool): whether to set training mode (``True``) or evaluation
                     mode (``False``). Default: ``True``.

    Returns:
        Module: self
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self
type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T inherited

Casts all parameters and buffers to :attr:dst_type.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
dst_type type or string

the desired type

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def type(self: T, dst_type: Union[dtype, str]) -> T:
    r"""Casts all parameters and buffers to :attr:`dst_type`.

    .. note::
        This method modifies the module in-place.

    Args:
        dst_type (type or string): the desired type

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))
xpu(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the XPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on XPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.xpu(device))
zero_grad(self, set_to_none: bool = False) -> None inherited

Sets gradients of all model parameters to zero. See similar function under :class:torch.optim.Optimizer for more context.

Parameters:

Name Type Description Default
set_to_none bool

instead of setting to zero, set the grads to None. See :meth:torch.optim.Optimizer.zero_grad for details.

False
Source code in zamba/pytorch/transforms.py
def zero_grad(self, set_to_none: bool = False) -> None:
    r"""Sets gradients of all model parameters to zero. See similar function
    under :class:`torch.optim.Optimizer` for more context.

    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            See :meth:`torch.optim.Optimizer.zero_grad` for details.
    """
    if getattr(self, '_is_replica', False):
        warnings.warn(
            "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
            "The parameters are copied (in a differentiable manner) from the original module. "
            "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
            "If you need gradients in your forward method, consider using autograd.grad instead.")

    for p in self.parameters():
        if p.grad is not None:
            if set_to_none:
                p.grad = None
            else:
                if p.grad.grad_fn is not None:
                    p.grad.detach_()
                else:
                    p.grad.requires_grad_(False)
                p.grad.zero_()

ConvertTHWCtoTCHW (Module)

Convert tensor from (T, H, W, C) to (T, C, H, W)

Source code in zamba/pytorch/transforms.py
class ConvertTHWCtoTCHW(torch.nn.Module):
    """Convert tensor from (T, H, W, C) to (T, C, H, W)"""

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.permute(0, 3, 1, 2)

Attributes

T_destination inherited
dump_patches: bool inherited

This allows better BC support for :meth:load_state_dict. In :meth:state_dict, the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.

If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module's _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.

Methods

__init__(self) -> None inherited special
Source code in zamba/pytorch/transforms.py
def __init__(self) -> None:
    """
    Initializes internal Module state, shared by both nn.Module and ScriptModule.
    """
    torch._C._log_api_usage_once("python.nn_module")

    self.training = True
    self._parameters: Dict[str, Optional[Parameter]] = OrderedDict()
    self._buffers: Dict[str, Optional[Tensor]] = OrderedDict()
    self._non_persistent_buffers_set: Set[str] = set()
    self._backward_hooks: Dict[int, Callable] = OrderedDict()
    self._is_full_backward_hook = None
    self._forward_hooks: Dict[int, Callable] = OrderedDict()
    self._forward_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._state_dict_hooks: Dict[int, Callable] = OrderedDict()
    self._load_state_dict_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._modules: Dict[str, Optional['Module']] = OrderedDict()
add_module(self, name: str, module: Optional[Module]) -> None inherited

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

Name Type Description Default
name string

name of the child module. The child module can be accessed from this module using the given name

required
module Module

child module to be added to the module.

required
Source code in zamba/pytorch/transforms.py
def add_module(self, name: str, module: Optional['Module']) -> None:
    r"""Adds a child module to the current module.

    The module can be accessed as an attribute using the given name.

    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\", got: {}".format(name))
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module
apply(self: ~T, fn: Callable[[Module], NoneType]) -> ~T inherited

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also :ref:nn-init-doc).

Parameters:

Name Type Description Default
fn

class:Module -> None): function to be applied to each submodule

required

Returns:

Type Description
Module

self

Example::

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Source code in zamba/pytorch/transforms.py
def apply(self: T, fn: Callable[['Module'], None]) -> T:
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`nn-init-doc`).

    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule

    Returns:
        Module: self

    Example::

        >>> @torch.no_grad()
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self
bfloat16(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to bfloat16 datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def bfloat16(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
buffers(self, recurse: bool = True) -> Iterator[torch.Tensor] inherited

Returns an iterator over module buffers.

Parameters:

Name Type Description Default
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
torch.Tensor

module buffer

Example::

>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
    r"""Returns an iterator over module buffers.

    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        torch.Tensor: module buffer

    Example::

        >>> for buf in model.buffers():
        >>>     print(type(buf), buf.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for _, buf in self.named_buffers(recurse=recurse):
        yield buf
children(self) -> Iterator[Module] inherited

Returns an iterator over immediate children modules.

Yields:

Type Description
Module

a child module

Source code in zamba/pytorch/transforms.py
def children(self) -> Iterator['Module']:
    r"""Returns an iterator over immediate children modules.

    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module
cpu(self: ~T) -> ~T inherited

Moves all model parameters and buffers to the CPU.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())
cuda(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))
double(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to double datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def double(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``double`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)
eval(self: ~T) -> ~T inherited

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

This is equivalent with :meth:self.train(False) <torch.nn.Module.train>.

See :ref:locally-disable-grad-doc for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def eval(self: T) -> T:
    r"""Sets the module in evaluation mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.eval()` and several similar mechanisms that may be confused with it.

    Returns:
        Module: self
    """
    return self.train(False)
extra_repr(self) -> str inherited

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Source code in zamba/pytorch/transforms.py
def extra_repr(self) -> str:
    r"""Set the extra representation of the module

    To print customized extra information, you should re-implement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''
float(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to float datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def float(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``float`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)
forward(self, vid: Tensor) -> Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

.. note:: Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Source code in zamba/pytorch/transforms.py
def forward(self, vid: torch.Tensor) -> torch.Tensor:
    return vid.permute(0, 3, 1, 2)
get_buffer(self, target: str) -> Tensor inherited

Returns the buffer given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.Tensor

The buffer referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not a buffer

Source code in zamba/pytorch/transforms.py
def get_buffer(self, target: str) -> "Tensor":
    """
    Returns the buffer given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the buffer
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.Tensor: The buffer referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not a
            buffer
    """
    module_path, _, buffer_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, buffer_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + buffer_name + "`")

    buffer: torch.Tensor = getattr(mod, buffer_name)

    if buffer_name not in mod._buffers:
        raise AttributeError("`" + buffer_name + "` is not a buffer")

    return buffer
get_extra_state(self) -> Any inherited

Returns any extra state to include in the module's state_dict. Implement this and a corresponding :func:set_extra_state for your module if you need to store extra state. This function is called when building the module's state_dict().

Note that extra state should be pickleable to ensure working serialization of the state_dict. We only provide provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Type Description
object

Any extra state to store in the module's state_dict

Source code in zamba/pytorch/transforms.py
def get_extra_state(self) -> Any:
    """
    Returns any extra state to include in the module's state_dict.
    Implement this and a corresponding :func:`set_extra_state` for your module
    if you need to store extra state. This function is called when building the
    module's `state_dict()`.

    Note that extra state should be pickleable to ensure working serialization
    of the state_dict. We only provide provide backwards compatibility guarantees
    for serializing Tensors; other objects may break backwards compatibility if
    their serialized pickled form changes.

    Returns:
        object: Any extra state to store in the module's state_dict
    """
    raise RuntimeError(
        "Reached a code path in Module.get_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
get_parameter(self, target: str) -> Parameter inherited

Returns the parameter given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Parameter

The Parameter referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Parameter

Source code in zamba/pytorch/transforms.py
def get_parameter(self, target: str) -> "Parameter":
    """
    Returns the parameter given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the Parameter
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Parameter: The Parameter referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Parameter``
    """
    module_path, _, param_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, param_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + param_name + "`")

    param: torch.nn.Parameter = getattr(mod, param_name)

    if not isinstance(param, torch.nn.Parameter):
        raise AttributeError("`" + param_name + "` is not an "
                             "nn.Parameter")

    return param
get_submodule(self, target: str) -> Module inherited

Returns the submodule given by target if it exists, otherwise throws an error.

For example, let's say you have an nn.Module A that looks like this:

.. code-block::text

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Module

The submodule referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Module

Source code in zamba/pytorch/transforms.py
def get_submodule(self, target: str) -> "Module":
    """
    Returns the submodule given by ``target`` if it exists,
    otherwise throws an error.

    For example, let's say you have an ``nn.Module`` ``A`` that
    looks like this:

    .. code-block::text

        A(
            (net_b): Module(
                (net_c): Module(
                    (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
                )
                (linear): Linear(in_features=100, out_features=200, bias=True)
            )
        )

    (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
    submodule ``net_b``, which itself has two submodules ``net_c``
    and ``linear``. ``net_c`` then has a submodule ``conv``.)

    To check whether or not we have the ``linear`` submodule, we
    would call ``get_submodule("net_b.linear")``. To check whether
    we have the ``conv`` submodule, we would call
    ``get_submodule("net_b.net_c.conv")``.

    The runtime of ``get_submodule`` is bounded by the degree
    of module nesting in ``target``. A query against
    ``named_modules`` achieves the same result, but it is O(N) in
    the number of transitive modules. So, for a simple check to see
    if some submodule exists, ``get_submodule`` should always be
    used.

    Args:
        target: The fully-qualified string name of the submodule
            to look for. (See above example for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Module: The submodule referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Module``
    """
    if target == "":
        return self

    atoms: List[str] = target.split(".")
    mod: torch.nn.Module = self

    for item in atoms:

        if not hasattr(mod, item):
            raise AttributeError(mod._get_name() + " has no "
                                 "attribute `" + item + "`")

        mod = getattr(mod, item)

        if not isinstance(mod, torch.nn.Module):
            raise AttributeError("`" + item + "` is not "
                                 "an nn.Module")

    return mod
half(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to half datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def half(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``half`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)
load_state_dict(self, state_dict: OrderedDict[str, Tensor], strict: bool = True) inherited

Copies parameters and buffers from :attr:state_dict into this module and its descendants. If :attr:strict is True, then the keys of :attr:state_dict must exactly match the keys returned by this module's :meth:~torch.nn.Module.state_dict function.

Parameters:

Name Type Description Default
state_dict dict

a dict containing parameters and persistent buffers.

required
strict bool

whether to strictly enforce that the keys in :attr:state_dict match the keys returned by this module's :meth:~torch.nn.Module.state_dict function. Default: True

True

Returns:

Type Description
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields
  • missing_keys is a list of str containing the missing keys
    • unexpected_keys is a list of str containing the unexpected keys

!!! note If a parameter or buffer is registered as None and its corresponding key exists in :attr:state_dict, :meth:load_state_dict will raise a RuntimeError.

Source code in zamba/pytorch/transforms.py
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                    strict: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``

    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Note:
        If a parameter or buffer is registered as ``None`` and its corresponding key
        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
        ``RuntimeError``.
    """
    missing_keys: List[str] = []
    unexpected_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        # mypy isn't aware that "_metadata" exists in state_dict
        state_dict._metadata = metadata  # type: ignore[attr-defined]

    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')

    load(self)
    del load

    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))

    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)
modules(self) -> Iterator[Module] inherited

Returns an iterator over all modules in the network.

Yields:

Type Description
Module

a module in the network

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
Source code in zamba/pytorch/transforms.py
def modules(self) -> Iterator['Module']:
    r"""Returns an iterator over all modules in the network.

    Yields:
        Module: a module in the network

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)

        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)

    """
    for _, module in self.named_modules():
        yield module
named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]] inherited

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all buffer names.

''
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
(string, torch.Tensor)

Tuple containing the name and buffer

Example::

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
Source code in zamba/pytorch/transforms.py
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.

    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer

    Example::

        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())

    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
named_children(self) -> Iterator[Tuple[str, Module]] inherited

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

Type Description
(string, Module)

Tuple containing a name and child module

Example::

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Source code in zamba/pytorch/transforms.py
def named_children(self) -> Iterator[Tuple[str, 'Module']]:
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.

    Yields:
        (string, Module): Tuple containing a name and child module

    Example::

        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)

    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module
named_modules(self, memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True) inherited

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

Name Type Description Default
memo Optional[Set[Module]]

a memo to store the set of modules already added to the result

None
prefix str

a prefix that will be added to the name of the module

''
remove_duplicate bool

whether to remove the duplicated module instances in the result

True

Yields:

Type Description
(string, Module)

Tuple of name and module

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
Source code in zamba/pytorch/transforms.py
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = '', remove_duplicate: bool = True):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.

    Args:
        memo: a memo to store the set of modules already added to the result
        prefix: a prefix that will be added to the name of the module
        remove_duplicate: whether to remove the duplicated module instances in the result
        or not

    Yields:
        (string, Module): Tuple of name and module

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)

        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

    """

    if memo is None:
        memo = set()
    if self not in memo:
        if remove_duplicate:
            memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix, remove_duplicate):
                yield m
named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]] inherited

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all parameter names.

''
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
(string, Parameter)

Tuple containing the name and parameter

Example::

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
Source code in zamba/pytorch/transforms.py
def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.

    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        (string, Parameter): Tuple containing the name and parameter

    Example::

        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())

    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter] inherited

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

Name Type Description Default
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
Parameter

module parameter

Example::

>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters.

    This is typically passed to an optimizer.

    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        Parameter: module parameter

    Example::

        >>> for param in model.parameters():
        >>>     print(type(param), param.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param
register_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

This function is deprecated in favor of :meth:~torch.nn.Module.register_full_backward_hook and the behavior of this function will change in future versions.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
    the behavior of this function will change in future versions.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is True:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = False

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None inherited

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

Name Type Description Default
name string

name of the buffer. The buffer can be accessed from this module using the given name

required
tensor Tensor or None

buffer to be registered. If None, then operations that run on buffers, such as :attr:cuda, are ignored. If None, the buffer is not included in the module's :attr:state_dict.

required
persistent bool

whether the buffer is part of this module's :attr:state_dict.

True

Example::

>>> self.register_buffer('running_mean', torch.zeros(num_features))
Source code in zamba/pytorch/transforms.py
def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    r"""Adds a buffer to the module.

    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the module's state. Buffers, by
    default, are persistent and will be saved alongside parameters. This
    behavior can be changed by setting :attr:`persistent` to ``False``. The
    only difference between a persistent buffer and a non-persistent buffer
    is that the latter will not be a part of this module's
    :attr:`state_dict`.

    Buffers can be accessed as attributes using given names.

    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor or None): buffer to be registered. If ``None``, then operations
            that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
            the buffer is **not** included in the module's :attr:`state_dict`.
        persistent (bool): whether the buffer is part of this module's
            :attr:`state_dict`.

    Example::

        >>> self.register_buffer('running_mean', torch.zeros(num_features))

    """
    if persistent is False and isinstance(self, torch.jit.ScriptModule):
        raise RuntimeError("ScriptModule does not support non-persistent buffers")

    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor
        if persistent:
            self._non_persistent_buffers_set.discard(name)
        else:
            self._non_persistent_buffers_set.add(name)
register_forward_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward hook on the module.

The hook will be called every time after :func:forward has computed an output. It should have the following signature::

hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after :func:forward is called.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward hook on the module.

    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::

        hook(module, input, output) -> None or modified output

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the output. It can modify the input inplace but
    it will not have effect on forward since this is called after
    :func:`forward` is called.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle
register_forward_pre_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward pre-hook on the module.

The hook will be called every time before :func:forward is invoked. It should have the following signature::

hook(module, input) -> None or modified input

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward pre-hook on the module.

    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::

        hook(module, input) -> None or modified input

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the input. User can either return a tuple or a
    single modified value in the hook. We will wrap the value into a tuple
    if a single value is returned(unless that value is already a tuple).

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle
register_full_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature::

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The :attr:grad_input and :attr:grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of :attr:grad_input in subsequent computations. :attr:grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in :attr:grad_input and :attr:grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module's forward function.

.. warning :: Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_full_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::

        hook(module, grad_input, grad_output) -> tuple(Tensor) or None

    The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
    with respect to the inputs and outputs respectively. The hook should
    not modify its arguments, but it can optionally return a new gradient with
    respect to the input that will be used in place of :attr:`grad_input` in
    subsequent computations. :attr:`grad_input` will only correspond to the inputs given
    as positional arguments and all kwarg arguments are ignored. Entries
    in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
    arguments.

    For technical reasons, when this hook is applied to a Module, its forward function will
    receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
    of each Tensor returned by the Module's forward function.

    .. warning ::
        Modifying inputs or outputs inplace is not allowed when using backward hooks and
        will raise an error.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is False:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = True

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None inherited

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

Name Type Description Default
name string

name of the parameter. The parameter can be accessed from this module using the given name

required
param Parameter or None

parameter to be added to the module. If None, then operations that run on parameters, such as :attr:cuda, are ignored. If None, the parameter is not included in the module's :attr:state_dict.

required
Source code in zamba/pytorch/transforms.py
def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    r"""Adds a parameter to the module.

    The parameter can be accessed as an attribute using given name.

    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter or None): parameter to be added to the module. If
            ``None``, then operations that run on parameters, such as :attr:`cuda`,
            are ignored. If ``None``, the parameter is **not** included in the
            module's :attr:`state_dict`.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")

    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))

    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param
requires_grad_(self: ~T, requires_grad: bool = True) -> ~T inherited

Change if autograd should record operations on parameters in this module.

This method sets the parameters' :attr:requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See :ref:locally-disable-grad-doc for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

Name Type Description Default
requires_grad bool

whether autograd should record operations on parameters in this module. Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def requires_grad_(self: T, requires_grad: bool = True) -> T:
    r"""Change if autograd should record operations on parameters in this
    module.

    This method sets the parameters' :attr:`requires_grad` attributes
    in-place.

    This method is helpful for freezing part of the module for finetuning
    or training parts of a model individually (e.g., GAN training).

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.requires_grad_()` and several similar mechanisms that may be confused with it.

    Args:
        requires_grad (bool): whether autograd should record operations on
                              parameters in this module. Default: ``True``.

    Returns:
        Module: self
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self
set_extra_state(self, state: Any) inherited

This function is called from :func:load_state_dict to handle any extra state found within the state_dict. Implement this function and a corresponding :func:get_extra_state for your module if you need to store extra state within its state_dict.

Parameters:

Name Type Description Default
state dict

Extra state from the state_dict

required
Source code in zamba/pytorch/transforms.py
def set_extra_state(self, state: Any):
    """
    This function is called from :func:`load_state_dict` to handle any extra state
    found within the `state_dict`. Implement this function and a corresponding
    :func:`get_extra_state` for your module if you need to store extra state within its
    `state_dict`.

    Args:
        state (dict): Extra state from the `state_dict`
    """
    raise RuntimeError(
        "Reached a code path in Module.set_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
share_memory(self: ~T) -> ~T inherited

See :meth:torch.Tensor.share_memory_

Source code in zamba/pytorch/transforms.py
def share_memory(self: T) -> T:
    r"""See :meth:`torch.Tensor.share_memory_`"""
    return self._apply(lambda t: t.share_memory_())
state_dict(self, destination = None, prefix = '', keep_vars = False) inherited

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Returns:

Type Description
dict

a dictionary containing a whole state of the module

Example::

>>> module.state_dict().keys()
['bias', 'weight']
Source code in zamba/pytorch/transforms.py
def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.

    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Parameters and buffers set to ``None`` are not included.

    Returns:
        dict:
            a dictionary containing a whole state of the module

    Example::

        >>> module.state_dict().keys()
        ['bias', 'weight']

    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    self._save_to_state_dict(destination, prefix, keep_vars)
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination
to(self, *args, **kwargs) inherited

Moves and/or casts the parameters and buffers.

This can be called as

.. function:: to(device=None, dtype=None, non_blocking=False) :noindex:

.. function:: to(dtype, non_blocking=False) :noindex:

.. function:: to(tensor, non_blocking=False) :noindex:

.. function:: to(memory_format=torch.channels_last) :noindex:

Its signature is similar to :meth:torch.Tensor.to, but only accepts floating point or complex :attr:dtype\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:dtype (if given). The integral parameters and buffers will be moved :attr:device, if that is given, but with dtypes unchanged. When :attr:non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device

class:torch.device): the desired device of the parameters and buffers in this module

required
dtype

class:torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module

required
tensor torch.Tensor

Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

required
memory_format

class:torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword only argument)

required

Returns:

Type Description
Module

self

Examples::

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
Source code in zamba/pytorch/transforms.py
def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.

    This can be called as

    .. function:: to(device=None, dtype=None, non_blocking=False)
       :noindex:

    .. function:: to(dtype, non_blocking=False)
       :noindex:

    .. function:: to(tensor, non_blocking=False)
       :noindex:

    .. function:: to(memory_format=torch.channels_last)
       :noindex:

    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`\ s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.

    See below for examples.

    .. note::
        This method modifies the module in-place.

    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)

    Returns:
        Module: self

    Examples::

        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)

        >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.3741+0.j,  0.2382+0.j],
                [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
        >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
        tensor([[0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

    """

    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)

    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError('nn.Module.to only accepts floating point or complex '
                            'dtypes, but got desired dtype={}'.format(dtype))
        if dtype.is_complex:
            warnings.warn(
                "Complex modules are a new feature under active development whose design may change, "
                "and some modules might not work as expected when using complex tensors as parameters or buffers. "
                "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
                "if a complex module does not work as expected.")

    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

    return self._apply(convert)
to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T inherited

Moves the parameters and buffers to the specified device without copying storage.

Parameters:

Name Type Description Default
device

class:torch.device): The desired device of the parameters and buffers in this module.

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def to_empty(self: T, *, device: Union[str, device]) -> T:
    r"""Moves the parameters and buffers to the specified device without copying storage.

    Args:
        device (:class:`torch.device`): The desired device of the parameters
            and buffers in this module.

    Returns:
        Module: self
    """
    return self._apply(lambda t: torch.empty_like(t, device=device))
train(self: ~T, mode: bool = True) -> ~T inherited

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Parameters:

Name Type Description Default
mode bool

whether to set training mode (True) or evaluation mode (False). Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def train(self: T, mode: bool = True) -> T:
    r"""Sets the module in training mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    Args:
        mode (bool): whether to set training mode (``True``) or evaluation
                     mode (``False``). Default: ``True``.

    Returns:
        Module: self
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self
type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T inherited

Casts all parameters and buffers to :attr:dst_type.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
dst_type type or string

the desired type

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def type(self: T, dst_type: Union[dtype, str]) -> T:
    r"""Casts all parameters and buffers to :attr:`dst_type`.

    .. note::
        This method modifies the module in-place.

    Args:
        dst_type (type or string): the desired type

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))
xpu(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the XPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on XPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.xpu(device))
zero_grad(self, set_to_none: bool = False) -> None inherited

Sets gradients of all model parameters to zero. See similar function under :class:torch.optim.Optimizer for more context.

Parameters:

Name Type Description Default
set_to_none bool

instead of setting to zero, set the grads to None. See :meth:torch.optim.Optimizer.zero_grad for details.

False
Source code in zamba/pytorch/transforms.py
def zero_grad(self, set_to_none: bool = False) -> None:
    r"""Sets gradients of all model parameters to zero. See similar function
    under :class:`torch.optim.Optimizer` for more context.

    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            See :meth:`torch.optim.Optimizer.zero_grad` for details.
    """
    if getattr(self, '_is_replica', False):
        warnings.warn(
            "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
            "The parameters are copied (in a differentiable manner) from the original module. "
            "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
            "If you need gradients in your forward method, consider using autograd.grad instead.")

    for p in self.parameters():
        if p.grad is not None:
            if set_to_none:
                p.grad = None
            else:
                if p.grad.grad_fn is not None:
                    p.grad.detach_()
                else:
                    p.grad.requires_grad_(False)
                p.grad.zero_()

PackSlowFastPathways (Module)

Creates the slow and fast pathway inputs for the slowfast model.

Source code in zamba/pytorch/transforms.py
class PackSlowFastPathways(torch.nn.Module):
    """Creates the slow and fast pathway inputs for the slowfast model."""

    def __init__(self, alpha: int = 4):
        super().__init__()
        self.alpha = alpha

    def forward(self, frames: torch.Tensor):
        fast_pathway = frames
        # Perform temporal sampling from the fast pathway.
        slow_pathway = torch.index_select(
            frames,
            1,
            torch.linspace(0, frames.shape[1] - 1, frames.shape[1] // self.alpha).long(),
        )
        frame_list = [slow_pathway, fast_pathway]
        return frame_list

Attributes

T_destination inherited
dump_patches: bool inherited

This allows better BC support for :meth:load_state_dict. In :meth:state_dict, the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.

If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module's _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.

Methods

__init__(self, alpha: int = 4) special
Source code in zamba/pytorch/transforms.py
def __init__(self, alpha: int = 4):
    super().__init__()
    self.alpha = alpha
add_module(self, name: str, module: Optional[Module]) -> None inherited

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

Name Type Description Default
name string

name of the child module. The child module can be accessed from this module using the given name

required
module Module

child module to be added to the module.

required
Source code in zamba/pytorch/transforms.py
def add_module(self, name: str, module: Optional['Module']) -> None:
    r"""Adds a child module to the current module.

    The module can be accessed as an attribute using the given name.

    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\", got: {}".format(name))
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module
apply(self: ~T, fn: Callable[[Module], NoneType]) -> ~T inherited

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also :ref:nn-init-doc).

Parameters:

Name Type Description Default
fn

class:Module -> None): function to be applied to each submodule

required

Returns:

Type Description
Module

self

Example::

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Source code in zamba/pytorch/transforms.py
def apply(self: T, fn: Callable[['Module'], None]) -> T:
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`nn-init-doc`).

    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule

    Returns:
        Module: self

    Example::

        >>> @torch.no_grad()
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self
bfloat16(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to bfloat16 datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def bfloat16(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
buffers(self, recurse: bool = True) -> Iterator[torch.Tensor] inherited

Returns an iterator over module buffers.

Parameters:

Name Type Description Default
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
torch.Tensor

module buffer

Example::

>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
    r"""Returns an iterator over module buffers.

    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        torch.Tensor: module buffer

    Example::

        >>> for buf in model.buffers():
        >>>     print(type(buf), buf.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for _, buf in self.named_buffers(recurse=recurse):
        yield buf
children(self) -> Iterator[Module] inherited

Returns an iterator over immediate children modules.

Yields:

Type Description
Module

a child module

Source code in zamba/pytorch/transforms.py
def children(self) -> Iterator['Module']:
    r"""Returns an iterator over immediate children modules.

    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module
cpu(self: ~T) -> ~T inherited

Moves all model parameters and buffers to the CPU.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())
cuda(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))
double(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to double datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def double(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``double`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)
eval(self: ~T) -> ~T inherited

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

This is equivalent with :meth:self.train(False) <torch.nn.Module.train>.

See :ref:locally-disable-grad-doc for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def eval(self: T) -> T:
    r"""Sets the module in evaluation mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.eval()` and several similar mechanisms that may be confused with it.

    Returns:
        Module: self
    """
    return self.train(False)
extra_repr(self) -> str inherited

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Source code in zamba/pytorch/transforms.py
def extra_repr(self) -> str:
    r"""Set the extra representation of the module

    To print customized extra information, you should re-implement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''
float(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to float datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def float(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``float`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)
forward(self, frames: Tensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

.. note:: Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Source code in zamba/pytorch/transforms.py
def forward(self, frames: torch.Tensor):
    fast_pathway = frames
    # Perform temporal sampling from the fast pathway.
    slow_pathway = torch.index_select(
        frames,
        1,
        torch.linspace(0, frames.shape[1] - 1, frames.shape[1] // self.alpha).long(),
    )
    frame_list = [slow_pathway, fast_pathway]
    return frame_list
get_buffer(self, target: str) -> Tensor inherited

Returns the buffer given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.Tensor

The buffer referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not a buffer

Source code in zamba/pytorch/transforms.py
def get_buffer(self, target: str) -> "Tensor":
    """
    Returns the buffer given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the buffer
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.Tensor: The buffer referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not a
            buffer
    """
    module_path, _, buffer_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, buffer_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + buffer_name + "`")

    buffer: torch.Tensor = getattr(mod, buffer_name)

    if buffer_name not in mod._buffers:
        raise AttributeError("`" + buffer_name + "` is not a buffer")

    return buffer
get_extra_state(self) -> Any inherited

Returns any extra state to include in the module's state_dict. Implement this and a corresponding :func:set_extra_state for your module if you need to store extra state. This function is called when building the module's state_dict().

Note that extra state should be pickleable to ensure working serialization of the state_dict. We only provide provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Type Description
object

Any extra state to store in the module's state_dict

Source code in zamba/pytorch/transforms.py
def get_extra_state(self) -> Any:
    """
    Returns any extra state to include in the module's state_dict.
    Implement this and a corresponding :func:`set_extra_state` for your module
    if you need to store extra state. This function is called when building the
    module's `state_dict()`.

    Note that extra state should be pickleable to ensure working serialization
    of the state_dict. We only provide provide backwards compatibility guarantees
    for serializing Tensors; other objects may break backwards compatibility if
    their serialized pickled form changes.

    Returns:
        object: Any extra state to store in the module's state_dict
    """
    raise RuntimeError(
        "Reached a code path in Module.get_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
get_parameter(self, target: str) -> Parameter inherited

Returns the parameter given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Parameter

The Parameter referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Parameter

Source code in zamba/pytorch/transforms.py
def get_parameter(self, target: str) -> "Parameter":
    """
    Returns the parameter given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the Parameter
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Parameter: The Parameter referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Parameter``
    """
    module_path, _, param_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, param_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + param_name + "`")

    param: torch.nn.Parameter = getattr(mod, param_name)

    if not isinstance(param, torch.nn.Parameter):
        raise AttributeError("`" + param_name + "` is not an "
                             "nn.Parameter")

    return param
get_submodule(self, target: str) -> Module inherited

Returns the submodule given by target if it exists, otherwise throws an error.

For example, let's say you have an nn.Module A that looks like this:

.. code-block::text

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Module

The submodule referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Module

Source code in zamba/pytorch/transforms.py
def get_submodule(self, target: str) -> "Module":
    """
    Returns the submodule given by ``target`` if it exists,
    otherwise throws an error.

    For example, let's say you have an ``nn.Module`` ``A`` that
    looks like this:

    .. code-block::text

        A(
            (net_b): Module(
                (net_c): Module(
                    (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
                )
                (linear): Linear(in_features=100, out_features=200, bias=True)
            )
        )

    (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
    submodule ``net_b``, which itself has two submodules ``net_c``
    and ``linear``. ``net_c`` then has a submodule ``conv``.)

    To check whether or not we have the ``linear`` submodule, we
    would call ``get_submodule("net_b.linear")``. To check whether
    we have the ``conv`` submodule, we would call
    ``get_submodule("net_b.net_c.conv")``.

    The runtime of ``get_submodule`` is bounded by the degree
    of module nesting in ``target``. A query against
    ``named_modules`` achieves the same result, but it is O(N) in
    the number of transitive modules. So, for a simple check to see
    if some submodule exists, ``get_submodule`` should always be
    used.

    Args:
        target: The fully-qualified string name of the submodule
            to look for. (See above example for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Module: The submodule referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Module``
    """
    if target == "":
        return self

    atoms: List[str] = target.split(".")
    mod: torch.nn.Module = self

    for item in atoms:

        if not hasattr(mod, item):
            raise AttributeError(mod._get_name() + " has no "
                                 "attribute `" + item + "`")

        mod = getattr(mod, item)

        if not isinstance(mod, torch.nn.Module):
            raise AttributeError("`" + item + "` is not "
                                 "an nn.Module")

    return mod
half(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to half datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def half(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``half`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)
load_state_dict(self, state_dict: OrderedDict[str, Tensor], strict: bool = True) inherited

Copies parameters and buffers from :attr:state_dict into this module and its descendants. If :attr:strict is True, then the keys of :attr:state_dict must exactly match the keys returned by this module's :meth:~torch.nn.Module.state_dict function.

Parameters:

Name Type Description Default
state_dict dict

a dict containing parameters and persistent buffers.

required
strict bool

whether to strictly enforce that the keys in :attr:state_dict match the keys returned by this module's :meth:~torch.nn.Module.state_dict function. Default: True

True

Returns:

Type Description
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields
  • missing_keys is a list of str containing the missing keys
    • unexpected_keys is a list of str containing the unexpected keys

!!! note If a parameter or buffer is registered as None and its corresponding key exists in :attr:state_dict, :meth:load_state_dict will raise a RuntimeError.

Source code in zamba/pytorch/transforms.py
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                    strict: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``

    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Note:
        If a parameter or buffer is registered as ``None`` and its corresponding key
        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
        ``RuntimeError``.
    """
    missing_keys: List[str] = []
    unexpected_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        # mypy isn't aware that "_metadata" exists in state_dict
        state_dict._metadata = metadata  # type: ignore[attr-defined]

    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')

    load(self)
    del load

    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))

    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)
modules(self) -> Iterator[Module] inherited

Returns an iterator over all modules in the network.

Yields:

Type Description
Module

a module in the network

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
Source code in zamba/pytorch/transforms.py
def modules(self) -> Iterator['Module']:
    r"""Returns an iterator over all modules in the network.

    Yields:
        Module: a module in the network

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)

        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)

    """
    for _, module in self.named_modules():
        yield module
named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]] inherited

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all buffer names.

''
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
(string, torch.Tensor)

Tuple containing the name and buffer

Example::

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
Source code in zamba/pytorch/transforms.py
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.

    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer

    Example::

        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())

    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
named_children(self) -> Iterator[Tuple[str, Module]] inherited

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

Type Description
(string, Module)

Tuple containing a name and child module

Example::

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Source code in zamba/pytorch/transforms.py
def named_children(self) -> Iterator[Tuple[str, 'Module']]:
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.

    Yields:
        (string, Module): Tuple containing a name and child module

    Example::

        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)

    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module
named_modules(self, memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True) inherited

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

Name Type Description Default
memo Optional[Set[Module]]

a memo to store the set of modules already added to the result

None
prefix str

a prefix that will be added to the name of the module

''
remove_duplicate bool

whether to remove the duplicated module instances in the result

True

Yields:

Type Description
(string, Module)

Tuple of name and module

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
Source code in zamba/pytorch/transforms.py
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = '', remove_duplicate: bool = True):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.

    Args:
        memo: a memo to store the set of modules already added to the result
        prefix: a prefix that will be added to the name of the module
        remove_duplicate: whether to remove the duplicated module instances in the result
        or not

    Yields:
        (string, Module): Tuple of name and module

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)

        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

    """

    if memo is None:
        memo = set()
    if self not in memo:
        if remove_duplicate:
            memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix, remove_duplicate):
                yield m
named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]] inherited

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all parameter names.

''
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
(string, Parameter)

Tuple containing the name and parameter

Example::

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
Source code in zamba/pytorch/transforms.py
def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.

    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        (string, Parameter): Tuple containing the name and parameter

    Example::

        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())

    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter] inherited

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

Name Type Description Default
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
Parameter

module parameter

Example::

>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters.

    This is typically passed to an optimizer.

    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        Parameter: module parameter

    Example::

        >>> for param in model.parameters():
        >>>     print(type(param), param.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param
register_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

This function is deprecated in favor of :meth:~torch.nn.Module.register_full_backward_hook and the behavior of this function will change in future versions.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
    the behavior of this function will change in future versions.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is True:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = False

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None inherited

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

Name Type Description Default
name string

name of the buffer. The buffer can be accessed from this module using the given name

required
tensor Tensor or None

buffer to be registered. If None, then operations that run on buffers, such as :attr:cuda, are ignored. If None, the buffer is not included in the module's :attr:state_dict.

required
persistent bool

whether the buffer is part of this module's :attr:state_dict.

True

Example::

>>> self.register_buffer('running_mean', torch.zeros(num_features))
Source code in zamba/pytorch/transforms.py
def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    r"""Adds a buffer to the module.

    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the module's state. Buffers, by
    default, are persistent and will be saved alongside parameters. This
    behavior can be changed by setting :attr:`persistent` to ``False``. The
    only difference between a persistent buffer and a non-persistent buffer
    is that the latter will not be a part of this module's
    :attr:`state_dict`.

    Buffers can be accessed as attributes using given names.

    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor or None): buffer to be registered. If ``None``, then operations
            that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
            the buffer is **not** included in the module's :attr:`state_dict`.
        persistent (bool): whether the buffer is part of this module's
            :attr:`state_dict`.

    Example::

        >>> self.register_buffer('running_mean', torch.zeros(num_features))

    """
    if persistent is False and isinstance(self, torch.jit.ScriptModule):
        raise RuntimeError("ScriptModule does not support non-persistent buffers")

    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor
        if persistent:
            self._non_persistent_buffers_set.discard(name)
        else:
            self._non_persistent_buffers_set.add(name)
register_forward_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward hook on the module.

The hook will be called every time after :func:forward has computed an output. It should have the following signature::

hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after :func:forward is called.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward hook on the module.

    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::

        hook(module, input, output) -> None or modified output

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the output. It can modify the input inplace but
    it will not have effect on forward since this is called after
    :func:`forward` is called.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle
register_forward_pre_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward pre-hook on the module.

The hook will be called every time before :func:forward is invoked. It should have the following signature::

hook(module, input) -> None or modified input

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward pre-hook on the module.

    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::

        hook(module, input) -> None or modified input

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the input. User can either return a tuple or a
    single modified value in the hook. We will wrap the value into a tuple
    if a single value is returned(unless that value is already a tuple).

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle
register_full_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature::

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The :attr:grad_input and :attr:grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of :attr:grad_input in subsequent computations. :attr:grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in :attr:grad_input and :attr:grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module's forward function.

.. warning :: Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_full_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::

        hook(module, grad_input, grad_output) -> tuple(Tensor) or None

    The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
    with respect to the inputs and outputs respectively. The hook should
    not modify its arguments, but it can optionally return a new gradient with
    respect to the input that will be used in place of :attr:`grad_input` in
    subsequent computations. :attr:`grad_input` will only correspond to the inputs given
    as positional arguments and all kwarg arguments are ignored. Entries
    in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
    arguments.

    For technical reasons, when this hook is applied to a Module, its forward function will
    receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
    of each Tensor returned by the Module's forward function.

    .. warning ::
        Modifying inputs or outputs inplace is not allowed when using backward hooks and
        will raise an error.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is False:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = True

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None inherited

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

Name Type Description Default
name string

name of the parameter. The parameter can be accessed from this module using the given name

required
param Parameter or None

parameter to be added to the module. If None, then operations that run on parameters, such as :attr:cuda, are ignored. If None, the parameter is not included in the module's :attr:state_dict.

required
Source code in zamba/pytorch/transforms.py
def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    r"""Adds a parameter to the module.

    The parameter can be accessed as an attribute using given name.

    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter or None): parameter to be added to the module. If
            ``None``, then operations that run on parameters, such as :attr:`cuda`,
            are ignored. If ``None``, the parameter is **not** included in the
            module's :attr:`state_dict`.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")

    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))

    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param
requires_grad_(self: ~T, requires_grad: bool = True) -> ~T inherited

Change if autograd should record operations on parameters in this module.

This method sets the parameters' :attr:requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See :ref:locally-disable-grad-doc for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

Name Type Description Default
requires_grad bool

whether autograd should record operations on parameters in this module. Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def requires_grad_(self: T, requires_grad: bool = True) -> T:
    r"""Change if autograd should record operations on parameters in this
    module.

    This method sets the parameters' :attr:`requires_grad` attributes
    in-place.

    This method is helpful for freezing part of the module for finetuning
    or training parts of a model individually (e.g., GAN training).

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.requires_grad_()` and several similar mechanisms that may be confused with it.

    Args:
        requires_grad (bool): whether autograd should record operations on
                              parameters in this module. Default: ``True``.

    Returns:
        Module: self
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self
set_extra_state(self, state: Any) inherited

This function is called from :func:load_state_dict to handle any extra state found within the state_dict. Implement this function and a corresponding :func:get_extra_state for your module if you need to store extra state within its state_dict.

Parameters:

Name Type Description Default
state dict

Extra state from the state_dict

required
Source code in zamba/pytorch/transforms.py
def set_extra_state(self, state: Any):
    """
    This function is called from :func:`load_state_dict` to handle any extra state
    found within the `state_dict`. Implement this function and a corresponding
    :func:`get_extra_state` for your module if you need to store extra state within its
    `state_dict`.

    Args:
        state (dict): Extra state from the `state_dict`
    """
    raise RuntimeError(
        "Reached a code path in Module.set_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
share_memory(self: ~T) -> ~T inherited

See :meth:torch.Tensor.share_memory_

Source code in zamba/pytorch/transforms.py
def share_memory(self: T) -> T:
    r"""See :meth:`torch.Tensor.share_memory_`"""
    return self._apply(lambda t: t.share_memory_())
state_dict(self, destination = None, prefix = '', keep_vars = False) inherited

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Returns:

Type Description
dict

a dictionary containing a whole state of the module

Example::

>>> module.state_dict().keys()
['bias', 'weight']
Source code in zamba/pytorch/transforms.py
def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.

    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Parameters and buffers set to ``None`` are not included.

    Returns:
        dict:
            a dictionary containing a whole state of the module

    Example::

        >>> module.state_dict().keys()
        ['bias', 'weight']

    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    self._save_to_state_dict(destination, prefix, keep_vars)
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination
to(self, *args, **kwargs) inherited

Moves and/or casts the parameters and buffers.

This can be called as

.. function:: to(device=None, dtype=None, non_blocking=False) :noindex:

.. function:: to(dtype, non_blocking=False) :noindex:

.. function:: to(tensor, non_blocking=False) :noindex:

.. function:: to(memory_format=torch.channels_last) :noindex:

Its signature is similar to :meth:torch.Tensor.to, but only accepts floating point or complex :attr:dtype\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:dtype (if given). The integral parameters and buffers will be moved :attr:device, if that is given, but with dtypes unchanged. When :attr:non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device

class:torch.device): the desired device of the parameters and buffers in this module

required
dtype

class:torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module

required
tensor torch.Tensor

Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

required
memory_format

class:torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword only argument)

required

Returns:

Type Description
Module

self

Examples::

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
Source code in zamba/pytorch/transforms.py
def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.

    This can be called as

    .. function:: to(device=None, dtype=None, non_blocking=False)
       :noindex:

    .. function:: to(dtype, non_blocking=False)
       :noindex:

    .. function:: to(tensor, non_blocking=False)
       :noindex:

    .. function:: to(memory_format=torch.channels_last)
       :noindex:

    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`\ s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.

    See below for examples.

    .. note::
        This method modifies the module in-place.

    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)

    Returns:
        Module: self

    Examples::

        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)

        >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.3741+0.j,  0.2382+0.j],
                [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
        >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
        tensor([[0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

    """

    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)

    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError('nn.Module.to only accepts floating point or complex '
                            'dtypes, but got desired dtype={}'.format(dtype))
        if dtype.is_complex:
            warnings.warn(
                "Complex modules are a new feature under active development whose design may change, "
                "and some modules might not work as expected when using complex tensors as parameters or buffers. "
                "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
                "if a complex module does not work as expected.")

    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

    return self._apply(convert)
to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T inherited

Moves the parameters and buffers to the specified device without copying storage.

Parameters:

Name Type Description Default
device

class:torch.device): The desired device of the parameters and buffers in this module.

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def to_empty(self: T, *, device: Union[str, device]) -> T:
    r"""Moves the parameters and buffers to the specified device without copying storage.

    Args:
        device (:class:`torch.device`): The desired device of the parameters
            and buffers in this module.

    Returns:
        Module: self
    """
    return self._apply(lambda t: torch.empty_like(t, device=device))
train(self: ~T, mode: bool = True) -> ~T inherited

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Parameters:

Name Type Description Default
mode bool

whether to set training mode (True) or evaluation mode (False). Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def train(self: T, mode: bool = True) -> T:
    r"""Sets the module in training mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    Args:
        mode (bool): whether to set training mode (``True``) or evaluation
                     mode (``False``). Default: ``True``.

    Returns:
        Module: self
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self
type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T inherited

Casts all parameters and buffers to :attr:dst_type.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
dst_type type or string

the desired type

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def type(self: T, dst_type: Union[dtype, str]) -> T:
    r"""Casts all parameters and buffers to :attr:`dst_type`.

    .. note::
        This method modifies the module in-place.

    Args:
        dst_type (type or string): the desired type

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))
xpu(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the XPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on XPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.xpu(device))
zero_grad(self, set_to_none: bool = False) -> None inherited

Sets gradients of all model parameters to zero. See similar function under :class:torch.optim.Optimizer for more context.

Parameters:

Name Type Description Default
set_to_none bool

instead of setting to zero, set the grads to None. See :meth:torch.optim.Optimizer.zero_grad for details.

False
Source code in zamba/pytorch/transforms.py
def zero_grad(self, set_to_none: bool = False) -> None:
    r"""Sets gradients of all model parameters to zero. See similar function
    under :class:`torch.optim.Optimizer` for more context.

    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            See :meth:`torch.optim.Optimizer.zero_grad` for details.
    """
    if getattr(self, '_is_replica', False):
        warnings.warn(
            "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
            "The parameters are copied (in a differentiable manner) from the original module. "
            "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
            "If you need gradients in your forward method, consider using autograd.grad instead.")

    for p in self.parameters():
        if p.grad is not None:
            if set_to_none:
                p.grad = None
            else:
                if p.grad.grad_fn is not None:
                    p.grad.detach_()
                else:
                    p.grad.requires_grad_(False)
                p.grad.zero_()

PadDimensions (Module)

Pads a tensor to ensure a fixed output dimension for a give axis.

Attributes:

Name Type Description
dimension_sizes

A tuple of int or None the same length as the number of dimensions in the input tensor. If int, pad that dimension to at least that size. If None, do not pad.

Source code in zamba/pytorch/transforms.py
class PadDimensions(torch.nn.Module):
    """Pads a tensor to ensure a fixed output dimension for a give axis.

    Attributes:
        dimension_sizes: A tuple of int or None the same length as the number of dimensions in the
            input tensor. If int, pad that dimension to at least that size. If None, do not pad.
    """

    def __init__(self, dimension_sizes: Tuple[Optional[int]]):
        super().__init__()
        self.dimension_sizes = dimension_sizes

    @staticmethod
    def compute_left_and_right_pad(original_size: int, padded_size: int) -> Tuple[int, int]:
        """Computes left and right pad size.

        Args:
            original_size (list, int): The original tensor size
            padded_size (list, int): The desired tensor size

        Returns:
           Tuple[int]: Pad size for right and left. For odd padding size, the right = left + 1
        """
        if original_size >= padded_size:
            return 0, 0
        pad = padded_size - original_size
        quotient, remainder = divmod(pad, 2)
        return quotient, quotient + remainder

    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        padding = tuple(
            itertools.chain.from_iterable(
                (0, 0)
                if padded_size is None
                else self.compute_left_and_right_pad(original_size, padded_size)
                for original_size, padded_size in zip(vid.shape, self.dimension_sizes)
            )
        )
        return torch.nn.functional.pad(vid, padding[::-1])

Attributes

T_destination inherited
dump_patches: bool inherited

This allows better BC support for :meth:load_state_dict. In :meth:state_dict, the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.

If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module's _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.

Methods

__init__(self, dimension_sizes: Tuple[Optional[int]]) special
Source code in zamba/pytorch/transforms.py
def __init__(self, dimension_sizes: Tuple[Optional[int]]):
    super().__init__()
    self.dimension_sizes = dimension_sizes
add_module(self, name: str, module: Optional[Module]) -> None inherited

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

Name Type Description Default
name string

name of the child module. The child module can be accessed from this module using the given name

required
module Module

child module to be added to the module.

required
Source code in zamba/pytorch/transforms.py
def add_module(self, name: str, module: Optional['Module']) -> None:
    r"""Adds a child module to the current module.

    The module can be accessed as an attribute using the given name.

    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\", got: {}".format(name))
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module
apply(self: ~T, fn: Callable[[Module], NoneType]) -> ~T inherited

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also :ref:nn-init-doc).

Parameters:

Name Type Description Default
fn

class:Module -> None): function to be applied to each submodule

required

Returns:

Type Description
Module

self

Example::

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Source code in zamba/pytorch/transforms.py
def apply(self: T, fn: Callable[['Module'], None]) -> T:
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`nn-init-doc`).

    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule

    Returns:
        Module: self

    Example::

        >>> @torch.no_grad()
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self
bfloat16(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to bfloat16 datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def bfloat16(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
buffers(self, recurse: bool = True) -> Iterator[torch.Tensor] inherited

Returns an iterator over module buffers.

Parameters:

Name Type Description Default
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
torch.Tensor

module buffer

Example::

>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
    r"""Returns an iterator over module buffers.

    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        torch.Tensor: module buffer

    Example::

        >>> for buf in model.buffers():
        >>>     print(type(buf), buf.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for _, buf in self.named_buffers(recurse=recurse):
        yield buf
children(self) -> Iterator[Module] inherited

Returns an iterator over immediate children modules.

Yields:

Type Description
Module

a child module

Source code in zamba/pytorch/transforms.py
def children(self) -> Iterator['Module']:
    r"""Returns an iterator over immediate children modules.

    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module
compute_left_and_right_pad(original_size: int, padded_size: int) -> Tuple[int, int] staticmethod

Computes left and right pad size.

Parameters:

Name Type Description Default
original_size list, int

The original tensor size

required
padded_size list, int

The desired tensor size

required

Returns:

Type Description
Tuple[int]

Pad size for right and left. For odd padding size, the right = left + 1

Source code in zamba/pytorch/transforms.py
@staticmethod
def compute_left_and_right_pad(original_size: int, padded_size: int) -> Tuple[int, int]:
    """Computes left and right pad size.

    Args:
        original_size (list, int): The original tensor size
        padded_size (list, int): The desired tensor size

    Returns:
       Tuple[int]: Pad size for right and left. For odd padding size, the right = left + 1
    """
    if original_size >= padded_size:
        return 0, 0
    pad = padded_size - original_size
    quotient, remainder = divmod(pad, 2)
    return quotient, quotient + remainder
cpu(self: ~T) -> ~T inherited

Moves all model parameters and buffers to the CPU.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())
cuda(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))
double(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to double datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def double(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``double`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)
eval(self: ~T) -> ~T inherited

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

This is equivalent with :meth:self.train(False) <torch.nn.Module.train>.

See :ref:locally-disable-grad-doc for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def eval(self: T) -> T:
    r"""Sets the module in evaluation mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.eval()` and several similar mechanisms that may be confused with it.

    Returns:
        Module: self
    """
    return self.train(False)
extra_repr(self) -> str inherited

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Source code in zamba/pytorch/transforms.py
def extra_repr(self) -> str:
    r"""Set the extra representation of the module

    To print customized extra information, you should re-implement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''
float(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to float datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def float(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``float`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)
forward(self, vid: Tensor) -> Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

.. note:: Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Source code in zamba/pytorch/transforms.py
def forward(self, vid: torch.Tensor) -> torch.Tensor:
    padding = tuple(
        itertools.chain.from_iterable(
            (0, 0)
            if padded_size is None
            else self.compute_left_and_right_pad(original_size, padded_size)
            for original_size, padded_size in zip(vid.shape, self.dimension_sizes)
        )
    )
    return torch.nn.functional.pad(vid, padding[::-1])
get_buffer(self, target: str) -> Tensor inherited

Returns the buffer given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.Tensor

The buffer referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not a buffer

Source code in zamba/pytorch/transforms.py
def get_buffer(self, target: str) -> "Tensor":
    """
    Returns the buffer given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the buffer
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.Tensor: The buffer referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not a
            buffer
    """
    module_path, _, buffer_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, buffer_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + buffer_name + "`")

    buffer: torch.Tensor = getattr(mod, buffer_name)

    if buffer_name not in mod._buffers:
        raise AttributeError("`" + buffer_name + "` is not a buffer")

    return buffer
get_extra_state(self) -> Any inherited

Returns any extra state to include in the module's state_dict. Implement this and a corresponding :func:set_extra_state for your module if you need to store extra state. This function is called when building the module's state_dict().

Note that extra state should be pickleable to ensure working serialization of the state_dict. We only provide provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Type Description
object

Any extra state to store in the module's state_dict

Source code in zamba/pytorch/transforms.py
def get_extra_state(self) -> Any:
    """
    Returns any extra state to include in the module's state_dict.
    Implement this and a corresponding :func:`set_extra_state` for your module
    if you need to store extra state. This function is called when building the
    module's `state_dict()`.

    Note that extra state should be pickleable to ensure working serialization
    of the state_dict. We only provide provide backwards compatibility guarantees
    for serializing Tensors; other objects may break backwards compatibility if
    their serialized pickled form changes.

    Returns:
        object: Any extra state to store in the module's state_dict
    """
    raise RuntimeError(
        "Reached a code path in Module.get_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
get_parameter(self, target: str) -> Parameter inherited

Returns the parameter given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Parameter

The Parameter referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Parameter

Source code in zamba/pytorch/transforms.py
def get_parameter(self, target: str) -> "Parameter":
    """
    Returns the parameter given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the Parameter
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Parameter: The Parameter referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Parameter``
    """
    module_path, _, param_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, param_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + param_name + "`")

    param: torch.nn.Parameter = getattr(mod, param_name)

    if not isinstance(param, torch.nn.Parameter):
        raise AttributeError("`" + param_name + "` is not an "
                             "nn.Parameter")

    return param
get_submodule(self, target: str) -> Module inherited

Returns the submodule given by target if it exists, otherwise throws an error.

For example, let's say you have an nn.Module A that looks like this:

.. code-block::text

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Module

The submodule referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Module

Source code in zamba/pytorch/transforms.py
def get_submodule(self, target: str) -> "Module":
    """
    Returns the submodule given by ``target`` if it exists,
    otherwise throws an error.

    For example, let's say you have an ``nn.Module`` ``A`` that
    looks like this:

    .. code-block::text

        A(
            (net_b): Module(
                (net_c): Module(
                    (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
                )
                (linear): Linear(in_features=100, out_features=200, bias=True)
            )
        )

    (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
    submodule ``net_b``, which itself has two submodules ``net_c``
    and ``linear``. ``net_c`` then has a submodule ``conv``.)

    To check whether or not we have the ``linear`` submodule, we
    would call ``get_submodule("net_b.linear")``. To check whether
    we have the ``conv`` submodule, we would call
    ``get_submodule("net_b.net_c.conv")``.

    The runtime of ``get_submodule`` is bounded by the degree
    of module nesting in ``target``. A query against
    ``named_modules`` achieves the same result, but it is O(N) in
    the number of transitive modules. So, for a simple check to see
    if some submodule exists, ``get_submodule`` should always be
    used.

    Args:
        target: The fully-qualified string name of the submodule
            to look for. (See above example for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Module: The submodule referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Module``
    """
    if target == "":
        return self

    atoms: List[str] = target.split(".")
    mod: torch.nn.Module = self

    for item in atoms:

        if not hasattr(mod, item):
            raise AttributeError(mod._get_name() + " has no "
                                 "attribute `" + item + "`")

        mod = getattr(mod, item)

        if not isinstance(mod, torch.nn.Module):
            raise AttributeError("`" + item + "` is not "
                                 "an nn.Module")

    return mod
half(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to half datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def half(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``half`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)
load_state_dict(self, state_dict: OrderedDict[str, Tensor], strict: bool = True) inherited

Copies parameters and buffers from :attr:state_dict into this module and its descendants. If :attr:strict is True, then the keys of :attr:state_dict must exactly match the keys returned by this module's :meth:~torch.nn.Module.state_dict function.

Parameters:

Name Type Description Default
state_dict dict

a dict containing parameters and persistent buffers.

required
strict bool

whether to strictly enforce that the keys in :attr:state_dict match the keys returned by this module's :meth:~torch.nn.Module.state_dict function. Default: True

True

Returns:

Type Description
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields
  • missing_keys is a list of str containing the missing keys
    • unexpected_keys is a list of str containing the unexpected keys

!!! note If a parameter or buffer is registered as None and its corresponding key exists in :attr:state_dict, :meth:load_state_dict will raise a RuntimeError.

Source code in zamba/pytorch/transforms.py
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                    strict: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``

    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Note:
        If a parameter or buffer is registered as ``None`` and its corresponding key
        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
        ``RuntimeError``.
    """
    missing_keys: List[str] = []
    unexpected_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        # mypy isn't aware that "_metadata" exists in state_dict
        state_dict._metadata = metadata  # type: ignore[attr-defined]

    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')

    load(self)
    del load

    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))

    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)
modules(self) -> Iterator[Module] inherited

Returns an iterator over all modules in the network.

Yields:

Type Description
Module

a module in the network

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
Source code in zamba/pytorch/transforms.py
def modules(self) -> Iterator['Module']:
    r"""Returns an iterator over all modules in the network.

    Yields:
        Module: a module in the network

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)

        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)

    """
    for _, module in self.named_modules():
        yield module
named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]] inherited

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all buffer names.

''
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
(string, torch.Tensor)

Tuple containing the name and buffer

Example::

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
Source code in zamba/pytorch/transforms.py
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.

    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer

    Example::

        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())

    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
named_children(self) -> Iterator[Tuple[str, Module]] inherited

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

Type Description
(string, Module)

Tuple containing a name and child module

Example::

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Source code in zamba/pytorch/transforms.py
def named_children(self) -> Iterator[Tuple[str, 'Module']]:
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.

    Yields:
        (string, Module): Tuple containing a name and child module

    Example::

        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)

    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module
named_modules(self, memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True) inherited

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

Name Type Description Default
memo Optional[Set[Module]]

a memo to store the set of modules already added to the result

None
prefix str

a prefix that will be added to the name of the module

''
remove_duplicate bool

whether to remove the duplicated module instances in the result

True

Yields:

Type Description
(string, Module)

Tuple of name and module

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
Source code in zamba/pytorch/transforms.py
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = '', remove_duplicate: bool = True):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.

    Args:
        memo: a memo to store the set of modules already added to the result
        prefix: a prefix that will be added to the name of the module
        remove_duplicate: whether to remove the duplicated module instances in the result
        or not

    Yields:
        (string, Module): Tuple of name and module

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)

        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

    """

    if memo is None:
        memo = set()
    if self not in memo:
        if remove_duplicate:
            memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix, remove_duplicate):
                yield m
named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]] inherited

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all parameter names.

''
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
(string, Parameter)

Tuple containing the name and parameter

Example::

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
Source code in zamba/pytorch/transforms.py
def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.

    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        (string, Parameter): Tuple containing the name and parameter

    Example::

        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())

    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter] inherited

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

Name Type Description Default
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
Parameter

module parameter

Example::

>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters.

    This is typically passed to an optimizer.

    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        Parameter: module parameter

    Example::

        >>> for param in model.parameters():
        >>>     print(type(param), param.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param
register_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

This function is deprecated in favor of :meth:~torch.nn.Module.register_full_backward_hook and the behavior of this function will change in future versions.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
    the behavior of this function will change in future versions.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is True:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = False

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None inherited

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

Name Type Description Default
name string

name of the buffer. The buffer can be accessed from this module using the given name

required
tensor Tensor or None

buffer to be registered. If None, then operations that run on buffers, such as :attr:cuda, are ignored. If None, the buffer is not included in the module's :attr:state_dict.

required
persistent bool

whether the buffer is part of this module's :attr:state_dict.

True

Example::

>>> self.register_buffer('running_mean', torch.zeros(num_features))
Source code in zamba/pytorch/transforms.py
def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    r"""Adds a buffer to the module.

    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the module's state. Buffers, by
    default, are persistent and will be saved alongside parameters. This
    behavior can be changed by setting :attr:`persistent` to ``False``. The
    only difference between a persistent buffer and a non-persistent buffer
    is that the latter will not be a part of this module's
    :attr:`state_dict`.

    Buffers can be accessed as attributes using given names.

    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor or None): buffer to be registered. If ``None``, then operations
            that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
            the buffer is **not** included in the module's :attr:`state_dict`.
        persistent (bool): whether the buffer is part of this module's
            :attr:`state_dict`.

    Example::

        >>> self.register_buffer('running_mean', torch.zeros(num_features))

    """
    if persistent is False and isinstance(self, torch.jit.ScriptModule):
        raise RuntimeError("ScriptModule does not support non-persistent buffers")

    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor
        if persistent:
            self._non_persistent_buffers_set.discard(name)
        else:
            self._non_persistent_buffers_set.add(name)
register_forward_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward hook on the module.

The hook will be called every time after :func:forward has computed an output. It should have the following signature::

hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after :func:forward is called.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward hook on the module.

    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::

        hook(module, input, output) -> None or modified output

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the output. It can modify the input inplace but
    it will not have effect on forward since this is called after
    :func:`forward` is called.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle
register_forward_pre_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward pre-hook on the module.

The hook will be called every time before :func:forward is invoked. It should have the following signature::

hook(module, input) -> None or modified input

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward pre-hook on the module.

    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::

        hook(module, input) -> None or modified input

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the input. User can either return a tuple or a
    single modified value in the hook. We will wrap the value into a tuple
    if a single value is returned(unless that value is already a tuple).

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle
register_full_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature::

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The :attr:grad_input and :attr:grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of :attr:grad_input in subsequent computations. :attr:grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in :attr:grad_input and :attr:grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module's forward function.

.. warning :: Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_full_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::

        hook(module, grad_input, grad_output) -> tuple(Tensor) or None

    The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
    with respect to the inputs and outputs respectively. The hook should
    not modify its arguments, but it can optionally return a new gradient with
    respect to the input that will be used in place of :attr:`grad_input` in
    subsequent computations. :attr:`grad_input` will only correspond to the inputs given
    as positional arguments and all kwarg arguments are ignored. Entries
    in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
    arguments.

    For technical reasons, when this hook is applied to a Module, its forward function will
    receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
    of each Tensor returned by the Module's forward function.

    .. warning ::
        Modifying inputs or outputs inplace is not allowed when using backward hooks and
        will raise an error.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is False:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = True

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None inherited

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

Name Type Description Default
name string

name of the parameter. The parameter can be accessed from this module using the given name

required
param Parameter or None

parameter to be added to the module. If None, then operations that run on parameters, such as :attr:cuda, are ignored. If None, the parameter is not included in the module's :attr:state_dict.

required
Source code in zamba/pytorch/transforms.py
def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    r"""Adds a parameter to the module.

    The parameter can be accessed as an attribute using given name.

    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter or None): parameter to be added to the module. If
            ``None``, then operations that run on parameters, such as :attr:`cuda`,
            are ignored. If ``None``, the parameter is **not** included in the
            module's :attr:`state_dict`.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")

    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))

    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param
requires_grad_(self: ~T, requires_grad: bool = True) -> ~T inherited

Change if autograd should record operations on parameters in this module.

This method sets the parameters' :attr:requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See :ref:locally-disable-grad-doc for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

Name Type Description Default
requires_grad bool

whether autograd should record operations on parameters in this module. Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def requires_grad_(self: T, requires_grad: bool = True) -> T:
    r"""Change if autograd should record operations on parameters in this
    module.

    This method sets the parameters' :attr:`requires_grad` attributes
    in-place.

    This method is helpful for freezing part of the module for finetuning
    or training parts of a model individually (e.g., GAN training).

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.requires_grad_()` and several similar mechanisms that may be confused with it.

    Args:
        requires_grad (bool): whether autograd should record operations on
                              parameters in this module. Default: ``True``.

    Returns:
        Module: self
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self
set_extra_state(self, state: Any) inherited

This function is called from :func:load_state_dict to handle any extra state found within the state_dict. Implement this function and a corresponding :func:get_extra_state for your module if you need to store extra state within its state_dict.

Parameters:

Name Type Description Default
state dict

Extra state from the state_dict

required
Source code in zamba/pytorch/transforms.py
def set_extra_state(self, state: Any):
    """
    This function is called from :func:`load_state_dict` to handle any extra state
    found within the `state_dict`. Implement this function and a corresponding
    :func:`get_extra_state` for your module if you need to store extra state within its
    `state_dict`.

    Args:
        state (dict): Extra state from the `state_dict`
    """
    raise RuntimeError(
        "Reached a code path in Module.set_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
share_memory(self: ~T) -> ~T inherited

See :meth:torch.Tensor.share_memory_

Source code in zamba/pytorch/transforms.py
def share_memory(self: T) -> T:
    r"""See :meth:`torch.Tensor.share_memory_`"""
    return self._apply(lambda t: t.share_memory_())
state_dict(self, destination = None, prefix = '', keep_vars = False) inherited

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Returns:

Type Description
dict

a dictionary containing a whole state of the module

Example::

>>> module.state_dict().keys()
['bias', 'weight']
Source code in zamba/pytorch/transforms.py
def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.

    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Parameters and buffers set to ``None`` are not included.

    Returns:
        dict:
            a dictionary containing a whole state of the module

    Example::

        >>> module.state_dict().keys()
        ['bias', 'weight']

    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    self._save_to_state_dict(destination, prefix, keep_vars)
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination
to(self, *args, **kwargs) inherited

Moves and/or casts the parameters and buffers.

This can be called as

.. function:: to(device=None, dtype=None, non_blocking=False) :noindex:

.. function:: to(dtype, non_blocking=False) :noindex:

.. function:: to(tensor, non_blocking=False) :noindex:

.. function:: to(memory_format=torch.channels_last) :noindex:

Its signature is similar to :meth:torch.Tensor.to, but only accepts floating point or complex :attr:dtype\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:dtype (if given). The integral parameters and buffers will be moved :attr:device, if that is given, but with dtypes unchanged. When :attr:non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device

class:torch.device): the desired device of the parameters and buffers in this module

required
dtype

class:torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module

required
tensor torch.Tensor

Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

required
memory_format

class:torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword only argument)

required

Returns:

Type Description
Module

self

Examples::

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
Source code in zamba/pytorch/transforms.py
def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.

    This can be called as

    .. function:: to(device=None, dtype=None, non_blocking=False)
       :noindex:

    .. function:: to(dtype, non_blocking=False)
       :noindex:

    .. function:: to(tensor, non_blocking=False)
       :noindex:

    .. function:: to(memory_format=torch.channels_last)
       :noindex:

    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`\ s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.

    See below for examples.

    .. note::
        This method modifies the module in-place.

    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)

    Returns:
        Module: self

    Examples::

        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)

        >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.3741+0.j,  0.2382+0.j],
                [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
        >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
        tensor([[0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

    """

    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)

    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError('nn.Module.to only accepts floating point or complex '
                            'dtypes, but got desired dtype={}'.format(dtype))
        if dtype.is_complex:
            warnings.warn(
                "Complex modules are a new feature under active development whose design may change, "
                "and some modules might not work as expected when using complex tensors as parameters or buffers. "
                "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
                "if a complex module does not work as expected.")

    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

    return self._apply(convert)
to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T inherited

Moves the parameters and buffers to the specified device without copying storage.

Parameters:

Name Type Description Default
device

class:torch.device): The desired device of the parameters and buffers in this module.

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def to_empty(self: T, *, device: Union[str, device]) -> T:
    r"""Moves the parameters and buffers to the specified device without copying storage.

    Args:
        device (:class:`torch.device`): The desired device of the parameters
            and buffers in this module.

    Returns:
        Module: self
    """
    return self._apply(lambda t: torch.empty_like(t, device=device))
train(self: ~T, mode: bool = True) -> ~T inherited

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Parameters:

Name Type Description Default
mode bool

whether to set training mode (True) or evaluation mode (False). Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def train(self: T, mode: bool = True) -> T:
    r"""Sets the module in training mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    Args:
        mode (bool): whether to set training mode (``True``) or evaluation
                     mode (``False``). Default: ``True``.

    Returns:
        Module: self
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self
type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T inherited

Casts all parameters and buffers to :attr:dst_type.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
dst_type type or string

the desired type

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def type(self: T, dst_type: Union[dtype, str]) -> T:
    r"""Casts all parameters and buffers to :attr:`dst_type`.

    .. note::
        This method modifies the module in-place.

    Args:
        dst_type (type or string): the desired type

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))
xpu(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the XPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on XPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.xpu(device))
zero_grad(self, set_to_none: bool = False) -> None inherited

Sets gradients of all model parameters to zero. See similar function under :class:torch.optim.Optimizer for more context.

Parameters:

Name Type Description Default
set_to_none bool

instead of setting to zero, set the grads to None. See :meth:torch.optim.Optimizer.zero_grad for details.

False
Source code in zamba/pytorch/transforms.py
def zero_grad(self, set_to_none: bool = False) -> None:
    r"""Sets gradients of all model parameters to zero. See similar function
    under :class:`torch.optim.Optimizer` for more context.

    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            See :meth:`torch.optim.Optimizer.zero_grad` for details.
    """
    if getattr(self, '_is_replica', False):
        warnings.warn(
            "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
            "The parameters are copied (in a differentiable manner) from the original module. "
            "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
            "If you need gradients in your forward method, consider using autograd.grad instead.")

    for p in self.parameters():
        if p.grad is not None:
            if set_to_none:
                p.grad = None
            else:
                if p.grad.grad_fn is not None:
                    p.grad.detach_()
                else:
                    p.grad.requires_grad_(False)
                p.grad.zero_()

Uint8ToFloat (Module)

Source code in zamba/pytorch/transforms.py
class Uint8ToFloat(torch.nn.Module):
    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
        return tensor / 255.0

Attributes

T_destination inherited
dump_patches: bool inherited

This allows better BC support for :meth:load_state_dict. In :meth:state_dict, the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.

If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module's _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.

Methods

__init__(self) -> None inherited special
Source code in zamba/pytorch/transforms.py
def __init__(self) -> None:
    """
    Initializes internal Module state, shared by both nn.Module and ScriptModule.
    """
    torch._C._log_api_usage_once("python.nn_module")

    self.training = True
    self._parameters: Dict[str, Optional[Parameter]] = OrderedDict()
    self._buffers: Dict[str, Optional[Tensor]] = OrderedDict()
    self._non_persistent_buffers_set: Set[str] = set()
    self._backward_hooks: Dict[int, Callable] = OrderedDict()
    self._is_full_backward_hook = None
    self._forward_hooks: Dict[int, Callable] = OrderedDict()
    self._forward_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._state_dict_hooks: Dict[int, Callable] = OrderedDict()
    self._load_state_dict_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._modules: Dict[str, Optional['Module']] = OrderedDict()
add_module(self, name: str, module: Optional[Module]) -> None inherited

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

Name Type Description Default
name string

name of the child module. The child module can be accessed from this module using the given name

required
module Module

child module to be added to the module.

required
Source code in zamba/pytorch/transforms.py
def add_module(self, name: str, module: Optional['Module']) -> None:
    r"""Adds a child module to the current module.

    The module can be accessed as an attribute using the given name.

    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\", got: {}".format(name))
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module
apply(self: ~T, fn: Callable[[Module], NoneType]) -> ~T inherited

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also :ref:nn-init-doc).

Parameters:

Name Type Description Default
fn

class:Module -> None): function to be applied to each submodule

required

Returns:

Type Description
Module

self

Example::

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Source code in zamba/pytorch/transforms.py
def apply(self: T, fn: Callable[['Module'], None]) -> T:
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`nn-init-doc`).

    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule

    Returns:
        Module: self

    Example::

        >>> @torch.no_grad()
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self
bfloat16(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to bfloat16 datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def bfloat16(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
buffers(self, recurse: bool = True) -> Iterator[torch.Tensor] inherited

Returns an iterator over module buffers.

Parameters:

Name Type Description Default
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
torch.Tensor

module buffer

Example::

>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
    r"""Returns an iterator over module buffers.

    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        torch.Tensor: module buffer

    Example::

        >>> for buf in model.buffers():
        >>>     print(type(buf), buf.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for _, buf in self.named_buffers(recurse=recurse):
        yield buf
children(self) -> Iterator[Module] inherited

Returns an iterator over immediate children modules.

Yields:

Type Description
Module

a child module

Source code in zamba/pytorch/transforms.py
def children(self) -> Iterator['Module']:
    r"""Returns an iterator over immediate children modules.

    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module
cpu(self: ~T) -> ~T inherited

Moves all model parameters and buffers to the CPU.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())
cuda(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))
double(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to double datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def double(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``double`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)
eval(self: ~T) -> ~T inherited

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

This is equivalent with :meth:self.train(False) <torch.nn.Module.train>.

See :ref:locally-disable-grad-doc for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def eval(self: T) -> T:
    r"""Sets the module in evaluation mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.eval()` and several similar mechanisms that may be confused with it.

    Returns:
        Module: self
    """
    return self.train(False)
extra_repr(self) -> str inherited

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Source code in zamba/pytorch/transforms.py
def extra_repr(self) -> str:
    r"""Set the extra representation of the module

    To print customized extra information, you should re-implement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''
float(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to float datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def float(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``float`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)
forward(self, tensor: Tensor) -> Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

.. note:: Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Source code in zamba/pytorch/transforms.py
def forward(self, tensor: torch.Tensor) -> torch.Tensor:
    return tensor / 255.0
get_buffer(self, target: str) -> Tensor inherited

Returns the buffer given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.Tensor

The buffer referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not a buffer

Source code in zamba/pytorch/transforms.py
def get_buffer(self, target: str) -> "Tensor":
    """
    Returns the buffer given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the buffer
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.Tensor: The buffer referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not a
            buffer
    """
    module_path, _, buffer_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, buffer_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + buffer_name + "`")

    buffer: torch.Tensor = getattr(mod, buffer_name)

    if buffer_name not in mod._buffers:
        raise AttributeError("`" + buffer_name + "` is not a buffer")

    return buffer
get_extra_state(self) -> Any inherited

Returns any extra state to include in the module's state_dict. Implement this and a corresponding :func:set_extra_state for your module if you need to store extra state. This function is called when building the module's state_dict().

Note that extra state should be pickleable to ensure working serialization of the state_dict. We only provide provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Type Description
object

Any extra state to store in the module's state_dict

Source code in zamba/pytorch/transforms.py
def get_extra_state(self) -> Any:
    """
    Returns any extra state to include in the module's state_dict.
    Implement this and a corresponding :func:`set_extra_state` for your module
    if you need to store extra state. This function is called when building the
    module's `state_dict()`.

    Note that extra state should be pickleable to ensure working serialization
    of the state_dict. We only provide provide backwards compatibility guarantees
    for serializing Tensors; other objects may break backwards compatibility if
    their serialized pickled form changes.

    Returns:
        object: Any extra state to store in the module's state_dict
    """
    raise RuntimeError(
        "Reached a code path in Module.get_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
get_parameter(self, target: str) -> Parameter inherited

Returns the parameter given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Parameter

The Parameter referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Parameter

Source code in zamba/pytorch/transforms.py
def get_parameter(self, target: str) -> "Parameter":
    """
    Returns the parameter given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the Parameter
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Parameter: The Parameter referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Parameter``
    """
    module_path, _, param_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, param_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + param_name + "`")

    param: torch.nn.Parameter = getattr(mod, param_name)

    if not isinstance(param, torch.nn.Parameter):
        raise AttributeError("`" + param_name + "` is not an "
                             "nn.Parameter")

    return param
get_submodule(self, target: str) -> Module inherited

Returns the submodule given by target if it exists, otherwise throws an error.

For example, let's say you have an nn.Module A that looks like this:

.. code-block::text

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Module

The submodule referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Module

Source code in zamba/pytorch/transforms.py
def get_submodule(self, target: str) -> "Module":
    """
    Returns the submodule given by ``target`` if it exists,
    otherwise throws an error.

    For example, let's say you have an ``nn.Module`` ``A`` that
    looks like this:

    .. code-block::text

        A(
            (net_b): Module(
                (net_c): Module(
                    (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
                )
                (linear): Linear(in_features=100, out_features=200, bias=True)
            )
        )

    (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
    submodule ``net_b``, which itself has two submodules ``net_c``
    and ``linear``. ``net_c`` then has a submodule ``conv``.)

    To check whether or not we have the ``linear`` submodule, we
    would call ``get_submodule("net_b.linear")``. To check whether
    we have the ``conv`` submodule, we would call
    ``get_submodule("net_b.net_c.conv")``.

    The runtime of ``get_submodule`` is bounded by the degree
    of module nesting in ``target``. A query against
    ``named_modules`` achieves the same result, but it is O(N) in
    the number of transitive modules. So, for a simple check to see
    if some submodule exists, ``get_submodule`` should always be
    used.

    Args:
        target: The fully-qualified string name of the submodule
            to look for. (See above example for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Module: The submodule referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Module``
    """
    if target == "":
        return self

    atoms: List[str] = target.split(".")
    mod: torch.nn.Module = self

    for item in atoms:

        if not hasattr(mod, item):
            raise AttributeError(mod._get_name() + " has no "
                                 "attribute `" + item + "`")

        mod = getattr(mod, item)

        if not isinstance(mod, torch.nn.Module):
            raise AttributeError("`" + item + "` is not "
                                 "an nn.Module")

    return mod
half(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to half datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def half(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``half`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)
load_state_dict(self, state_dict: OrderedDict[str, Tensor], strict: bool = True) inherited

Copies parameters and buffers from :attr:state_dict into this module and its descendants. If :attr:strict is True, then the keys of :attr:state_dict must exactly match the keys returned by this module's :meth:~torch.nn.Module.state_dict function.

Parameters:

Name Type Description Default
state_dict dict

a dict containing parameters and persistent buffers.

required
strict bool

whether to strictly enforce that the keys in :attr:state_dict match the keys returned by this module's :meth:~torch.nn.Module.state_dict function. Default: True

True

Returns:

Type Description
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields
  • missing_keys is a list of str containing the missing keys
    • unexpected_keys is a list of str containing the unexpected keys

!!! note If a parameter or buffer is registered as None and its corresponding key exists in :attr:state_dict, :meth:load_state_dict will raise a RuntimeError.

Source code in zamba/pytorch/transforms.py
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                    strict: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``

    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Note:
        If a parameter or buffer is registered as ``None`` and its corresponding key
        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
        ``RuntimeError``.
    """
    missing_keys: List[str] = []
    unexpected_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        # mypy isn't aware that "_metadata" exists in state_dict
        state_dict._metadata = metadata  # type: ignore[attr-defined]

    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')

    load(self)
    del load

    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))

    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)
modules(self) -> Iterator[Module] inherited

Returns an iterator over all modules in the network.

Yields:

Type Description
Module

a module in the network

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
Source code in zamba/pytorch/transforms.py
def modules(self) -> Iterator['Module']:
    r"""Returns an iterator over all modules in the network.

    Yields:
        Module: a module in the network

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)

        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)

    """
    for _, module in self.named_modules():
        yield module
named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]] inherited

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all buffer names.

''
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
(string, torch.Tensor)

Tuple containing the name and buffer

Example::

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
Source code in zamba/pytorch/transforms.py
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.

    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer

    Example::

        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())

    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
named_children(self) -> Iterator[Tuple[str, Module]] inherited

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

Type Description
(string, Module)

Tuple containing a name and child module

Example::

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Source code in zamba/pytorch/transforms.py
def named_children(self) -> Iterator[Tuple[str, 'Module']]:
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.

    Yields:
        (string, Module): Tuple containing a name and child module

    Example::

        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)

    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module
named_modules(self, memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True) inherited

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

Name Type Description Default
memo Optional[Set[Module]]

a memo to store the set of modules already added to the result

None
prefix str

a prefix that will be added to the name of the module

''
remove_duplicate bool

whether to remove the duplicated module instances in the result

True

Yields:

Type Description
(string, Module)

Tuple of name and module

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
Source code in zamba/pytorch/transforms.py
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = '', remove_duplicate: bool = True):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.

    Args:
        memo: a memo to store the set of modules already added to the result
        prefix: a prefix that will be added to the name of the module
        remove_duplicate: whether to remove the duplicated module instances in the result
        or not

    Yields:
        (string, Module): Tuple of name and module

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)

        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

    """

    if memo is None:
        memo = set()
    if self not in memo:
        if remove_duplicate:
            memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix, remove_duplicate):
                yield m
named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]] inherited

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all parameter names.

''
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
(string, Parameter)

Tuple containing the name and parameter

Example::

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
Source code in zamba/pytorch/transforms.py
def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.

    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        (string, Parameter): Tuple containing the name and parameter

    Example::

        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())

    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter] inherited

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

Name Type Description Default
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
Parameter

module parameter

Example::

>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters.

    This is typically passed to an optimizer.

    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        Parameter: module parameter

    Example::

        >>> for param in model.parameters():
        >>>     print(type(param), param.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param
register_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

This function is deprecated in favor of :meth:~torch.nn.Module.register_full_backward_hook and the behavior of this function will change in future versions.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
    the behavior of this function will change in future versions.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is True:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = False

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None inherited

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

Name Type Description Default
name string

name of the buffer. The buffer can be accessed from this module using the given name

required
tensor Tensor or None

buffer to be registered. If None, then operations that run on buffers, such as :attr:cuda, are ignored. If None, the buffer is not included in the module's :attr:state_dict.

required
persistent bool

whether the buffer is part of this module's :attr:state_dict.

True

Example::

>>> self.register_buffer('running_mean', torch.zeros(num_features))
Source code in zamba/pytorch/transforms.py
def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    r"""Adds a buffer to the module.

    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the module's state. Buffers, by
    default, are persistent and will be saved alongside parameters. This
    behavior can be changed by setting :attr:`persistent` to ``False``. The
    only difference between a persistent buffer and a non-persistent buffer
    is that the latter will not be a part of this module's
    :attr:`state_dict`.

    Buffers can be accessed as attributes using given names.

    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor or None): buffer to be registered. If ``None``, then operations
            that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
            the buffer is **not** included in the module's :attr:`state_dict`.
        persistent (bool): whether the buffer is part of this module's
            :attr:`state_dict`.

    Example::

        >>> self.register_buffer('running_mean', torch.zeros(num_features))

    """
    if persistent is False and isinstance(self, torch.jit.ScriptModule):
        raise RuntimeError("ScriptModule does not support non-persistent buffers")

    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor
        if persistent:
            self._non_persistent_buffers_set.discard(name)
        else:
            self._non_persistent_buffers_set.add(name)
register_forward_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward hook on the module.

The hook will be called every time after :func:forward has computed an output. It should have the following signature::

hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after :func:forward is called.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward hook on the module.

    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::

        hook(module, input, output) -> None or modified output

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the output. It can modify the input inplace but
    it will not have effect on forward since this is called after
    :func:`forward` is called.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle
register_forward_pre_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward pre-hook on the module.

The hook will be called every time before :func:forward is invoked. It should have the following signature::

hook(module, input) -> None or modified input

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward pre-hook on the module.

    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::

        hook(module, input) -> None or modified input

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the input. User can either return a tuple or a
    single modified value in the hook. We will wrap the value into a tuple
    if a single value is returned(unless that value is already a tuple).

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle
register_full_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature::

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The :attr:grad_input and :attr:grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of :attr:grad_input in subsequent computations. :attr:grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in :attr:grad_input and :attr:grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module's forward function.

.. warning :: Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_full_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::

        hook(module, grad_input, grad_output) -> tuple(Tensor) or None

    The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
    with respect to the inputs and outputs respectively. The hook should
    not modify its arguments, but it can optionally return a new gradient with
    respect to the input that will be used in place of :attr:`grad_input` in
    subsequent computations. :attr:`grad_input` will only correspond to the inputs given
    as positional arguments and all kwarg arguments are ignored. Entries
    in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
    arguments.

    For technical reasons, when this hook is applied to a Module, its forward function will
    receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
    of each Tensor returned by the Module's forward function.

    .. warning ::
        Modifying inputs or outputs inplace is not allowed when using backward hooks and
        will raise an error.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is False:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = True

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None inherited

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

Name Type Description Default
name string

name of the parameter. The parameter can be accessed from this module using the given name

required
param Parameter or None

parameter to be added to the module. If None, then operations that run on parameters, such as :attr:cuda, are ignored. If None, the parameter is not included in the module's :attr:state_dict.

required
Source code in zamba/pytorch/transforms.py
def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    r"""Adds a parameter to the module.

    The parameter can be accessed as an attribute using given name.

    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter or None): parameter to be added to the module. If
            ``None``, then operations that run on parameters, such as :attr:`cuda`,
            are ignored. If ``None``, the parameter is **not** included in the
            module's :attr:`state_dict`.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")

    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))

    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param
requires_grad_(self: ~T, requires_grad: bool = True) -> ~T inherited

Change if autograd should record operations on parameters in this module.

This method sets the parameters' :attr:requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See :ref:locally-disable-grad-doc for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

Name Type Description Default
requires_grad bool

whether autograd should record operations on parameters in this module. Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def requires_grad_(self: T, requires_grad: bool = True) -> T:
    r"""Change if autograd should record operations on parameters in this
    module.

    This method sets the parameters' :attr:`requires_grad` attributes
    in-place.

    This method is helpful for freezing part of the module for finetuning
    or training parts of a model individually (e.g., GAN training).

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.requires_grad_()` and several similar mechanisms that may be confused with it.

    Args:
        requires_grad (bool): whether autograd should record operations on
                              parameters in this module. Default: ``True``.

    Returns:
        Module: self
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self
set_extra_state(self, state: Any) inherited

This function is called from :func:load_state_dict to handle any extra state found within the state_dict. Implement this function and a corresponding :func:get_extra_state for your module if you need to store extra state within its state_dict.

Parameters:

Name Type Description Default
state dict

Extra state from the state_dict

required
Source code in zamba/pytorch/transforms.py
def set_extra_state(self, state: Any):
    """
    This function is called from :func:`load_state_dict` to handle any extra state
    found within the `state_dict`. Implement this function and a corresponding
    :func:`get_extra_state` for your module if you need to store extra state within its
    `state_dict`.

    Args:
        state (dict): Extra state from the `state_dict`
    """
    raise RuntimeError(
        "Reached a code path in Module.set_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
share_memory(self: ~T) -> ~T inherited

See :meth:torch.Tensor.share_memory_

Source code in zamba/pytorch/transforms.py
def share_memory(self: T) -> T:
    r"""See :meth:`torch.Tensor.share_memory_`"""
    return self._apply(lambda t: t.share_memory_())
state_dict(self, destination = None, prefix = '', keep_vars = False) inherited

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Returns:

Type Description
dict

a dictionary containing a whole state of the module

Example::

>>> module.state_dict().keys()
['bias', 'weight']
Source code in zamba/pytorch/transforms.py
def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.

    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Parameters and buffers set to ``None`` are not included.

    Returns:
        dict:
            a dictionary containing a whole state of the module

    Example::

        >>> module.state_dict().keys()
        ['bias', 'weight']

    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    self._save_to_state_dict(destination, prefix, keep_vars)
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination
to(self, *args, **kwargs) inherited

Moves and/or casts the parameters and buffers.

This can be called as

.. function:: to(device=None, dtype=None, non_blocking=False) :noindex:

.. function:: to(dtype, non_blocking=False) :noindex:

.. function:: to(tensor, non_blocking=False) :noindex:

.. function:: to(memory_format=torch.channels_last) :noindex:

Its signature is similar to :meth:torch.Tensor.to, but only accepts floating point or complex :attr:dtype\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:dtype (if given). The integral parameters and buffers will be moved :attr:device, if that is given, but with dtypes unchanged. When :attr:non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device

class:torch.device): the desired device of the parameters and buffers in this module

required
dtype

class:torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module

required
tensor torch.Tensor

Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

required
memory_format

class:torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword only argument)

required

Returns:

Type Description
Module

self

Examples::

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
Source code in zamba/pytorch/transforms.py
def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.

    This can be called as

    .. function:: to(device=None, dtype=None, non_blocking=False)
       :noindex:

    .. function:: to(dtype, non_blocking=False)
       :noindex:

    .. function:: to(tensor, non_blocking=False)
       :noindex:

    .. function:: to(memory_format=torch.channels_last)
       :noindex:

    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`\ s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.

    See below for examples.

    .. note::
        This method modifies the module in-place.

    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)

    Returns:
        Module: self

    Examples::

        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)

        >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.3741+0.j,  0.2382+0.j],
                [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
        >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
        tensor([[0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

    """

    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)

    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError('nn.Module.to only accepts floating point or complex '
                            'dtypes, but got desired dtype={}'.format(dtype))
        if dtype.is_complex:
            warnings.warn(
                "Complex modules are a new feature under active development whose design may change, "
                "and some modules might not work as expected when using complex tensors as parameters or buffers. "
                "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
                "if a complex module does not work as expected.")

    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

    return self._apply(convert)
to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T inherited

Moves the parameters and buffers to the specified device without copying storage.

Parameters:

Name Type Description Default
device

class:torch.device): The desired device of the parameters and buffers in this module.

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def to_empty(self: T, *, device: Union[str, device]) -> T:
    r"""Moves the parameters and buffers to the specified device without copying storage.

    Args:
        device (:class:`torch.device`): The desired device of the parameters
            and buffers in this module.

    Returns:
        Module: self
    """
    return self._apply(lambda t: torch.empty_like(t, device=device))
train(self: ~T, mode: bool = True) -> ~T inherited

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Parameters:

Name Type Description Default
mode bool

whether to set training mode (True) or evaluation mode (False). Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def train(self: T, mode: bool = True) -> T:
    r"""Sets the module in training mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    Args:
        mode (bool): whether to set training mode (``True``) or evaluation
                     mode (``False``). Default: ``True``.

    Returns:
        Module: self
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self
type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T inherited

Casts all parameters and buffers to :attr:dst_type.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
dst_type type or string

the desired type

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def type(self: T, dst_type: Union[dtype, str]) -> T:
    r"""Casts all parameters and buffers to :attr:`dst_type`.

    .. note::
        This method modifies the module in-place.

    Args:
        dst_type (type or string): the desired type

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))
xpu(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the XPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on XPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.xpu(device))
zero_grad(self, set_to_none: bool = False) -> None inherited

Sets gradients of all model parameters to zero. See similar function under :class:torch.optim.Optimizer for more context.

Parameters:

Name Type Description Default
set_to_none bool

instead of setting to zero, set the grads to None. See :meth:torch.optim.Optimizer.zero_grad for details.

False
Source code in zamba/pytorch/transforms.py
def zero_grad(self, set_to_none: bool = False) -> None:
    r"""Sets gradients of all model parameters to zero. See similar function
    under :class:`torch.optim.Optimizer` for more context.

    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            See :meth:`torch.optim.Optimizer.zero_grad` for details.
    """
    if getattr(self, '_is_replica', False):
        warnings.warn(
            "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
            "The parameters are copied (in a differentiable manner) from the original module. "
            "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
            "If you need gradients in your forward method, consider using autograd.grad instead.")

    for p in self.parameters():
        if p.grad is not None:
            if set_to_none:
                p.grad = None
            else:
                if p.grad.grad_fn is not None:
                    p.grad.detach_()
                else:
                    p.grad.requires_grad_(False)
                p.grad.zero_()

VideotoImg (Module)

Source code in zamba/pytorch/transforms.py
class VideotoImg(torch.nn.Module):
    def forward(self, vid: torch.Tensor) -> torch.Tensor:
        return vid.squeeze(0)

Attributes

T_destination inherited
dump_patches: bool inherited

This allows better BC support for :meth:load_state_dict. In :meth:state_dict, the version number will be saved as in the attribute _metadata of the returned state dict, and thus pickled. _metadata is a dictionary with keys that follow the naming convention of state dict. See _load_from_state_dict on how to use this information in loading.

If new parameters/buffers are added/removed from a module, this number shall be bumped, and the module's _load_from_state_dict method can compare the version number and do appropriate changes if the state dict is from before the change.

Methods

__init__(self) -> None inherited special
Source code in zamba/pytorch/transforms.py
def __init__(self) -> None:
    """
    Initializes internal Module state, shared by both nn.Module and ScriptModule.
    """
    torch._C._log_api_usage_once("python.nn_module")

    self.training = True
    self._parameters: Dict[str, Optional[Parameter]] = OrderedDict()
    self._buffers: Dict[str, Optional[Tensor]] = OrderedDict()
    self._non_persistent_buffers_set: Set[str] = set()
    self._backward_hooks: Dict[int, Callable] = OrderedDict()
    self._is_full_backward_hook = None
    self._forward_hooks: Dict[int, Callable] = OrderedDict()
    self._forward_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._state_dict_hooks: Dict[int, Callable] = OrderedDict()
    self._load_state_dict_pre_hooks: Dict[int, Callable] = OrderedDict()
    self._modules: Dict[str, Optional['Module']] = OrderedDict()
add_module(self, name: str, module: Optional[Module]) -> None inherited

Adds a child module to the current module.

The module can be accessed as an attribute using the given name.

Parameters:

Name Type Description Default
name string

name of the child module. The child module can be accessed from this module using the given name

required
module Module

child module to be added to the module.

required
Source code in zamba/pytorch/transforms.py
def add_module(self, name: str, module: Optional['Module']) -> None:
    r"""Adds a child module to the current module.

    The module can be accessed as an attribute using the given name.

    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\", got: {}".format(name))
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module
apply(self: ~T, fn: Callable[[Module], NoneType]) -> ~T inherited

Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also :ref:nn-init-doc).

Parameters:

Name Type Description Default
fn

class:Module -> None): function to be applied to each submodule

required

Returns:

Type Description
Module

self

Example::

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Source code in zamba/pytorch/transforms.py
def apply(self: T, fn: Callable[['Module'], None]) -> T:
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`nn-init-doc`).

    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule

    Returns:
        Module: self

    Example::

        >>> @torch.no_grad()
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self
bfloat16(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to bfloat16 datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def bfloat16(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``bfloat16`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
buffers(self, recurse: bool = True) -> Iterator[torch.Tensor] inherited

Returns an iterator over module buffers.

Parameters:

Name Type Description Default
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
torch.Tensor

module buffer

Example::

>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def buffers(self, recurse: bool = True) -> Iterator[Tensor]:
    r"""Returns an iterator over module buffers.

    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        torch.Tensor: module buffer

    Example::

        >>> for buf in model.buffers():
        >>>     print(type(buf), buf.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for _, buf in self.named_buffers(recurse=recurse):
        yield buf
children(self) -> Iterator[Module] inherited

Returns an iterator over immediate children modules.

Yields:

Type Description
Module

a child module

Source code in zamba/pytorch/transforms.py
def children(self) -> Iterator['Module']:
    r"""Returns an iterator over immediate children modules.

    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module
cpu(self: ~T) -> ~T inherited

Moves all model parameters and buffers to the CPU.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cpu(self: T) -> T:
    r"""Moves all model parameters and buffers to the CPU.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())
cuda(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the GPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def cuda(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the GPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Args:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))
double(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to double datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def double(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``double`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)
eval(self: ~T) -> ~T inherited

Sets the module in evaluation mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

This is equivalent with :meth:self.train(False) <torch.nn.Module.train>.

See :ref:locally-disable-grad-doc for a comparison between .eval() and several similar mechanisms that may be confused with it.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def eval(self: T) -> T:
    r"""Sets the module in evaluation mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.eval()` and several similar mechanisms that may be confused with it.

    Returns:
        Module: self
    """
    return self.train(False)
extra_repr(self) -> str inherited

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Source code in zamba/pytorch/transforms.py
def extra_repr(self) -> str:
    r"""Set the extra representation of the module

    To print customized extra information, you should re-implement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''
float(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to float datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def float(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``float`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)
forward(self, vid: Tensor) -> Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

.. note:: Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Source code in zamba/pytorch/transforms.py
def forward(self, vid: torch.Tensor) -> torch.Tensor:
    return vid.squeeze(0)
get_buffer(self, target: str) -> Tensor inherited

Returns the buffer given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the buffer to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.Tensor

The buffer referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not a buffer

Source code in zamba/pytorch/transforms.py
def get_buffer(self, target: str) -> "Tensor":
    """
    Returns the buffer given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the buffer
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.Tensor: The buffer referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not a
            buffer
    """
    module_path, _, buffer_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, buffer_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + buffer_name + "`")

    buffer: torch.Tensor = getattr(mod, buffer_name)

    if buffer_name not in mod._buffers:
        raise AttributeError("`" + buffer_name + "` is not a buffer")

    return buffer
get_extra_state(self) -> Any inherited

Returns any extra state to include in the module's state_dict. Implement this and a corresponding :func:set_extra_state for your module if you need to store extra state. This function is called when building the module's state_dict().

Note that extra state should be pickleable to ensure working serialization of the state_dict. We only provide provide backwards compatibility guarantees for serializing Tensors; other objects may break backwards compatibility if their serialized pickled form changes.

Returns:

Type Description
object

Any extra state to store in the module's state_dict

Source code in zamba/pytorch/transforms.py
def get_extra_state(self) -> Any:
    """
    Returns any extra state to include in the module's state_dict.
    Implement this and a corresponding :func:`set_extra_state` for your module
    if you need to store extra state. This function is called when building the
    module's `state_dict()`.

    Note that extra state should be pickleable to ensure working serialization
    of the state_dict. We only provide provide backwards compatibility guarantees
    for serializing Tensors; other objects may break backwards compatibility if
    their serialized pickled form changes.

    Returns:
        object: Any extra state to store in the module's state_dict
    """
    raise RuntimeError(
        "Reached a code path in Module.get_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
get_parameter(self, target: str) -> Parameter inherited

Returns the parameter given by target if it exists, otherwise throws an error.

See the docstring for get_submodule for a more detailed explanation of this method's functionality as well as how to correctly specify target.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the Parameter to look for. (See get_submodule for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Parameter

The Parameter referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Parameter

Source code in zamba/pytorch/transforms.py
def get_parameter(self, target: str) -> "Parameter":
    """
    Returns the parameter given by ``target`` if it exists,
    otherwise throws an error.

    See the docstring for ``get_submodule`` for a more detailed
    explanation of this method's functionality as well as how to
    correctly specify ``target``.

    Args:
        target: The fully-qualified string name of the Parameter
            to look for. (See ``get_submodule`` for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Parameter: The Parameter referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Parameter``
    """
    module_path, _, param_name = target.rpartition(".")

    mod: torch.nn.Module = self.get_submodule(module_path)

    if not hasattr(mod, param_name):
        raise AttributeError(mod._get_name() + " has no attribute `"
                             + param_name + "`")

    param: torch.nn.Parameter = getattr(mod, param_name)

    if not isinstance(param, torch.nn.Parameter):
        raise AttributeError("`" + param_name + "` is not an "
                             "nn.Parameter")

    return param
get_submodule(self, target: str) -> Module inherited

Returns the submodule given by target if it exists, otherwise throws an error.

For example, let's say you have an nn.Module A that looks like this:

.. code-block::text

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(The diagram shows an nn.Module A. A has a nested submodule net_b, which itself has two submodules net_c and linear. net_c then has a submodule conv.)

To check whether or not we have the linear submodule, we would call get_submodule("net_b.linear"). To check whether we have the conv submodule, we would call get_submodule("net_b.net_c.conv").

The runtime of get_submodule is bounded by the degree of module nesting in target. A query against named_modules achieves the same result, but it is O(N) in the number of transitive modules. So, for a simple check to see if some submodule exists, get_submodule should always be used.

Parameters:

Name Type Description Default
target str

The fully-qualified string name of the submodule to look for. (See above example for how to specify a fully-qualified string.)

required

Returns:

Type Description
torch.nn.Module

The submodule referenced by target

Exceptions:

Type Description
AttributeError

If the target string references an invalid path or resolves to something that is not an nn.Module

Source code in zamba/pytorch/transforms.py
def get_submodule(self, target: str) -> "Module":
    """
    Returns the submodule given by ``target`` if it exists,
    otherwise throws an error.

    For example, let's say you have an ``nn.Module`` ``A`` that
    looks like this:

    .. code-block::text

        A(
            (net_b): Module(
                (net_c): Module(
                    (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
                )
                (linear): Linear(in_features=100, out_features=200, bias=True)
            )
        )

    (The diagram shows an ``nn.Module`` ``A``. ``A`` has a nested
    submodule ``net_b``, which itself has two submodules ``net_c``
    and ``linear``. ``net_c`` then has a submodule ``conv``.)

    To check whether or not we have the ``linear`` submodule, we
    would call ``get_submodule("net_b.linear")``. To check whether
    we have the ``conv`` submodule, we would call
    ``get_submodule("net_b.net_c.conv")``.

    The runtime of ``get_submodule`` is bounded by the degree
    of module nesting in ``target``. A query against
    ``named_modules`` achieves the same result, but it is O(N) in
    the number of transitive modules. So, for a simple check to see
    if some submodule exists, ``get_submodule`` should always be
    used.

    Args:
        target: The fully-qualified string name of the submodule
            to look for. (See above example for how to specify a
            fully-qualified string.)

    Returns:
        torch.nn.Module: The submodule referenced by ``target``

    Raises:
        AttributeError: If the target string references an invalid
            path or resolves to something that is not an
            ``nn.Module``
    """
    if target == "":
        return self

    atoms: List[str] = target.split(".")
    mod: torch.nn.Module = self

    for item in atoms:

        if not hasattr(mod, item):
            raise AttributeError(mod._get_name() + " has no "
                                 "attribute `" + item + "`")

        mod = getattr(mod, item)

        if not isinstance(mod, torch.nn.Module):
            raise AttributeError("`" + item + "` is not "
                                 "an nn.Module")

    return mod
half(self: ~T) -> ~T inherited

Casts all floating point parameters and buffers to half datatype.

.. note:: This method modifies the module in-place.

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def half(self: T) -> T:
    r"""Casts all floating point parameters and buffers to ``half`` datatype.

    .. note::
        This method modifies the module in-place.

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)
load_state_dict(self, state_dict: OrderedDict[str, Tensor], strict: bool = True) inherited

Copies parameters and buffers from :attr:state_dict into this module and its descendants. If :attr:strict is True, then the keys of :attr:state_dict must exactly match the keys returned by this module's :meth:~torch.nn.Module.state_dict function.

Parameters:

Name Type Description Default
state_dict dict

a dict containing parameters and persistent buffers.

required
strict bool

whether to strictly enforce that the keys in :attr:state_dict match the keys returned by this module's :meth:~torch.nn.Module.state_dict function. Default: True

True

Returns:

Type Description
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields
  • missing_keys is a list of str containing the missing keys
    • unexpected_keys is a list of str containing the unexpected keys

!!! note If a parameter or buffer is registered as None and its corresponding key exists in :attr:state_dict, :meth:load_state_dict will raise a RuntimeError.

Source code in zamba/pytorch/transforms.py
def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                    strict: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``

    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys

    Note:
        If a parameter or buffer is registered as ``None`` and its corresponding key
        exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a
        ``RuntimeError``.
    """
    missing_keys: List[str] = []
    unexpected_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        # mypy isn't aware that "_metadata" exists in state_dict
        state_dict._metadata = metadata  # type: ignore[attr-defined]

    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')

    load(self)
    del load

    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))

    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)
modules(self) -> Iterator[Module] inherited

Returns an iterator over all modules in the network.

Yields:

Type Description
Module

a module in the network

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
Source code in zamba/pytorch/transforms.py
def modules(self) -> Iterator['Module']:
    r"""Returns an iterator over all modules in the network.

    Yields:
        Module: a module in the network

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)

        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)

    """
    for _, module in self.named_modules():
        yield module
named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.Tensor]] inherited

Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all buffer names.

''
recurse bool

if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.

True

Yields:

Type Description
(string, torch.Tensor)

Tuple containing the name and buffer

Example::

>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
Source code in zamba/pytorch/transforms.py
def named_buffers(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Tensor]]:
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.

    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.

    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer

    Example::

        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())

    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
named_children(self) -> Iterator[Tuple[str, Module]] inherited

Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

Yields:

Type Description
(string, Module)

Tuple containing a name and child module

Example::

>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
Source code in zamba/pytorch/transforms.py
def named_children(self) -> Iterator[Tuple[str, 'Module']]:
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.

    Yields:
        (string, Module): Tuple containing a name and child module

    Example::

        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)

    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module
named_modules(self, memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True) inherited

Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

Parameters:

Name Type Description Default
memo Optional[Set[Module]]

a memo to store the set of modules already added to the result

None
prefix str

a prefix that will be added to the name of the module

''
remove_duplicate bool

whether to remove the duplicated module instances in the result

True

Yields:

Type Description
(string, Module)

Tuple of name and module

!!! note Duplicate modules are returned only once. In the following example, l will be returned only once.

Example::

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
Source code in zamba/pytorch/transforms.py
def named_modules(self, memo: Optional[Set['Module']] = None, prefix: str = '', remove_duplicate: bool = True):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.

    Args:
        memo: a memo to store the set of modules already added to the result
        prefix: a prefix that will be added to the name of the module
        remove_duplicate: whether to remove the duplicated module instances in the result
        or not

    Yields:
        (string, Module): Tuple of name and module

    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.

    Example::

        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)

        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

    """

    if memo is None:
        memo = set()
    if self not in memo:
        if remove_duplicate:
            memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix, remove_duplicate):
                yield m
named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, torch.nn.parameter.Parameter]] inherited

Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

Parameters:

Name Type Description Default
prefix str

prefix to prepend to all parameter names.

''
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
(string, Parameter)

Tuple containing the name and parameter

Example::

>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
Source code in zamba/pytorch/transforms.py
def named_parameters(self, prefix: str = '', recurse: bool = True) -> Iterator[Tuple[str, Parameter]]:
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.

    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        (string, Parameter): Tuple containing the name and parameter

    Example::

        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())

    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem
parameters(self, recurse: bool = True) -> Iterator[torch.nn.parameter.Parameter] inherited

Returns an iterator over module parameters.

This is typically passed to an optimizer.

Parameters:

Name Type Description Default
recurse bool

if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.

True

Yields:

Type Description
Parameter

module parameter

Example::

>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
Source code in zamba/pytorch/transforms.py
def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
    r"""Returns an iterator over module parameters.

    This is typically passed to an optimizer.

    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.

    Yields:
        Parameter: module parameter

    Example::

        >>> for param in model.parameters():
        >>>     print(type(param), param.size())
        <class 'torch.Tensor'> (20L,)
        <class 'torch.Tensor'> (20L, 1L, 5L, 5L)

    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param
register_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

This function is deprecated in favor of :meth:~torch.nn.Module.register_full_backward_hook and the behavior of this function will change in future versions.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    This function is deprecated in favor of :meth:`~torch.nn.Module.register_full_backward_hook` and
    the behavior of this function will change in future versions.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is True:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = False

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_buffer(self, name: str, tensor: Optional[torch.Tensor], persistent: bool = True) -> None inherited

Adds a buffer to the module.

This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm's running_mean is not a parameter, but is part of the module's state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting :attr:persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module's :attr:state_dict.

Buffers can be accessed as attributes using given names.

Parameters:

Name Type Description Default
name string

name of the buffer. The buffer can be accessed from this module using the given name

required
tensor Tensor or None

buffer to be registered. If None, then operations that run on buffers, such as :attr:cuda, are ignored. If None, the buffer is not included in the module's :attr:state_dict.

required
persistent bool

whether the buffer is part of this module's :attr:state_dict.

True

Example::

>>> self.register_buffer('running_mean', torch.zeros(num_features))
Source code in zamba/pytorch/transforms.py
def register_buffer(self, name: str, tensor: Optional[Tensor], persistent: bool = True) -> None:
    r"""Adds a buffer to the module.

    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the module's state. Buffers, by
    default, are persistent and will be saved alongside parameters. This
    behavior can be changed by setting :attr:`persistent` to ``False``. The
    only difference between a persistent buffer and a non-persistent buffer
    is that the latter will not be a part of this module's
    :attr:`state_dict`.

    Buffers can be accessed as attributes using given names.

    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor or None): buffer to be registered. If ``None``, then operations
            that run on buffers, such as :attr:`cuda`, are ignored. If ``None``,
            the buffer is **not** included in the module's :attr:`state_dict`.
        persistent (bool): whether the buffer is part of this module's
            :attr:`state_dict`.

    Example::

        >>> self.register_buffer('running_mean', torch.zeros(num_features))

    """
    if persistent is False and isinstance(self, torch.jit.ScriptModule):
        raise RuntimeError("ScriptModule does not support non-persistent buffers")

    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor
        if persistent:
            self._non_persistent_buffers_set.discard(name)
        else:
            self._non_persistent_buffers_set.add(name)
register_forward_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward hook on the module.

The hook will be called every time after :func:forward has computed an output. It should have the following signature::

hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after :func:forward is called.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward hook on the module.

    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::

        hook(module, input, output) -> None or modified output

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the output. It can modify the input inplace but
    it will not have effect on forward since this is called after
    :func:`forward` is called.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle
register_forward_pre_hook(self, hook: Callable[..., NoneType]) -> RemovableHandle inherited

Registers a forward pre-hook on the module.

The hook will be called every time before :func:forward is invoked. It should have the following signature::

hook(module, input) -> None or modified input

The input contains only the positional arguments given to the module. Keyword arguments won't be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_forward_pre_hook(self, hook: Callable[..., None]) -> RemovableHandle:
    r"""Registers a forward pre-hook on the module.

    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::

        hook(module, input) -> None or modified input

    The input contains only the positional arguments given to the module.
    Keyword arguments won't be passed to the hooks and only to the ``forward``.
    The hook can modify the input. User can either return a tuple or a
    single modified value in the hook. We will wrap the value into a tuple
    if a single value is returned(unless that value is already a tuple).

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle
register_full_backward_hook(self, hook: Callable[[Module, Union[Tuple[torch.Tensor, ...], torch.Tensor], Union[Tuple[torch.Tensor, ...], torch.Tensor]], Union[NoneType, torch.Tensor]]) -> RemovableHandle inherited

Registers a backward hook on the module.

The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature::

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

The :attr:grad_input and :attr:grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of :attr:grad_input in subsequent computations. :attr:grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in :attr:grad_input and :attr:grad_output will be None for all non-Tensor arguments.

For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module's forward function.

.. warning :: Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.

Returns:

Type Description

class:torch.utils.hooks.RemovableHandle: a handle that can be used to remove the added hook by calling handle.remove()

Source code in zamba/pytorch/transforms.py
def register_full_backward_hook(
    self, hook: Callable[['Module', _grad_t, _grad_t], Union[None, Tensor]]
) -> RemovableHandle:
    r"""Registers a backward hook on the module.

    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::

        hook(module, grad_input, grad_output) -> tuple(Tensor) or None

    The :attr:`grad_input` and :attr:`grad_output` are tuples that contain the gradients
    with respect to the inputs and outputs respectively. The hook should
    not modify its arguments, but it can optionally return a new gradient with
    respect to the input that will be used in place of :attr:`grad_input` in
    subsequent computations. :attr:`grad_input` will only correspond to the inputs given
    as positional arguments and all kwarg arguments are ignored. Entries
    in :attr:`grad_input` and :attr:`grad_output` will be ``None`` for all non-Tensor
    arguments.

    For technical reasons, when this hook is applied to a Module, its forward function will
    receive a view of each Tensor passed to the Module. Similarly the caller will receive a view
    of each Tensor returned by the Module's forward function.

    .. warning ::
        Modifying inputs or outputs inplace is not allowed when using backward hooks and
        will raise an error.

    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``

    """
    if self._is_full_backward_hook is False:
        raise RuntimeError("Cannot use both regular backward hooks and full backward hooks on a "
                           "single Module. Please use only one of them.")

    self._is_full_backward_hook = True

    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle
register_parameter(self, name: str, param: Optional[torch.nn.parameter.Parameter]) -> None inherited

Adds a parameter to the module.

The parameter can be accessed as an attribute using given name.

Parameters:

Name Type Description Default
name string

name of the parameter. The parameter can be accessed from this module using the given name

required
param Parameter or None

parameter to be added to the module. If None, then operations that run on parameters, such as :attr:cuda, are ignored. If None, the parameter is not included in the module's :attr:state_dict.

required
Source code in zamba/pytorch/transforms.py
def register_parameter(self, name: str, param: Optional[Parameter]) -> None:
    r"""Adds a parameter to the module.

    The parameter can be accessed as an attribute using given name.

    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter or None): parameter to be added to the module. If
            ``None``, then operations that run on parameters, such as :attr:`cuda`,
            are ignored. If ``None``, the parameter is **not** included in the
            module's :attr:`state_dict`.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")

    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))

    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param
requires_grad_(self: ~T, requires_grad: bool = True) -> ~T inherited

Change if autograd should record operations on parameters in this module.

This method sets the parameters' :attr:requires_grad attributes in-place.

This method is helpful for freezing part of the module for finetuning or training parts of a model individually (e.g., GAN training).

See :ref:locally-disable-grad-doc for a comparison between .requires_grad_() and several similar mechanisms that may be confused with it.

Parameters:

Name Type Description Default
requires_grad bool

whether autograd should record operations on parameters in this module. Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def requires_grad_(self: T, requires_grad: bool = True) -> T:
    r"""Change if autograd should record operations on parameters in this
    module.

    This method sets the parameters' :attr:`requires_grad` attributes
    in-place.

    This method is helpful for freezing part of the module for finetuning
    or training parts of a model individually (e.g., GAN training).

    See :ref:`locally-disable-grad-doc` for a comparison between
    `.requires_grad_()` and several similar mechanisms that may be confused with it.

    Args:
        requires_grad (bool): whether autograd should record operations on
                              parameters in this module. Default: ``True``.

    Returns:
        Module: self
    """
    for p in self.parameters():
        p.requires_grad_(requires_grad)
    return self
set_extra_state(self, state: Any) inherited

This function is called from :func:load_state_dict to handle any extra state found within the state_dict. Implement this function and a corresponding :func:get_extra_state for your module if you need to store extra state within its state_dict.

Parameters:

Name Type Description Default
state dict

Extra state from the state_dict

required
Source code in zamba/pytorch/transforms.py
def set_extra_state(self, state: Any):
    """
    This function is called from :func:`load_state_dict` to handle any extra state
    found within the `state_dict`. Implement this function and a corresponding
    :func:`get_extra_state` for your module if you need to store extra state within its
    `state_dict`.

    Args:
        state (dict): Extra state from the `state_dict`
    """
    raise RuntimeError(
        "Reached a code path in Module.set_extra_state() that should never be called. "
        "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
        "to report this bug.")
share_memory(self: ~T) -> ~T inherited

See :meth:torch.Tensor.share_memory_

Source code in zamba/pytorch/transforms.py
def share_memory(self: T) -> T:
    r"""See :meth:`torch.Tensor.share_memory_`"""
    return self._apply(lambda t: t.share_memory_())
state_dict(self, destination = None, prefix = '', keep_vars = False) inherited

Returns a dictionary containing a whole state of the module.

Both parameters and persistent buffers (e.g. running averages) are included. Keys are corresponding parameter and buffer names. Parameters and buffers set to None are not included.

Returns:

Type Description
dict

a dictionary containing a whole state of the module

Example::

>>> module.state_dict().keys()
['bias', 'weight']
Source code in zamba/pytorch/transforms.py
def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.

    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Parameters and buffers set to ``None`` are not included.

    Returns:
        dict:
            a dictionary containing a whole state of the module

    Example::

        >>> module.state_dict().keys()
        ['bias', 'weight']

    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    self._save_to_state_dict(destination, prefix, keep_vars)
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination
to(self, *args, **kwargs) inherited

Moves and/or casts the parameters and buffers.

This can be called as

.. function:: to(device=None, dtype=None, non_blocking=False) :noindex:

.. function:: to(dtype, non_blocking=False) :noindex:

.. function:: to(tensor, non_blocking=False) :noindex:

.. function:: to(memory_format=torch.channels_last) :noindex:

Its signature is similar to :meth:torch.Tensor.to, but only accepts floating point or complex :attr:dtype\ s. In addition, this method will only cast the floating point or complex parameters and buffers to :attr:dtype (if given). The integral parameters and buffers will be moved :attr:device, if that is given, but with dtypes unchanged. When :attr:non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device

class:torch.device): the desired device of the parameters and buffers in this module

required
dtype

class:torch.dtype): the desired floating point or complex dtype of the parameters and buffers in this module

required
tensor torch.Tensor

Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

required
memory_format

class:torch.memory_format): the desired memory format for 4D parameters and buffers in this module (keyword only argument)

required

Returns:

Type Description
Module

self

Examples::

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
Source code in zamba/pytorch/transforms.py
def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.

    This can be called as

    .. function:: to(device=None, dtype=None, non_blocking=False)
       :noindex:

    .. function:: to(dtype, non_blocking=False)
       :noindex:

    .. function:: to(tensor, non_blocking=False)
       :noindex:

    .. function:: to(memory_format=torch.channels_last)
       :noindex:

    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point or complex :attr:`dtype`\ s. In addition, this method will
    only cast the floating point or complex parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.

    See below for examples.

    .. note::
        This method modifies the module in-place.

    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point or complex dtype of
            the parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
        memory_format (:class:`torch.memory_format`): the desired memory
            format for 4D parameters and buffers in this module (keyword
            only argument)

    Returns:
        Module: self

    Examples::

        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)

        >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.3741+0.j,  0.2382+0.j],
                [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
        >>> linear(torch.ones(3, 2, dtype=torch.cdouble))
        tensor([[0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j],
                [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

    """

    device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)

    if dtype is not None:
        if not (dtype.is_floating_point or dtype.is_complex):
            raise TypeError('nn.Module.to only accepts floating point or complex '
                            'dtypes, but got desired dtype={}'.format(dtype))
        if dtype.is_complex:
            warnings.warn(
                "Complex modules are a new feature under active development whose design may change, "
                "and some modules might not work as expected when using complex tensors as parameters or buffers. "
                "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.md "
                "if a complex module does not work as expected.")

    def convert(t):
        if convert_to_format is not None and t.dim() in (4, 5):
            return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None,
                        non_blocking, memory_format=convert_to_format)
        return t.to(device, dtype if t.is_floating_point() or t.is_complex() else None, non_blocking)

    return self._apply(convert)
to_empty(self: ~T, *, device: Union[str, torch.device]) -> ~T inherited

Moves the parameters and buffers to the specified device without copying storage.

Parameters:

Name Type Description Default
device

class:torch.device): The desired device of the parameters and buffers in this module.

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def to_empty(self: T, *, device: Union[str, device]) -> T:
    r"""Moves the parameters and buffers to the specified device without copying storage.

    Args:
        device (:class:`torch.device`): The desired device of the parameters
            and buffers in this module.

    Returns:
        Module: self
    """
    return self._apply(lambda t: torch.empty_like(t, device=device))
train(self: ~T, mode: bool = True) -> ~T inherited

Sets the module in training mode.

This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. :class:Dropout, :class:BatchNorm, etc.

Parameters:

Name Type Description Default
mode bool

whether to set training mode (True) or evaluation mode (False). Default: True.

True

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def train(self: T, mode: bool = True) -> T:
    r"""Sets the module in training mode.

    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.

    Args:
        mode (bool): whether to set training mode (``True``) or evaluation
                     mode (``False``). Default: ``True``.

    Returns:
        Module: self
    """
    if not isinstance(mode, bool):
        raise ValueError("training mode is expected to be boolean")
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self
type(self: ~T, dst_type: Union[torch.dtype, str]) -> ~T inherited

Casts all parameters and buffers to :attr:dst_type.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
dst_type type or string

the desired type

required

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def type(self: T, dst_type: Union[dtype, str]) -> T:
    r"""Casts all parameters and buffers to :attr:`dst_type`.

    .. note::
        This method modifies the module in-place.

    Args:
        dst_type (type or string): the desired type

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))
xpu(self: ~T, device: Union[int, torch.device] = None) -> ~T inherited

Moves all model parameters and buffers to the XPU.

This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on XPU while being optimized.

.. note:: This method modifies the module in-place.

Parameters:

Name Type Description Default
device int

if specified, all parameters will be copied to that device

None

Returns:

Type Description
Module

self

Source code in zamba/pytorch/transforms.py
def xpu(self: T, device: Optional[Union[int, device]] = None) -> T:
    r"""Moves all model parameters and buffers to the XPU.

    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on XPU while being optimized.

    .. note::
        This method modifies the module in-place.

    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device

    Returns:
        Module: self
    """
    return self._apply(lambda t: t.xpu(device))
zero_grad(self, set_to_none: bool = False) -> None inherited

Sets gradients of all model parameters to zero. See similar function under :class:torch.optim.Optimizer for more context.

Parameters:

Name Type Description Default
set_to_none bool

instead of setting to zero, set the grads to None. See :meth:torch.optim.Optimizer.zero_grad for details.

False
Source code in zamba/pytorch/transforms.py
def zero_grad(self, set_to_none: bool = False) -> None:
    r"""Sets gradients of all model parameters to zero. See similar function
    under :class:`torch.optim.Optimizer` for more context.

    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            See :meth:`torch.optim.Optimizer.zero_grad` for details.
    """
    if getattr(self, '_is_replica', False):
        warnings.warn(
            "Calling .zero_grad() from a module created with nn.DataParallel() has no effect. "
            "The parameters are copied (in a differentiable manner) from the original module. "
            "This means they are not leaf nodes in autograd and so don't accumulate gradients. "
            "If you need gradients in your forward method, consider using autograd.grad instead.")

    for p in self.parameters():
        if p.grad is not None:
            if set_to_none:
                p.grad = None
            else:
                if p.grad.grad_fn is not None:
                    p.grad.detach_()
                else:
                    p.grad.requires_grad_(False)
                p.grad.zero_()

slowfast_transforms()

Source code in zamba/pytorch/transforms.py
def slowfast_transforms():
    return transforms.Compose(
        [
            ConvertTHWCtoTCHW(),
            Uint8ToFloat(),
            Normalize(mean=[0.45, 0.45, 0.45], std=[0.225, 0.225, 0.225]),
            ConvertTCHWtoCTHW(),
            PadDimensions((None, 32, None, None)),
            PackSlowFastPathways(),
        ]
    )

zamba_image_model_transforms(single_frame = False, normalization_values = {'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}, channels_first = False)

Source code in zamba/pytorch/transforms.py
def zamba_image_model_transforms(
    single_frame=False, normalization_values=imagenet_normalization_values, channels_first=False
):
    img_transforms = [
        ConvertTHWCtoTCHW(),
        Uint8ToFloat(),
        transforms.Normalize(**imagenet_normalization_values),
    ]

    if single_frame:
        img_transforms += [VideotoImg()]  # squeeze dim

    if channels_first:
        img_transforms += [ConvertTCHWtoCTHW()]

    return transforms.Compose(img_transforms)