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__()
Methods¶
add_module¶
Module.
add_module
(
name,
module
)[source]¶Add a submodule to the module.
- Parameters:
- name (str) – The buffer name.
- module (dragon.vm.torch.nn.Module, optional) – The submodule to be registered.
apply¶
buffers¶
children¶
cpu¶
cuda¶
double¶
eval¶
float¶
forward¶
half¶
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
toFalse
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
andunexpected_keys
.
modules¶
named_buffers¶
named_children¶
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¶
register_buffer¶
Module.
register_buffer
(
name,
tensor
)[source]¶Add a buffer to the module.
- Parameters:
- name (str) – The buffer name.
- tensor (dragon.vm.torch.Tensor) – The tensor to be registered.
register_forward_hook¶
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:
- name (str) – The buffer name.
- param (dragon.vm.torch.Tensor, optional) – The tensor to be registered.
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.