Train Model

Train Model

Implement and execute model training loops including forward/backward passes, gradient updates, and checkpoint management.

When to Use

• Running full training pipeline on datasets
• Fine-tuning pretrained models
• Experimenting with hyperparameter variations
• Reproducing paper results

Quick Reference

# Mojo training loop pattern
struct Trainer:
    var model: NeuralNetwork
    var optimizer: Optimizer
    var loss_fn: LossFn

    fn train_epoch(mut self, mut dataloader: BatchLoader) -> Float32:
        var total_loss: Float32 = 0.0
        var batches: Int = 0
        for batch in dataloader:
            var predictions = self.model(batch.inputs)
            var loss = self.loss_fn(predictions, batch.targets)
            # Backward pass and optimization
            total_loss += loss
            batches += 1
        return total_loss / Float32(batches)

Workflow

1. Prepare data pipeline: Load and batch training data
2. Initialize model: Create network with specified architecture
3. Set up optimizer: Choose optimizer (SGD, Adam) with learning rate
4. Implement training loop: Forward pass, compute loss, backward pass, update weights
5. Monitor progress: Log loss, save checkpoints, validate periodically

Output Format

Training report:

• Loss values per epoch
• Training time per epoch
• Validation metrics (accuracy, loss)
• Learning curves (loss vs epoch)
• Final model performance
• Checkpoint locations

References

• See prepare-dataset skill for data pipeline setup
• See evaluate-model skill for validation
• See CLAUDE.md > Mojo for training loop patterns