Training API Reference

API documentation for training and evaluation.

Training Functions

train

mlops_project.train.train(cfg)

Train specified model on skin lesion dataset with Hydra configuration and W&B logging.

Supports both normal training runs and W&B sweeps.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	Hydra configuration object	required

train_model

mlops_project.train.train_model(model, model_name, train_loader, val_loader, epochs=10, patience=7, output_dir='models', wandb_logger=None)

Train a given model and return the trained model and training metrics.

create_model

mlops_project.train.create_model(cfg)

Create a model based on the configuration.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	Hydra configuration containing model parameters	required

Returns:

Type	Description
LightningModule	Instantiated PyTorch Lightning model

Evaluation Functions

evaluate

mlops_project.evaluate.evaluate(model_name, model_checkpoint, model_config, data_path, image_size=224, batch_size=32, num_workers=4, train_ratio=0.525, val_ratio=0.175, test_ratio=0.3, random_seed=42)

Evaluate a trained model on the test set and print accuracy.

Model Export

export_model_to_onnx

mlops_project.train.export_model_to_onnx(model, checkpoint_path, image_size)

Export trained model to ONNX format.

Parameters:

Name	Type	Description	Default
model	LightningModule	Trained PyTorch Lightning model	required
checkpoint_path	Path	Path to the model checkpoint	required
image_size	int	Size of input images	required

Returns:

Type	Description
Path	Path to the exported ONNX model

Usage Examples

from mlops_project.train import train_model, create_model
from mlops_project.evaluate import evaluate

# Train with Hydra config
# uv run python -m mlops_project.train model=efficientnet

# Evaluate a trained model
results = evaluate(
    model_path="models/best_model.ckpt",
    data_dir="data/"
)