CLI

class ssm.TrainingCLI(config_file=None)

Bases: object

Command Line Interface for training models. This class is responsible for loading the model and dataset from the configuration file, initializing the trainer, and starting the training process. It allows a straightforward way to configure and run training jobs from the command line, by simply providing a configuration file in YAML format.

static argparsing()

Parse command line arguments. :return: Parsed arguments. :rtype: argparse.Namespace

fit()

Start the training process. :return: None

static init_model(model_config, n_classes, mem_tokens)

Load the model from the configuration. :param dict model_config: Model configuration dictionary. :return: Model instance. :rtype: torch.nn.Module

static init_trainer(trainer_config, model, dataset)

Initialize the trainer with the given configuration.

Parameters:
  • trainer_config (dict) – Trainer configuration dictionary.

  • model – Model instance.

  • dataset – Dataset instance.

Returns:

Trainer instance.

Return type:

Trainer

initialize_logger(config)

Initialize the logger for TensorBoard. :param dict config: Configuration dictionary.

load_config(config_file)

Configure the training parameters. :param str config_file: Path to the configuration file. :return: Configuration dictionary. :rtype: dict

test()

Start the testing process. :return: None

write_on_tensorboard(config)

Write the configuration on TensorBoard. :param dict config: Configuration dictionary.