Module#

Base class for deriving trainable modules.

class ivy.stateful.module.Module(*args, **kwargs)[source]#

Bases: ModuleHelpers, ModuleConverters, ModuleMeta

Module is a base class for deriving trainable modules.

__call__(*args, v=None, buffers=None, **kwargs)[source]#

Forward an input through current module.

Parameters:

  • args – Positional args to the build method.

  • v (default: None) – If given, use this container as internal variables temporarily. Default is None.

  • buffers (default: None) – If given, use this container as internal buffers temporarily. Default is None.

  • kwargs – Keyword arguments to the build method.

Returns:

ret

__init__(*args, v=None, buffers=None, build_mode='on_init', store_vars=True, with_partial_v=False, dynamic_backend=None, training=True, dtype=None, device=None, **kwargs)[source]#

Initialize Ivy layer, which is a stateful object consisting of trainable variables.

Parameters:

  • args – Positional arguments to the _build method.

  • v (default: None) – Ivy container of trainable variables. Created internally by default.

  • buffers (default: None) – Ivy container of buffers/non-trainable arrays in the state_dict.

  • build_mode (default: 'on_init') – How the Module is built, either on initialization (now), explicitly by the user by calling build(), or the first time the __call__ method is run. Default is on initialization.

  • store_vars (default: True) – Whether or not to store the variables created. Default is True.

  • with_partial_v (default: False) – Whether to allow partial specification of variables. Default is False.

  • dynamic_backend (default: None) – When the value is true, allow conversion of arrays from a different backend to the current backend if v passed in the input contains arrays created with different backend.

  • training (default: True) – specifies whether the module is in training or evaluation mode. Default is True.

  • dtype (default: None) – Data type to be used for creating model variables. (Default value = None).

  • device (default: None) – Device on which to create the module’s variables ‘cuda:0’, ‘cuda:1’, ‘cpu’ etc. (Default value = None).

  • kwargs – Keyword arguments to the _build method.

property buffers#
build(*args, from_call=False, device=None, dtype=None, dynamic_backend=None, **kwargs)[source]#

Build the internal layers and variables for this module.

Parameters:

  • args – Positional arguments to the _build method.

  • from_call (default: False) – If True, denote that this build is triggered by calling. Otherwise, triggered by initializing the module. Default is False.

  • device (default: None) – The device we want to build module on. None for default device. Default is None.

  • dtype (default: None) – The data type for building the module. Default is None.

  • dynamic_backend (default: None) – Whether to use dynamic backend setting to deal if variables are passed as input and created with a different backend to the current backend.

  • kwargs – Keyword arguments to the _build method.

Returns:

ret – True for successfully built a module.

property build_mode#
property built#
property device#
property dtype#
eval()[source]#

Disable training mode.

static load(filename)[source]#

Load a module object from disk using pickle.

Parameters:

filename (str) – The name of the file to load the module object from.

Returns:

Module – The loaded module object.

property module_dict#
register_buffer(name, value)[source]#

Register a buffer.

Parameters:

  • name – Name of the buffer

  • value – Value of the buffer

register_parameter(name, value)[source]#

Register a parameter.

Parameters:

  • name – Name of the buffer

  • value – Value of the buffer

save(filename)[source]#

Save the module object to disk using pickle.

Parameters:

filename (str) – The name of the file to save the module object to.

save_weights(weights_path, /)[source]#

Save the weights on the Module.

Parameters:

weights_path – The hdf5 file for saving the weights.

Returns:

None

show_graph(randomness_factor=0.1, save_to_disk=False, notebook=False, with_edge_labels=True, with_arg_labels=True, with_output_labels=True, output_connected_only=True, highlight_subgraph=None, fname=None)[source]#
property state_dict#

Return the state_dict which is a collection of the variables and buffers.

to_device(device)[source]#

Move the weights and buffers to the specified device.

trace_graph(args=None, kwargs=None, **trace_kwargs)[source]#

Trace the ivy.Module’s _unified_ivy_graph or _call method to the target backend.

Parameters:

  • args (Optional[Tuple], default: None) – arguments used to trace. Defaults to None.

  • kwargs (Optional[Dict], default: None) – keyword arguments used to trace. Defaults to None.

  • trace_kwargs – keyword arguments passed to the trace function.

train(mode=True)[source]#

Enable or disable training mode.

property training#
property v#
class ivy.stateful.module.ModuleMeta(*args, **kwargs)[source]#

Bases: object

This should have hopefully given you an overview of the module submodule, if you have any questions, please feel free to reach out on our discord in the module channel!