epoch-engine 0.1.0__tar.gz

epoch_engine-0.1.0/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Sergey Polivin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

epoch_engine-0.1.0/PKG-INFO

@@ -0,0 +1,68 @@
+Metadata-Version: 2.4
+Name: epoch-engine
+Version: 0.1.0
+Summary: Trainer and evaluator for PyTorch models with a focus on simplicity and flexibility.
+Author-email: Sergey Polivin <s.polivin@gmail.com>
+License: MIT
+Keywords: pytorch,deep-learning
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Development Status :: 3 - Alpha
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: torch==2.5.0
+Requires-Dist: torchvision==0.20.0
+Requires-Dist: tqdm==4.67.0
+Provides-Extra: build
+Requires-Dist: setuptools; extra == "build"
+Requires-Dist: wheel; extra == "build"
+Requires-Dist: build; extra == "build"
+Requires-Dist: twine; extra == "build"
+Provides-Extra: linters
+Requires-Dist: black; extra == "linters"
+Requires-Dist: isort; extra == "linters"
+Dynamic: license-file
+
+# Epoch Engine - Python Library for training PyTorch models
+
+This project represents my attempt to come up with a convenient way to train neural nets coded in Torch. While being aware of already existing libraries for training PyTorch models (e.g. PyTorch Lightning), my idea here is to make training of the models more visual and understandable as to what is going on during training.
+
+The project is currently in its raw form, more changes expected.
+
+## Features
+
+* TQDM-Progress bar support for both training and validation loops
+* Intemediate metrics computations after each forward pass (currently it is based on computing loss and accuracy only)
+* Saving/loading checkpoints from/into Trainer directly without having to touch model, optimizer or scheduler separately
+* Resuming training from the loaded checkpoint with epoch number being remembered automatically to avoid having to remember from which epoch the training originally started
+* Ready-to-use neural net architectures coded from scratch (currently only 4-layer Encoder-Decoder and ResNet with 20 layers architectures are available)
+
+## Installation
+
+After cloning this repo, the package can be installed in the development mode as follows:
+
+```bash
+# Installing the main package
+pip install -e .
+
+# Installing additional optional dependencies
+pip install -e .[build,linters]
+```
+
+## Python API
+
+The basics of the developed API are presented in the [test script](./run_trainer.py) I built. It can be run for instance as follows:
+
+```bash
+python run_trainer.py --model=resnet --epochs=3 --batch-size=16
+```
+> The training will be launched on the device automatically derived based on the CUDA availability and the final training checkpoint will be saved in `checkpoints` directory.
+
+One can also resume the training from the saved checkpoint:
+
+```bash
+python run_trainer.py --model=resnet --epochs=4 --resume-training=True --ckpt-path=checkpoints/ckpt_3.pt
+```
+> The training will be resumed from the loaded checkpoint with TQDM-progress bar showing the next training epoch.

epoch_engine-0.1.0/README.md

@@ -0,0 +1,41 @@
+# Epoch Engine - Python Library for training PyTorch models
+
+This project represents my attempt to come up with a convenient way to train neural nets coded in Torch. While being aware of already existing libraries for training PyTorch models (e.g. PyTorch Lightning), my idea here is to make training of the models more visual and understandable as to what is going on during training.
+
+The project is currently in its raw form, more changes expected.
+
+## Features
+
+* TQDM-Progress bar support for both training and validation loops
+* Intemediate metrics computations after each forward pass (currently it is based on computing loss and accuracy only)
+* Saving/loading checkpoints from/into Trainer directly without having to touch model, optimizer or scheduler separately
+* Resuming training from the loaded checkpoint with epoch number being remembered automatically to avoid having to remember from which epoch the training originally started
+* Ready-to-use neural net architectures coded from scratch (currently only 4-layer Encoder-Decoder and ResNet with 20 layers architectures are available)
+
+## Installation
+
+After cloning this repo, the package can be installed in the development mode as follows:
+
+```bash
+# Installing the main package
+pip install -e .
+
+# Installing additional optional dependencies
+pip install -e .[build,linters]
+```
+
+## Python API
+
+The basics of the developed API are presented in the [test script](./run_trainer.py) I built. It can be run for instance as follows:
+
+```bash
+python run_trainer.py --model=resnet --epochs=3 --batch-size=16
+```
+> The training will be launched on the device automatically derived based on the CUDA availability and the final training checkpoint will be saved in `checkpoints` directory.
+
+One can also resume the training from the saved checkpoint:
+
+```bash
+python run_trainer.py --model=resnet --epochs=4 --resume-training=True --ckpt-path=checkpoints/ckpt_3.pt
+```
+> The training will be resumed from the loaded checkpoint with TQDM-progress bar showing the next training epoch.

epoch_engine-0.1.0/epoch_engine/core/checkpoint_handler.py

@@ -0,0 +1,40 @@
+"""Module for handling saving/loading of checkpoints."""
+
+# Author: Sergey Polivin <s.polivin@gmail.com>
+# License: MIT License
+
+import torch
+
+
+class CheckpointHandler:
+    """Class for saving and loading trainer checkpoints."""
+
+    def save_checkpoint(self, trainer, path: str) -> None:
+        """Saves the checkpoint: last trained epoch, model state and optimizer state.
+
+        Args:
+            trainer: Trainer instance.
+            path (str): Path of the checkpoint file to which data are to be saved.
+        """
+        torch.save(
+            {
+                "epoch": trainer.last_epoch,
+                "model_state_dict": trainer.model.state_dict(),
+                "optimizer_state_dict": trainer.optimizer.state_dict(),
+            },
+            path,
+        )
+
+    def load_checkpoint(self, trainer, path: str) -> None:
+        """Loads the checkpoint: last trained epoch, model state and optimizer state.
+
+        Args:
+            trainer: Trainer instance.
+            path (str): Path of the checkpoint file from which data are to be loaded.
+        """
+        # Loading the checkpoint file
+        checkpoint = torch.load(path, weights_only=True)
+        # Loading the data into the Trainer instance attributes
+        trainer.last_epoch = checkpoint["epoch"]
+        trainer.model.load_state_dict(checkpoint["model_state_dict"])
+        trainer.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

epoch_engine-0.1.0/epoch_engine/core/metrics_tracker.py

@@ -0,0 +1,100 @@
+"""Module for handling tracking and computation of metrics."""
+
+# Author: Sergey Polivin <s.polivin@gmail.com>
+# License: MIT License
+
+
+class MetricsTracker:
+    """Class for keeping track of metrics."""
+
+    def __init__(self) -> None:
+        """Initializes a class instance.
+
+        Attributes:
+            metrics (dict[str, list[float]]): Batch-level metrics in a dict format.
+            total_correct_train (int): Number of correct predictions in train set.
+            total_samples_train (int): Total number of examples in train set.
+            total_correct_valid (int): Number of correct predictions in valid set.
+            total_samples_valid (int): Total number of examples in valid set.
+        """
+        self.metrics = {}
+        self.total_correct_train = 0
+        self.total_samples_train = 0
+        self.total_correct_valid = 0
+        self.total_samples_valid = 0
+
+    def reset(self) -> None:
+        """Resets all metrics and counters."""
+        self.metrics = {}
+        self.total_correct_train = 0
+        self.total_samples_train = 0
+        self.total_correct_valid = 0
+        self.total_samples_valid = 0
+
+    def update(self, name: str, value: float) -> None:
+        """Updates a specific metric.
+
+        Args:
+            name (str): Metric name.
+            value (float): Value of a metric.
+        """
+        if name not in self.metrics:
+            self.metrics[name] = []
+        self.metrics[name].append(value)
+
+    def update_accuracy(self, correct: int, total: int, split: str) -> None:
+        """Updates accuracy counters.
+
+        Args:
+            correct (int): Number of correct predictions.
+            total (int): Number of total examples.
+            split (str): Train or valid split indicator.
+
+        Raises:
+            ValueError: Error in case some other split name is passed.
+        """
+        if split == "train":
+            self.total_correct_train += correct
+            self.total_samples_train += total
+        elif split == "valid":
+            self.total_correct_valid += correct
+            self.total_samples_valid += total
+        else:
+            raise ValueError
+
+    def compute_accuracy(self, split: str) -> float:
+        """Computes accuracy across all batches.
+
+        Args:
+            split (str): Train or valid split indicator.
+
+        Raises:
+            ValueError: Error in case some other split name is passed.
+
+        Returns:
+            float: Accuracy score.
+        """
+        if self.total_samples_train == 0 or self.total_samples_valid == 0:
+            return 0.0
+        if split == "train":
+            return self.total_correct_train / self.total_samples_train
+        elif split == "valid":
+            return self.total_correct_valid / self.total_samples_valid
+        else:
+            raise ValueError
+
+    def get_all_metrics(self) -> dict[str, float]:
+        """Computes all metrics as averages.
+
+        Returns:
+            dict[str, float]: Epoch-level metrics.
+        """
+        # Aggregating loss across batches
+        metrics = {
+            name: sum(values) / len(values) for name, values in self.metrics.items()
+        }
+        # Computing accuracy scores
+        metrics["accuracy/train"] = self.compute_accuracy(split="train")
+        metrics["accuracy/valid"] = self.compute_accuracy(split="valid")
+
+        return metrics

epoch_engine-0.1.0/epoch_engine/core/trainer.py

@@ -0,0 +1,294 @@
+"""Module for Trainer functionality."""
+
+# Author: Sergey Polivin <s.polivin@gmail.com>
+# License: MIT License
+
+from typing import Any, TypeAlias
+
+import torch
+from tqdm import tqdm
+
+from .checkpoint_handler import CheckpointHandler
+from .metrics_tracker import MetricsTracker
+
+TorchDataloader: TypeAlias = torch.utils.data.dataloader.DataLoader
+
+
+class Trainer:
+    """Trainer of PyTorch models."""
+
+    def __init__(
+        self,
+        model: Any,
+        criterion: Any,
+        optimizer: Any,
+        train_loader: TorchDataloader,
+        valid_loader: TorchDataloader,
+        train_on: str = "auto",
+    ) -> None:
+        """Initializes a class instance.
+
+        Args:
+            model (Any): PyTorch model instance.
+            criterion (Any): PyTorch loss function instance.
+            optimizer (Any): PyTorch optimizer instance.
+            train_loader (TorchDataloader): Torch Dataloader for training set.
+            valid_loader (TorchDataloader): Torch Dataloader for validation set.
+            train_on (str, optional): Indicator of CPU- or GPU-based training. Defaults to "auto".
+
+        Attributes:
+            model (Any): PyTorch model instance.
+            criterion (Any): PyTorch loss function instance.
+            optimizer (Any): PyTorch optimizer instance.
+            train_loader (TorchDataloader): Torch Dataloader for training set.
+            valid_loader (TorchDataloader): Torch Dataloader for validation set.
+            device (torch.device): Device to be used to train the model.
+            scheduler (Any): Instance of PyTorch scheduler for learning rate. Defaults to None.
+            scheduler_level (bool): Level on which to apply scheduler. Defaults to None.
+            last_epoch (int): Last epoch when training has been successfully finished. Defaults to 0.
+            metrics_tracker (MetricsTracker): Class for handling computing metrics.
+            ckpt_handler (CheckpointHandler): Class for saving/loading checkpoints.
+        """
+        # Using the CPU or GPU device or inferring one
+        self.device = (
+            torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if train_on == "auto"
+            else torch.device(train_on)
+        )
+        # Moving model to the device specified
+        self.model = model.to(self.device)
+        # Reinitializing the optimizer with model parameters having been moved to device
+        self.optimizer = type(optimizer)(self.model.parameters(), **optimizer.defaults)
+
+        # Setting loss function
+        self.criterion = criterion
+
+        # Setting dataloaders for training/validation sets
+        self.train_loader = train_loader
+        self.valid_loader = valid_loader
+
+        # Setting scheduler-related attributes
+        self.scheduler = None
+        self.scheduler_level = None
+
+        # Setting epoch indicator
+        self.last_epoch = 0
+
+        # Setting tracker for metrics and checkpoints handler
+        self.metrics_tracker = MetricsTracker()
+        self.ckpt_handler = CheckpointHandler()
+
+    def set_scheduler(
+        self,
+        scheduler_class: Any,
+        scheduler_params: dict[str, Any],
+        level: str = "epoch",
+    ) -> None:
+        """Sets and initializes a scheduler.
+
+        Args:
+            scheduler_class (Any): Scheduler class from PyTorch
+            scheduler_params (dict[str, Any]): Scheduler parameters.
+            level (str, optional): Level on which to apply scheduler. Defaults to "epoch".
+
+        Raises:
+            ValueError: Error raised if the specified level name is unexpected.
+        """
+        # Setting scheduler level on batch- or epoch-level
+        if level not in ("epoch", "batch"):
+            raise ValueError("Invalid value")
+        else:
+            self.scheduler_level = level
+        # Initializing a scheduler
+        self.scheduler = scheduler_class(self.optimizer, **scheduler_params)
+
+    def reset_scheduler(self) -> None:
+        """Resetting and removing scheduler."""
+        self.scheduler = None
+        self.scheduler_level = None
+
+    def __call__(
+        self, x_batch: torch.Tensor, y_batch: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Computes loss and model outputs.
+
+        Args:
+            x_batch (torch.Tensor): Batch of examples.
+            y_batch (torch.Tensor): Batch of labels.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: Tuple containing value of loss function and output tensor.
+        """
+        outputs = self.model(x_batch)
+        loss = self.criterion(outputs, y_batch.long())
+
+        return loss, outputs
+
+    def run(
+        self,
+        epochs: int,
+        seed: int = 42,
+        enable_tqdm: bool = True,
+    ) -> None:
+        """Launches the trainer.
+
+        Args:
+            epochs (int): Number of epochs to train.
+            seed (int, optional): Random number generator seed. Defaults to 42.
+            enable_tqdm (bool, optional): Indicator to turn on tqdm progress bar. Defaults to True.
+        """
+        # Setting seed
+        if seed:
+            torch.manual_seed(seed)
+        # Computing the number of total epochs to be trained (useful in case of additional trainings)
+        total_epochs = epochs + self.last_epoch
+        # Running training/validation loops
+        for epoch in range(self.last_epoch, total_epochs):
+            # Processing all training batches
+            self._train_one_epoch(
+                epoch=epoch, epochs=total_epochs, enable_tqdm=enable_tqdm
+            )
+            # Processing all validation batches
+            self._validate_one_epoch(
+                epoch=epoch, epochs=total_epochs, enable_tqdm=enable_tqdm
+            )
+            # Incrementing last epoch after successful training/validation
+            self.last_epoch = epoch + 1
+            # Resetting metrics counters after successful training/validation
+            self.metrics_tracker.reset()
+
+    def _train_one_epoch(self, epoch: int, epochs: int, enable_tqdm: bool) -> None:
+        """Processes all training batches for one epoch.
+
+        Args:
+            epoch (int): Current epoch.
+            epochs (int): Total number of epochs.
+            enable_tqdm (bool): Indicator to turn on tqdm progress bar.
+        """
+
+        with tqdm(
+            total=len(self.train_loader),
+            desc=f"Epoch {epoch + 1}/{epochs} [Training]",
+            position=0,
+            leave=True,
+            unit="batches",
+            disable=not enable_tqdm,
+        ) as pbar:
+            # Setting the learning rate in progress bar
+            pbar.set_postfix({"lr": self.optimizer.param_groups[0]["lr"]})
+            # Setting model in training mode
+            self.model.train()
+            # Processing all training batches
+            for batch in self.train_loader:
+                self._process_batch(batch=batch, split="train")
+                pbar.update(1)
+        # Adjusting learning rate if set
+        if self.scheduler and self.scheduler_level == "epoch":
+            self.scheduler.step()
+
+    def _validate_one_epoch(self, epoch: int, epochs: int, enable_tqdm: bool) -> None:
+        """Processes all validation batches for one epoch.
+
+        Args:
+            epoch (int): Current epoch.
+            epochs (int): Total number of epochs.
+            enable_tqdm (bool): Indicator to turn on tqdm progress bar.
+        """
+        with tqdm(
+            total=len(self.valid_loader),
+            desc=f"Epoch {epoch + 1}/{epochs} [Validation]",
+            position=0,
+            leave=True,
+            unit="batches",
+            disable=not enable_tqdm,
+        ) as pbar:
+            # Setting model to evalutation model
+            self.model.eval()
+            # Processing all validation batches (grad compute turned off)
+            with torch.no_grad():
+                for batch in self.valid_loader:
+                    self._process_batch(batch=batch, split="valid")
+                    pbar.update(1)
+
+            # Aggregating batch-level metrics (loss, accuracy) to epoch-level
+            self.metrics = self.metrics_tracker.get_all_metrics()
+            # Setting the epoch-level stats to the progress bar
+            pbar.set_postfix(
+                {
+                    "loss": (
+                        round(self.metrics["loss/train"], 4),
+                        round(self.metrics["loss/valid"], 4),
+                    ),
+                    "acc": (
+                        round(self.metrics["accuracy/train"], 4),
+                        round(self.metrics["accuracy/valid"], 4),
+                    ),
+                }
+            )
+
+    def _process_batch(self, batch: list[torch.Tensor], split: str) -> None:
+        """Backpropagates model during training and collects batch-level metrics.
+
+        Args:
+            batch (list[torch.Tensor]): Examples and labels.
+            split (str): Indicator of training or validation set.
+        """
+
+        # Moving examples and labels batches to device chosen
+        x_batch, y_batch = batch
+        x_batch, y_batch = x_batch.to(self.device), y_batch.to(self.device)
+        # Computing loss and output tensor for the batch
+        loss, outputs = self(x_batch=x_batch, y_batch=y_batch)
+
+        # If training set, backpropagate
+        if split == "train":
+            loss.backward()
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+            # Adjusting learning rate at batch-level
+            if self.scheduler and self.scheduler_level == "batch":
+                self.scheduler.step()
+
+        # Collecting batch-level metrics
+        self._collect_metrics(loss=loss, outputs=outputs, labels=y_batch, split=split)
+
+    def _collect_metrics(
+        self,
+        loss: torch.Tensor,
+        outputs: torch.Tensor,
+        labels: torch.Tensor,
+        split: str,
+    ) -> None:
+        """Records loss and accuracy.
+
+        Args:
+            loss (torch.Tensor): Value of loss function.
+            outputs (torch.Tensor): Output tensor.
+            labels (torch.Tensor): Batch of labels.
+            split (str): Indicator of training or validation set.
+        """
+        # Adding batch-level loss
+        loss = loss.item()
+        self.metrics_tracker.update(f"loss/{split}", loss)
+
+        # Adding batch-level accuracy
+        predictions = torch.argmax(outputs, dim=1)
+        correct = (predictions == labels).sum().cpu().item()
+        total = labels.size(0)
+        self.metrics_tracker.update_accuracy(correct, total, split=split)
+
+    def save_checkpoint(self, path: str) -> None:
+        """Saves trainer checkpoint.
+
+        Args:
+            path (str): Path to the checkpoint to be saved.
+        """
+        self.ckpt_handler.save_checkpoint(path=path, trainer=self)
+
+    def load_checkpoint(self, path: str) -> None:
+        """Loads trainer checkpoint.
+
+        Args:
+            path (str): Path to the checkpoint to be loaded.
+        """
+        self.ckpt_handler.load_checkpoint(path=path, trainer=self)

epoch_engine-0.1.0/epoch_engine/models/architectures.py

@@ -0,0 +1,313 @@
+"""Module containing building blocks of models and functions to build neural networks."""
+
+# Author: Sergey Polivin <s.polivin@gmail.com>
+# License: MIT License
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Encoder(nn.Module):
+    """Builds a series of convolutional blocks."""
+
+    def __init__(self, encoder_channels: tuple[int]) -> None:
+        """Initializes a class instance.
+
+        Args:
+            encoder_channels (tuple[int]): Tuple of convolution channels.
+        """
+        super().__init__()
+        # Creating a series of encoder blocks in accordance with channels in `encoder_channels`
+        encoder_blocks = [
+            self._make_encoder_block(in_channel, out_channel)
+            for in_channel, out_channel in zip(encoder_channels, encoder_channels[1:])
+        ]
+        # Sequentially connecting the generated convolution blocks
+        self.encoder_blocks = nn.Sequential(*encoder_blocks)
+
+    def _make_encoder_block(self, in_channels: int, out_channels: int) -> nn.Sequential:
+        """Creates a convolution block with ReLU activation and MaxPooling.
+
+        Args:
+            in_channels (int): Number of input channels for convolution.
+            out_channels (int): Number of output channels for convolution.
+
+        Returns:
+            nn.Sequential: Convolution block of sequentially connected layers.
+        """
+        return nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=3,
+                padding=1,
+            ),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=3),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Makes a forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor for Encoder.
+
+        Returns:
+            torch.Tensor: Output tensor of Encoder.
+        """
+        return self.encoder_blocks(x)
+
+
+class Decoder(nn.Module):
+    """Builds a series of feedforward blocks."""
+
+    def __init__(self, decoder_features: tuple[int], num_labels: int) -> None:
+        """Initializes a class instance.
+
+        Args:
+            decoder_features (tuple[int]): Tuple of features in the Decoder.
+            num_labels (int): Number of output labels in the last layer.
+        """
+        super().__init__()
+        # Creating a series of decoder blocks in accordance with features in `decoder_features`
+        decoder_blocks = [
+            self._make_decoder_block(in_feature, out_feature)
+            for in_feature, out_feature in zip(decoder_features, decoder_features[1:])
+        ]
+        # Sequentially connecting the generated feedforward blocks
+        self.decoder_blocks = nn.Sequential(*decoder_blocks)
+        # Creating the last layer of the Decoder
+        self.last = nn.Linear(decoder_features[-1], num_labels)
+
+    def _make_decoder_block(self, in_features: int, out_features: int) -> nn.Sequential:
+        """Creates a fully connected linear layer with Sigmoid activation.
+
+        Args:
+            in_features (int): Number of input features of the Decoder.
+            out_features (int): Number of output features of the Decoder.
+
+        Returns:
+            nn.Sequential: Decoder block of sequentially connected layers.
+        """
+        return nn.Sequential(
+            nn.Linear(in_features, out_features),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Makes a forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor for Decoder.
+
+        Returns:
+            torch.Tensor: Output tensor of Decoder.
+        """
+        x = self.decoder_blocks(x)
+        x = self.last(x)
+
+        return x
+
+
+class EDNet(nn.Module):
+    """Joins Encoder with Decoder to form a CNN."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        encoder_channels: tuple[int],
+        decoder_features: tuple[int],
+        num_labels: int,
+    ) -> None:
+        """Initializes a class instance.
+
+        Args:
+            in_channels (int): Number of input channels for the image.
+            encoder_channels (tuple[int]): Tuple of channels for the Encoder.
+            decoder_features (tuple[int]): Tuple of channels for the Decoder.
+            num_labels (int): Number of ouput labels.
+        """
+        super().__init__()
+        # Setting up Encoder block
+        self.encoder_channels = [in_channels, *encoder_channels]
+        self.encoder = Encoder(self.encoder_channels)
+        # Setting up Decoder block
+        self.decoder_features = decoder_features
+        self.decoder = Decoder(self.decoder_features, num_labels)
+        # Flatten layer
+        self.flatten = nn.Flatten()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Makes a forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor for a CNN.
+
+        Returns:
+            torch.Tensor: Output tensor of a CNN.
+        """
+        x = self.encoder(x)
+        x = self.flatten(x)
+        x = self.decoder(x)
+
+        return x
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1) -> None:
+        """Builds a block of ResNet layer with 2 inner blocks (CONV + BN + SKIP CONNECTION).
+
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels.
+            stride (int, optional): Value of stride. Defaults to 1.
+        """
+        super().__init__()
+        # BLOCK 1 (CONV + BN) #################################################
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm2d(out_channels)
+
+        # BLOCK 2 (CONV + BN) #################################################
+        self.conv2 = nn.Conv2d(
+            out_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=False,
+        )
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        # SKIP CONNECTION #################################################
+        self.shortcut = nn.Sequential()
+        # Specifying condition for applying residual connection
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=1,
+                    stride=stride,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(out_channels),
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Makes a forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor.
+        """
+        # Output for BLOCK 1
+        out = F.relu(self.bn1(self.conv1(x)))
+        # Output for BLOCK 2
+        out = self.bn2(self.conv2(out))
+        # Adding up SKIP CONNECTION if condition met
+        out += self.shortcut(x)
+        # Output of RESNET BLOCK
+        return F.relu(out)
+
+
+class ResNet(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        block: BasicBlock,
+        num_blocks: list[int],
+        num_classes: int = 10,
+    ) -> None:
+        """Builds a ResNet model with 3 layers.
+
+        Args:
+            in_channels (int): Number of input channels.
+            block (BasicBlock): Block of ResNet layer.
+            num_blocks (list[int]): Number of blocks inside Block of ResNet layer.
+            num_classes (int, optional): Number of output classes. Defaults to 10.
+        """
+        super().__init__()
+        # Specifying a number of initial input channels for ResNet block
+        self.in_block_channels = 16
+
+        # LAYER 0 (CONV + BN) #################################################################################
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            self.in_block_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm2d(self.in_block_channels)
+
+        # LAYERS 1, 2, 3 (`num_blocks` blocks with 2 CONV/BN BLOCKS for 1 layer) ##############################
+        self.layer1 = self._make_resnet_layer(
+            block=block, out_channels=16, num_blocks=num_blocks[0], stride=1
+        )
+        self.layer2 = self._make_resnet_layer(
+            block=block, out_channels=32, num_blocks=num_blocks[1], stride=2
+        )
+        self.layer3 = self._make_resnet_layer(
+            block=block, out_channels=64, num_blocks=num_blocks[2], stride=2
+        )
+
+        # GLOBAL POOLING LAYER ################################################################################
+        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+
+        # FC LAYER ############################################################################################
+        self.fc = nn.Linear(64, num_classes)
+
+    def _make_resnet_layer(
+        self,
+        block: BasicBlock,
+        out_channels: int,
+        num_blocks: int,
+        stride: int,
+    ) -> nn.Sequential:
+        """Creates a ResNet layer by stacking up BasicBlock-s.
+
+        Args:
+            block (BasicBlock): Block inside ResNet layer.
+            out_channels (int): Number of output channels.
+            num_blocks (int): Number of blocks inside ResNet layer.
+            stride (int): Value of stride.
+
+        Returns:
+            nn.Sequential: Layer of ResNet with blocks added inside.
+        """
+        # Specifying strides to be used for each BasicBlock
+        strides = [stride] + [1] * (num_blocks - 1)
+        # Stacking up ResNet blocks
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_block_channels, out_channels, stride))
+            self.in_block_channels = out_channels
+        return nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Makes a forward pass.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor.
+        """
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.avg_pool(out)
+        out = torch.flatten(out, 1)
+        return self.fc(out)

epoch_engine-0.1.0/epoch_engine.egg-info/PKG-INFO

@@ -0,0 +1,68 @@
+Metadata-Version: 2.4
+Name: epoch-engine
+Version: 0.1.0
+Summary: Trainer and evaluator for PyTorch models with a focus on simplicity and flexibility.
+Author-email: Sergey Polivin <s.polivin@gmail.com>
+License: MIT
+Keywords: pytorch,deep-learning
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Development Status :: 3 - Alpha
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: torch==2.5.0
+Requires-Dist: torchvision==0.20.0
+Requires-Dist: tqdm==4.67.0
+Provides-Extra: build
+Requires-Dist: setuptools; extra == "build"
+Requires-Dist: wheel; extra == "build"
+Requires-Dist: build; extra == "build"
+Requires-Dist: twine; extra == "build"
+Provides-Extra: linters
+Requires-Dist: black; extra == "linters"
+Requires-Dist: isort; extra == "linters"
+Dynamic: license-file
+
+# Epoch Engine - Python Library for training PyTorch models
+
+This project represents my attempt to come up with a convenient way to train neural nets coded in Torch. While being aware of already existing libraries for training PyTorch models (e.g. PyTorch Lightning), my idea here is to make training of the models more visual and understandable as to what is going on during training.
+
+The project is currently in its raw form, more changes expected.
+
+## Features
+
+* TQDM-Progress bar support for both training and validation loops
+* Intemediate metrics computations after each forward pass (currently it is based on computing loss and accuracy only)
+* Saving/loading checkpoints from/into Trainer directly without having to touch model, optimizer or scheduler separately
+* Resuming training from the loaded checkpoint with epoch number being remembered automatically to avoid having to remember from which epoch the training originally started
+* Ready-to-use neural net architectures coded from scratch (currently only 4-layer Encoder-Decoder and ResNet with 20 layers architectures are available)
+
+## Installation
+
+After cloning this repo, the package can be installed in the development mode as follows:
+
+```bash
+# Installing the main package
+pip install -e .
+
+# Installing additional optional dependencies
+pip install -e .[build,linters]
+```
+
+## Python API
+
+The basics of the developed API are presented in the [test script](./run_trainer.py) I built. It can be run for instance as follows:
+
+```bash
+python run_trainer.py --model=resnet --epochs=3 --batch-size=16
+```
+> The training will be launched on the device automatically derived based on the CUDA availability and the final training checkpoint will be saved in `checkpoints` directory.
+
+One can also resume the training from the saved checkpoint:
+
+```bash
+python run_trainer.py --model=resnet --epochs=4 --resume-training=True --ckpt-path=checkpoints/ckpt_3.pt
+```
+> The training will be resumed from the loaded checkpoint with TQDM-progress bar showing the next training epoch.

epoch_engine-0.1.0/epoch_engine.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+LICENSE.txt
+README.md
+pyproject.toml
+epoch_engine/__init__.py
+epoch_engine.egg-info/PKG-INFO
+epoch_engine.egg-info/SOURCES.txt
+epoch_engine.egg-info/dependency_links.txt
+epoch_engine.egg-info/requires.txt
+epoch_engine.egg-info/top_level.txt
+epoch_engine/core/__init__.py
+epoch_engine/core/checkpoint_handler.py
+epoch_engine/core/metrics_tracker.py
+epoch_engine/core/trainer.py
+epoch_engine/models/__init__.py
+epoch_engine/models/architectures.py

epoch_engine-0.1.0/epoch_engine.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

epoch_engine-0.1.0/epoch_engine.egg-info/requires.txt

@@ -0,0 +1,13 @@
+torch==2.5.0
+torchvision==0.20.0
+tqdm==4.67.0
+
+[build]
+setuptools
+wheel
+build
+twine
+
+[linters]
+black
+isort

epoch_engine-0.1.0/epoch_engine.egg-info/top_level.txt

@@ -0,0 +1 @@
+epoch_engine

epoch_engine-0.1.0/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "epoch-engine"
+version = "0.1.0"
+description = "Trainer and evaluator for PyTorch models with a focus on simplicity and flexibility."
+authors = [{name = "Sergey Polivin", email = "s.polivin@gmail.com"}]
+keywords = ["pytorch", "deep-learning"]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+  "Development Status :: 3 - Alpha",
+]
+readme = "README.md"
+requires-python = ">=3.8"
+license = {text = "MIT"}
+dependencies = [
+    "torch==2.5.0",
+    "torchvision==0.20.0",
+    "tqdm==4.67.0"
+]
+
+[project.optional-dependencies]
+build = ["setuptools", "wheel", "build", "twine"]
+linters = ["black", "isort"]
+
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["epoch_engine*"]

epoch_engine-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+