pytorch-TrainEngine 0.1.0

A flexible PyTorch training engine with interface-based design

PyTorch TrainEngine

A flexible and modular PyTorch training engine with interface-based design.

Features

• Interface-based Architecture: Define trainable models through the ITrainable interface contract
• Automatic Hardware Detection: Automatically detects and uses CUDA if available
• Lazy Optimizer Initialization: Optimizers are initialized only when training starts
• Standard Training Pipeline: Implements the canonical PyTorch training loop
• Easy Integration: Simple integration with custom models and datasets

Installation

pip install pytorch-TrainEngine

Quick Start

import torch
from pytorch_trainengine import ITrainable, TrainingEngine

# Define your model
class MyModel(ITrainable):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 1)
    
    def forward_pass(self, x):
        return self.linear(x)
    
    def parameters(self):
        return self.linear.parameters()

# Create training engine
criterion = torch.nn.MSELoss()
engine = TrainingEngine(criterion=criterion, lr=0.01)

# Train your model
model = MyModel()
x = torch.randn(32, 10)
y = torch.randn(32, 1)

loss = engine.train_step(model, x, y)
print(f"Loss: {loss}")

API Reference

ITrainable

Abstract base class for trainable models.

Methods:

• forward_pass(x): Forward pass logic
• parameters(): Return model parameters for optimizer

TrainingEngine

Main training engine class.

Parameters:

• criterion: Loss function
• optimizer_cls: Optimizer class (default: torch.optim.SGD)
• lr: Learning rate (default: 0.01)
• device: Computation device (auto-detected if None)

Methods:

• setup_optimizer(model): Initialize optimizer (called automatically)
• train_step(model, x, y): Single training step
• train_epoch(model, dataloader): Train one epoch

Requirements

• Python >= 3.8
• torch >= 1.9.0

License

MIT

Author

[e-mail]