yamle.models.model module

class yamle.models.model.BaseModel(inputs_dim, outputs_dim, task, seed, *args, **kwargs)

Bases: Module, ABC

This is the base class for all the models.

By default it should have an input and output layer in _input and _output respectively. All the intermediate layers should be in _layers. The depth of the model should be in _depth.

Parameters:

• inputs_dim (Tuple[int,...]) – The input dimensions.

• outputs_dim (int) – The output dimension.

• task (str) – The task to perform.

tasks = ['regression', 'classification', 'text_classification', 'segmentation', 'depth_estimation', 'pre_training', 'reconstruction']

abstract forward(x)

This method is used to perform a forward pass of the model.

Return type:

Tensor

abstract final_layer(x, **output_kwargs)

This function is used to get the final layer output.

Return type:

Tensor

classmethod add_specific_args(parent_parser)

This method adds model arguments to the given parser.

Return type:

ArgumentParser

reset()

This method is used to reset the model e.g. at the start of a new epoch.

Return type:

None

replace_layers_for_quantization()

Fuses all the operations in the network.

In this function we only need to fuse layers that are not in the blocks. e.g. the reshaping layers added by the method.

Return type:

None

add_method_specific_layers(method, **kwargs)

This method is used to add method specific layers to the model.

Parameters:

method (str) – The method to use.

Return type:

None

training: bool