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 parameter

  • value – Value of the parameter

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!