Skip to content

Model Trainer

tensoris.backend.trainers.trainer

Standard PyTorch Deep Learning Model Trainer.

Classes

ModelTrainer

Standard Deep Learning Model Trainer managing forward passes, backpropagation, and evaluation.

Execution Workflow Flowchart:

graph TD
    A["DataLoader Batch (inputs, targets)"] --> B["Transfer Tensors to Device (CUDA / MPS / CPU)"]
    B --> C["Optimizer.zero_grad()"]
    C --> D["Model Forward Pass: outputs = model(inputs)"]
    D --> E["Loss Calculation: loss = criterion(outputs, targets)"]
    E --> F["Backpropagation: loss.backward()"]
    F --> G["Optimizer Step: optimizer.step()"]
    G --> H["Accumulate & Return Epoch Mean Loss"]
Source code in src/tensoris/backend/trainers/trainer.py
class ModelTrainer:
    """Standard Deep Learning Model Trainer managing forward passes, backpropagation, and evaluation.

    Execution Workflow Flowchart:

    ```mermaid
    graph TD
        A["DataLoader Batch (inputs, targets)"] --> B["Transfer Tensors to Device (CUDA / MPS / CPU)"]
        B --> C["Optimizer.zero_grad()"]
        C --> D["Model Forward Pass: outputs = model(inputs)"]
        D --> E["Loss Calculation: loss = criterion(outputs, targets)"]
        E --> F["Backpropagation: loss.backward()"]
        F --> G["Optimizer Step: optimizer.step()"]
        G --> H["Accumulate & Return Epoch Mean Loss"]
    ```
    """

    def __init__(
        self,
        model: nn.Module,
        optimizer: torch.optim.Optimizer,
        criterion: nn.Module,
        device: str | torch.device = "cpu",
    ) -> None:
        """Initialize Model Trainer.

        Args:
            model: PyTorch neural network model instance.
            optimizer: PyTorch optimizer instance.
            criterion: PyTorch loss function instance.
            device: Compute device ('cuda', 'mps', or 'cpu').
        """
        self.model = model.to(device)
        self.optimizer = optimizer
        self.criterion = criterion
        self.device = device

    def train_epoch(self, dataloader: DataLoader) -> float:
        """Execute a single training epoch optimization loop.

        Args:
            dataloader: PyTorch DataLoader supplying training batches.

        Returns:
            Mean training loss across all batches.
        """
        self.model.train()
        total_loss = 0.0

        for inputs, targets in dataloader:
            inputs, targets = inputs.to(self.device), targets.to(self.device)

            self.optimizer.zero_grad()
            outputs = self.model(inputs)
            loss = self.criterion(outputs, targets)
            loss.backward()
            self.optimizer.step()

            total_loss += loss.item()

        return total_loss / max(1, len(dataloader))

    def evaluate(self, dataloader: DataLoader) -> float:
        """Execute evaluation loop without gradient computation.

        Args:
            dataloader: PyTorch DataLoader supplying evaluation batches.

        Returns:
            Mean evaluation loss across all batches.
        """
        self.model.eval()
        total_loss = 0.0

        with torch.no_grad():
            for inputs, targets in dataloader:
                inputs, targets = inputs.to(self.device), targets.to(self.device)
                outputs = self.model(inputs)
                loss = self.criterion(outputs, targets)
                total_loss += loss.item()

        return total_loss / max(1, len(dataloader))
Methods:
__init__
__init__(model, optimizer, criterion, device='cpu')

Initialize Model Trainer.

Parameters:

Name Type Description Default
model Module

PyTorch neural network model instance.

required
optimizer Optimizer

PyTorch optimizer instance.

required
criterion Module

PyTorch loss function instance.

required
device str | device

Compute device ('cuda', 'mps', or 'cpu').

'cpu'
Source code in src/tensoris/backend/trainers/trainer.py
def __init__(
    self,
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: str | torch.device = "cpu",
) -> None:
    """Initialize Model Trainer.

    Args:
        model: PyTorch neural network model instance.
        optimizer: PyTorch optimizer instance.
        criterion: PyTorch loss function instance.
        device: Compute device ('cuda', 'mps', or 'cpu').
    """
    self.model = model.to(device)
    self.optimizer = optimizer
    self.criterion = criterion
    self.device = device
evaluate
evaluate(dataloader)

Execute evaluation loop without gradient computation.

Parameters:

Name Type Description Default
dataloader DataLoader

PyTorch DataLoader supplying evaluation batches.

required

Returns:

Type Description
float

Mean evaluation loss across all batches.

Source code in src/tensoris/backend/trainers/trainer.py
def evaluate(self, dataloader: DataLoader) -> float:
    """Execute evaluation loop without gradient computation.

    Args:
        dataloader: PyTorch DataLoader supplying evaluation batches.

    Returns:
        Mean evaluation loss across all batches.
    """
    self.model.eval()
    total_loss = 0.0

    with torch.no_grad():
        for inputs, targets in dataloader:
            inputs, targets = inputs.to(self.device), targets.to(self.device)
            outputs = self.model(inputs)
            loss = self.criterion(outputs, targets)
            total_loss += loss.item()

    return total_loss / max(1, len(dataloader))
train_epoch
train_epoch(dataloader)

Execute a single training epoch optimization loop.

Parameters:

Name Type Description Default
dataloader DataLoader

PyTorch DataLoader supplying training batches.

required

Returns:

Type Description
float

Mean training loss across all batches.

Source code in src/tensoris/backend/trainers/trainer.py
def train_epoch(self, dataloader: DataLoader) -> float:
    """Execute a single training epoch optimization loop.

    Args:
        dataloader: PyTorch DataLoader supplying training batches.

    Returns:
        Mean training loss across all batches.
    """
    self.model.train()
    total_loss = 0.0

    for inputs, targets in dataloader:
        inputs, targets = inputs.to(self.device), targets.to(self.device)

        self.optimizer.zero_grad()
        outputs = self.model(inputs)
        loss = self.criterion(outputs, targets)
        loss.backward()
        self.optimizer.step()

        total_loss += loss.item()

    return total_loss / max(1, len(dataloader))