Training and Evaluation

Training and Evaluation

The MLFramework Template centralizes the execution of all machine learning workflows through a unified Command Line Interface (CLI) located at src/cli.py. This CLI utilizes Typer for command routing and Hydra for dynamic configuration injection.

Command Line Interface Architecture

The entrypoint script src/cli.py exposes strictly typed commands. The primary commands available are train and test. Both commands require a configuration file specified via the -c or --config flag.

# General syntax
python -m src.cli [COMMAND] -c [CONFIG_FILE.yaml] [OVERRIDES...]

Executing a Training Pipeline

To initiate a training run, invoke the train command. The pipeline will automatically instantiate the DataModule, Model Architecture, Lightning Task, and the designated Logger based on the provided configuration tree.

python -m src.cli train -c example_config.yaml

Parameter Overrides

Hyperparameters and pipeline behavior can be modified at runtime without editing the YAML files.

# Example: Overriding batch size, learning rate, and epochs
python -m src.cli train -c example_config.yaml \
    data.batch_size=256 \
    model.learning_rate=0.001 \
    trainer.max_epochs=50

Experiment Tracking Integration

By default, the template is designed to integrate with MLflow via the PyTorch Lightning MLFlowLogger. During the execution of the train command:

1. A new MLflow run is initialized.
2. All Hydra configuration parameters are logged as MLflow parameters.
3. Training metrics (loss, accuracy, etc.) are tracked continuously.
4. The best-performing model checkpoint (as defined by the ModelCheckpoint callback in the configuration) is saved automatically.

Executing an Evaluation (Testing) Pipeline

To evaluate a fully trained model on the designated test dataset, utilize the test command. This pipeline bypasses the optimization loop and strictly executes the test_step defined in the PyTorch Lightning Task.

python -m src.cli test -c example_config.yaml ckpt_path="path/to/checkpoint.ckpt"

Checkpoint Resolution Mechanism

The CLI implements an intelligent checkpoint resolution strategy via the internal resolve_checkpoint_path utility. The ckpt_path argument accepts multiple formats:

1. Local File Path: Direct execution from a local .ckpt file.

ckpt_path="/absolute/or/relative/path/to/model.ckpt"

2. MLflow Run URI: Direct resolution from a remote or local MLflow artifact store. The utility automatically interfaces with the MLflow tracking server to download and cache the requisite checkpoint.

ckpt_path="runs:/<mlflow_run_id>/checkpoints/best_model.ckpt"

This feature guarantees that evaluation pipelines can be executed locally or in CI/CD environments directly from centrally tracked artifacts, eliminating the need to manually transfer large checkpoint files.