pyprobound.base.Component

class Component(name='')

Bases: Module, ABC

Module that serves as a component in PyProBound.

Includes functions for loading from a checkpoint, freezing or unfreezing parameters, and defining sequential optimization procedures.

unfreezable

All possible values that can be passed to unfreeze().

_cache_fun

The name of a function in the module’s child components that will be cached to avoid recomputation.

_blocking

A mapping from the name of the cached function to the parent components waiting on that function’s output.

_caches

A mapping from the name of the cached function to a tuple of two optional elements, the input pointer and the output cache.

__init__(name='')

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Methods

check_length_consistency()	Checks that input lengths of Binding components are consistent.
components()	Iterator of child components.
forward(*input)	Define the computation performed at every call.
freeze()	Turns off gradient calculation for all parameters.
max_embedding_size()	The maximum number of bytes needed to encode a sequence.
optim_procedure([ancestry, current_order])	The sequential optimization procedure for all Binding components.
reload(checkpoint)	Loads the model from a checkpoint file.
reload_from_state_dict(state_dict)	Loads the model from a state dict.
save(checkpoint[, flank_lengths])	Saves the model to a file with "state_dict" and "metadata" fields.
unfreeze([parameter])	Turns on gradient calculation for the specified parameter.

Attributes

unfreezable

alias of Literal['all']

Non-Inherited Members

unfreezable

alias of Literal[‘all’]

save(checkpoint, flank_lengths=())

Saves the model to a file with “state_dict” and “metadata” fields.

Parameters:

• checkpoint (Union[str, PathLike, BinaryIO, IO[bytes]]) – The file where the model will be checkpointed to.

• flank_lengths (tuple[tuple[int, int], ...]) – The (left_flank_length, right_flank_length) of each table represented by the model, written to the metadata field.

Return type:

None

reload_from_state_dict(state_dict)

Loads the model from a state dict.

Parameters:

state_dict (dict[str, Any]) – The state dict, usually returned by self.state_dict().

Return type:

None

reload(checkpoint)

Loads the model from a checkpoint file.

Parameters:

checkpoint (Union[str, PathLike, BinaryIO, IO[bytes]]) – The file where the model state_dict was written to.

Return type:

dict[str, Any]

Returns:

The metadata field of the checkpoint file.

abstract components()

Iterator of child components.

Return type:

Iterator[Component]

max_embedding_size()

The maximum number of bytes needed to encode a sequence.

Used for splitting calculations to avoid GPU limits on tensor sizes.

Return type:

int

freeze()

Turns off gradient calculation for all parameters.

Return type:

None

unfreeze(parameter='all')

Turns on gradient calculation for the specified parameter.

Parameters:

parameter (unfreezable) – Parameter to be unfrozen, defaults to all parameters.

Return type:

None

check_length_consistency()

Checks that input lengths of Binding components are consistent.

Raises:

RuntimeError – There is an input mismatch between components.

Return type:

None

optim_procedure(ancestry=None, current_order=None)

The sequential optimization procedure for all Binding components.

The optimization procedure is generated recursively through iteration over the child components of each module. All Binding components with the same specification returned from key() are trained jointly.

Parameters:

• ancestry (tuple[Component, ...] | None) – The parent components from the root for which the procedure is being generated to the current component.

• current_order (dict[tuple[Spec, ...], BindingOptim] | None) – Mapping of Binding component specifications to the sequential optimization procedure for those Binding components.

Return type:

dict[tuple[Spec, ...], BindingOptim]

Returns:

The current_order updated with the optimization of the current component’s children.