SR-Forge

Structured Research Framework for Organized Research & Guided Experiments

SR-Forge is a modular, config-driven PyTorch framework for deep learning research. It handles the repetitive plumbing — data routing, component wiring, configuration management — so you can focus on what matters: your models, your data, and your experiments.

How It Works

Dataset --> [Entry] --> Transforms --> [Entry] --> Model --> [Entry] --> Metric

1. A Dataset loads raw data and wraps it in an Entry — a dictionary-like container that carries tensors, metadata, and any other fields through the pipeline
2. Transforms preprocess the data (normalize, augment, reshape)
3. A Model runs the neural network computation
4. Results are written back to the Entry
5. A Metric evaluates the prediction against the target — as a training loss or as an evaluation score

Every component reads from and writes to Entry objects. This uniform interface is what makes everything interchangeable — swap any component and the rest of the pipeline doesn't change.

Key Features

• Entry-based data flow — All data lives in Entry objects. Every component reads from Entry fields and writes back to them, so components are interchangeable without glue code.
• IO binding — Components declare abstract port names ("I need an image") that get mapped to concrete Entry fields ("read from input_rgb"). Write a component once, reuse it with any data layout.
• Pipeline composition — Chain models and transforms with a simple arrow syntax: image -> encoder -> features -> decoder -> output.
• Configuration over code — Define entire experiments in YAML. Reproduce any experiment by sharing a config file.
• Built-in models — FSRCNN, DSen2, RAMS, TR-MISR, MagNAt, plus a registry for custom architectures.
• Unified metrics — L1, L2, SSIM, LPIPS, schedulable loss combiners, and straightforward logging.

Installation

Prerequisites

Install PyTorch before installing SR-Forge. Follow the official instructions at pytorch.org to pick the right build for your OS and GPU.

# Example: CUDA 12.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

# Example: CPU only
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

Install SR-Forge

pip install srforge

Graph neural network support (optional)

For graph-based models (e.g., MagNAt), install PyTorch Geometric first, then the graph extra:

pip install torch-geometric
pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv \
    -f https://data.pyg.org/whl/torch-2.7.0+cu128.html

pip install srforge[graph]

Quick Start

mkdir my-experiment && cd my-experiment
srforge init

This generates a complete training script and config file:

my-experiment/
├── train.py              # Complete training script
└── configs/
    └── train-cfg.yaml    # Sample config with all settings

Edit configs/train-cfg.yaml to point at your data, then run:

python train.py

The generated files are a starting point — modify them to fit your workflow. The config supports multi-GPU, mixed precision, dataset caching, W&B tracking, loss scheduling, and checkpointing out of the box.

A Taste of SR-Forge

from srforge.models import Model
from srforge.data import Entry
import torch

class Upscaler(Model):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Conv2d(3, 3, 3, padding=1)

    def _forward(self, image):
        return self.net(image)

model = Upscaler()
model.set_io({"inputs": {"image": "input"}, "outputs": "prediction"})

entry = Entry({"input": torch.randn(1, 3, 64, 64)})
result = model(entry)
print(result.prediction.shape)  # torch.Size([1, 3, 64, 64])

The model reads from entry["input"], runs the network, and stores the result in entry["prediction"].

Documentation

Full documentation is available at tarasiewicztomasz.gitlab.io/sr-forge.

For Developers

Install PyTorch and PyG with CUDA wheels first (see above), then clone and install in editable mode:

git clone https://gitlab.com/tarasiewicztomasz/sr-forge.git
cd sr-forge
pip install -e ".[dev,graph]"

Run the test suite:

pytest tests/ -v

Some tests require optional dependencies (e.g., torch_geometric). These are automatically skipped if the dependency is missing.

Contributing

SR-Forge is under active development. Contributions welcome!

• Report issues on GitLab
• Submit merge requests
• Share your models and experiments