srforge 0.9.0
Super‑resolution research framework for PyTorch with a focus on simplicity and flexibility using config files.
Requires Python: >=3.10

# 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] --> Loss
```

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 **Loss** function evaluates the prediction against the target

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](https://pytorch.org/get-started/locally/) to pick the right build for your OS and GPU.

```bash
# 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

```bash
pip install srforge
```

### Graph neural network support (optional)

For graph-based models (e.g., MagNAt), install [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html) first, then the graph extra:

```bash
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

```bash
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:

```bash
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

```python
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](https://tarasiewicztomasz.gitlab.io/sr-forge/).

## For Developers

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

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

Run the test suite:

```bash
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](https://gitlab.com/tarasiewicztomasz/sr-forge/-/issues)
- Submit merge requests
- Share your models and experiments