Python v0.3.0

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.

buffers

Module.buffers(recurse=True)[source]

Return an iterator over all buffers.

Parameters:
  • recurse (bool, optional, default=True) – Yield parameters recursively or not.
Returns:

Iterator – The iterator of buffer.

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.cuda.set_default_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 computation performed at every call.

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
)[source]

Load the state dict from other module.

Typically, states can only be 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.
Returns:

namedtuple – The namedtuple with missing_keys and unexpected_keys.

modules

Module.modules()[source]

Return an iterator over all modules.

Returns:
Iterator – The iterator of module.

named_buffers

Module.named_buffers(
  prefix='',
  recurse=True
)[source]

Return an iterator over all buffers.

Parameters:
  • prefix (str, optional, default='') – The prefix added to the name.
  • recurse (bool, optional, default=True) – Yield buffers recursively or not.
Returns:

Iterator – The iterator of (name, buffer).

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 (name, module).

named_parameters

Module.named_parameters(
  prefix='',
  recurse=True
)[source]

Return an iterator over all parameters.

Parameters:
  • prefix (str, optional, default='') – The prefix added to the name.
  • recurse (bool, optional, default=True) – Yield parameters recursively or not.
Returns:

Iterator – The iterator of (name, param).

parameters

Module.parameters(recurse=True)[source]

Return an iterator over all parameters.

Parameters:
  • recurse (bool, optional, default=True) – Yield parameters recursively or not.
Returns:

Iterator – The iterator of param.

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.

to

Module.to(
  *args,
  **kwargs
)[source]

Convert states to the specified data type or device.

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

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.

zero_grad

Module.zero_grad()[source]

Set the gradient of parameters to zero.