Module#
Base class for deriving trainable modules.
- class ivy.stateful.module.Module(*args, device=None, v=None, build_mode='on_init', compile_on_next_step=False, store_vars=True, stateful=None, arg_stateful_idxs=None, kwarg_stateful_idxs=None, fallback_to_non_compiled=False, with_partial_v=False, devices=None, dtype=None, dynamic_backend=None, **kwargs)[source]#
Bases:
ModuleConverters
,ModuleHelpers
Module is a base class for deriving trainable modules.
- __call__(*args, v=None, stateful=None, arg_stateful_idxs=None, kwarg_stateful_idxs=None, track_submod_rets=False, submod_depth=None, submods_to_track=None, track_submod_call_order=False, expected_submod_rets=None, **kwargs)[source]#
Forward an input through current module.
- Parameters:
v – If given, use this container as internal varibles temporarily. Default is
None
.track_submod_rets – If True, will track the returns of submodules.
submod_depth – The depth of tracked submodules.
submods_to_track – If given, will only track submodules in submods_to_track.
track_submod_call_order – If True, will track the call order of submodules.
expected_submod_rets – If given, will raise exception if submodule returns are different from expected returns.
- Returns:
ret
- __init__(*args, device=None, v=None, build_mode='on_init', compile_on_next_step=False, store_vars=True, stateful=None, arg_stateful_idxs=None, kwarg_stateful_idxs=None, fallback_to_non_compiled=False, with_partial_v=False, devices=None, dtype=None, dynamic_backend=None, **kwargs)[source]#
Initialize Ivy layer, which is a stateful object consisting of trainable variables.
- Parameters:
device – device on which to create the module’s variables ‘cuda:0’, ‘cuda:1’, ‘cpu’ etc. (Default value = None)
v – Ivy container of trainable variables. Created internally by default.
build_mode – 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.
compile_on_next_step – Whether to compile the network on the next forward pass. Default is
False
.store_vars – Whether or not to store the variables created. Default is
True
.stateful – The constant id stateful items to track as part of the forward pass. Used when graph compiling, default is
None
.arg_stateful_idxs – The nested argument indices of stateful items to track as part of the forward pass. Used when graph compiling, default is
None
.kwarg_stateful_idxs – The nested keyword argument indices of stateful items to track as part of the forward pass. Used when graph compiling, default is
None
.fallback_to_non_compiled – Whether to fall back to non-compiled forward call in the case that an error is raised during the compiled forward pass. Default is
True
.with_partial_v – Whether to allow partial specification of variables. Default is
False
.devices – devices on which to distribute the module’s variables ‘cuda:0’, ‘cuda:1’, ‘cpu’ etc. (Default value = None)
- build(*args, from_call=False, device=None, dtype=None, dynamic_backend=None, **kwargs)[source]#
Build the internal layers and variables for this module.
- Parameters:
from_call – If True, denote that this build is triggered by calling. Otherwise, triggered by initializing the module. Default is
False
.device – The device we want to build module on. None for default device. Default is
None
.dtype – The data type for building the module. Default is
None
.
- Returns:
ret – True for successfully built a module.
- property build_mode#
- property built_#
- compile(args=None, kwargs=None, **compile_kwargs)[source]#
Compile the ivy.Module’s _unified_ivy_graph or _call method to the target backend.
- Parameters:
args (
Optional
[Tuple
]) – arguments used to compile. Defaults to None. (default:None
)kwargs (
Optional
[Dict
]) – keyword arguments used to compile. Defaults to None. (default:None
)compile_kwargs – keyword arguments passed to the compile function.
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 or in the module forum!