Python v0.3.0.dev+c1b8f91
  • Overview

    • Module

    class dragon.vm.torch.nn.Module[source]

    The base class of modules.

    Inherit this class to design a new module:

    class MyModule(torch.nn.Module):
    def __init__():
        super(MyModule, self).__init__()

    __init__

    Module.__init__()[source]

    Create a Module.

    Methods

    add_module

    Module.add_module(
      name,
      module
    )[source]

    Add a submodule to the module.

    Parameters:

    apply

    Module.apply(fn)[source]

    Apply the function over submodules.

    Parameters:
    • fn (callable) – The function to call.
    Returns:

    dragon.vm.torch.nn.Module – The self.

    children

    Module.children()[source]

    Return an iterator over immediate modules.

    Returns:
    Iterator – The iterator of module.

    cpu

    Module.cpu()[source]

    Switch the buffers and parameters to cpu device.

    Returns:
    dragon.vm.torch.nn.Module – The self.

    cuda

    Module.cuda(device=None)[source]

    Switch the buffers and parameters to cuda device.

    If device is not provided, use the value set by dragon.config.set_cuda_device().

    Parameters:
    • device (int, optional) – The optional device index.
    Returns:

    dragon.vm.torch.nn.Module – The self.

    double

    Module.double()[source]

    Switch the buffers and parameters to float64.

    Returns:
    dragon.vm.torch.nn.Module – The self.

    eval

    Module.eval()[source]

    Set to the evaluation mode.

    This method is identical to Module.train(False).

    Returns:
    dragon.vm.torch.nn.Module – The self.

    float

    Module.float()[source]

    Switch the buffers and parameters to float32.

    Returns:
    dragon.vm.torch.nn.Module – The self.

    forward

    Module.forward(
      *inputs,
      **kwargs
    )[source]

    Define the implementation of forward.

    All subclasses should override this method:

    class MyModule(torch.nn.Module):
    def forward(*inputs, **kwargs):
        pass

    half

    Module.half()[source]

    Switch the buffers and parameters to float16.

    Returns:
    dragon.vm.torch.nn.Module – The self.

    load_state_dict

    Module.load_state_dict(
      state_dict,
      strict=True,
      verbose=False
    )[source]

    Load the state dict from other module.

    Typically, states can only loaded from the same module class:

    mm = type(m)()
mm.load_state_dict(m.state_dict())

    Set strict to False to load a mismatched dict:

    # States matching the name successfully will be loaded
# Otherwise, we will directly ignore them
mm.load_state_dict(m.state_dict(), strict=False)
    Parameters:
    • state_dict (Dict) – The state dict.
    • strict (bool, optional, default=True) – True to verify the names strictly.
    • verbose (bool, optional, default=False) – True to print the state info.

    modules

    Module.modules()[source]

    Return an iterator over all modules.

    Returns:
    Iterator – The iterator of module.

    named_children

    Module.named_children()[source]

    Return an iterator over immediate modules, yield as (name, module).

    Returns:
    Iterator – The iterator of module.

    named_modules

    Module.named_modules(
      memo=None,
      prefix=''
    )[source]

    Return an iterator over all modules, yield as (name, module).

    Parameters:
    • memo (Set, optional) – The optional set to collect modules.
    • prefix (str, optional, default='') – The prefix added to the name.
    Returns:

    Iterator – The iterator of module.

    named_parameters

    Module.named_parameters(
      memo=None,
      prefix=''
    )[source]

    Return an iterator over all parameters.

    Parameters:
    • memo (Set, optional) – The optional set to collect parameters.
    • prefix (str, optional, default='') – The prefix added to the name.
    Returns:

    Iterator – The iterator of parameter.

    parameters

    Module.parameters()[source]

    Return an iterator over all parameters.

    Returns:
    Iterator – The iterator of parameter.

    register_buffer

    Module.register_buffer(
      name,
      tensor
    )[source]

    Add a buffer to the module.

    Parameters:

    register_forward_hook

    Module.register_forward_hook(hook)[source]

    Register forward hook on the module.

    Parameters:
    • hook (callable) – The hook function.
    Returns:

    RemovableHandle – The handle to remove this hook by calling handle.remove().

    register_parameter

    Module.register_parameter(
      name,
      param
    )[source]

    Add a parameter to the module.

    This method is identical to assign a parameter as attribute:

    m = torch.nn.Module()
weight = torch.nn.Parameter(torch.ones(1))
m.register_parameter('weight', weight)  # Style1
m.weight = weight  # Style2
    Parameters:

    state_dict

    Module.state_dict(
      destination=None,
      prefix='',
      to_numpy=False
    )[source]

    Return a dict stored the buffers and parameters.

    Usually, we will use this method to renew another module:

    m2.load_state_dict(m1.state_dict())

    Set to_numpy if you want to serialize these states:

    # Currently, ``torch.Tensor`` is not supported to pickle
# Convert tensors to numpy arrays before pickling
np_states = m.state_dict(to_numpy=True)

with open('states.pkl', 'wb') as f:
    pickle.dump(np_states, f, pickle.HIGHEST_PROTOCOL)
    Parameters:
    • destination (Dict, optional) – The optional output dict.
    • prefix (str, optional, default='') – The prefix added to the name of states.
    • to_numpy (bool, optional, default=False) – True to store the numpy array instead.
    Returns:

    Dict – The state dict.

    train

    Module.train(mode=True)[source]

    Set the training mode.

    Parameters:
    • mode (bool, optional, default=True) – The training mode.
    Returns:

    dragon.vm.torch.nn.Module – The self.