Quickstart
A base class for building image and volume learning apps with the WiDiT model and U-Net. WiDiTApp combines model creation, PyTorch data loaders, diffusion or supervised training, and Cluey command-line interfaces. It supports 2D images and 3D volumes.
Subclass WiDiTApp to supply your datasets and prediction workflow. The base
class provides training, but its datasets() and predict() methods raise
NotImplementedError until you implement them. This package does not install
a standalone command-line executable.
Installation
Requires Python 3.10–3.13. Install in a virtual environment using pip:
python -m pip install widitapp
Or install from the repository:
python -m pip install git+https://github.com/rbturnbull/WiDiTApp.git
The current project configuration fetches WiDiT from GitHub over SSH. Installing these dependencies requires Git and working GitHub SSH authentication.
Training currently requires a CUDA-capable GPU and a CUDA-enabled PyTorch installation; the training loop explicitly checks for CUDA. Model construction and CPU inference can be used independently of training.
Usage
Create a model
Build a small 2D regression model from Python:
import torch
from widitapp import WiDiTApp
app = WiDiTApp()
model = app.model(
dim=2,
use_diffusion=False,
hidden_size=64,
depth=1,
num_heads=4,
patch_size=2,
window_size=4,
use_flash_attention=False,
)
model.eval()
with torch.no_grad():
prediction = model(torch.zeros(1, 1, 32, 32))
Use dim=3 for volumes, or unet=True to build a U-Net. Inputs use
channel-first shapes: (N, C, H, W) for batched 2D images and
(N, C, D, H, W) for batched 3D volumes. Choose spatial sizes compatible with
the model’s patch and window settings.
Create your own training app
As in Supercat and Dosefusion, implement datasets() and decorate it with
Cluey’s @method. Return a pair of PyTorch datasets: training first, validation
second. The inherited train tool collects the options from model(),
dataloaders(), loss(), metrics(), and your datasets() method.
Save this example as myapp.py. It reads tensor pairs from train.pt and
validation.pt in a directory you provide:
from pathlib import Path
import torch
from cluey import method
from torch.utils.data import TensorDataset
from widitapp import WiDiTApp
class MyApp(WiDiTApp):
@method
def datasets(self, data_dir: Path = None, **kwargs) -> tuple:
"""Load paired input and target tensors."""
if data_dir is None:
raise ValueError("Provide data_dir (CLI: --data-dir).")
data_dir = Path(data_dir)
def load_split(name):
inputs, targets = torch.load(
data_dir / name, map_location="cpu", weights_only=True
)
return TensorDataset(inputs.float(), targets.float())
return load_split("train.pt"), load_split("validation.pt")
if __name__ == "__main__":
MyApp.tools()
Each file should contain a tuple (inputs, targets) saved with
torch.save((inputs, targets), path). For single-channel 2D data, both tensors
have shape (N, 1, H, W). Use separate samples for training and validation.
Preprocess and normalize your data consistently before saving it; WiDiTApp does
not read image formats or normalize intensities for you.
Inspect the generated options and train a small regression model:
python myapp.py train --help
python myapp.py train \
--data-dir /path/to/tensor-pairs \
--no-use-diffusion \
--dim 2 \
--hidden-size 64 \
--depth 1 \
--num-heads 4 \
--no-use-flash-attention \
--batch-size 4 \
--epochs 40 \
--results-dir results \
--run-name example
You can also call MyApp().train(data_dir=Path("/path/to/tensor-pairs"),
use_diffusion=False, ...) from Python with the corresponding keyword arguments.
Dataset and training modes
Each dataset item must return
(input, target)or(input, target, timestep). Individual images omit the batch dimension; the data loader adds it. Samples in a batch must have compatible shapes.With
use_diffusion=True(the default), the target is noised and the input is passed to the model asconditioned=input. The training loop samples diffusion timesteps; any dataset-provided timestep is replaced. WiDiT expects the conditioning input and target to have the same shape, so resize paired data to match before training.With
use_diffusion=False, the model predicts the target directly from the input using mean squared error. An optional dataset timestep is decremented by one before being passed to the model.Return
(training_dataset, None)to omit validation. The current training loop saves checkpoints during validation only, so provide a validation dataset when you need saved models.
Training data is shuffled; validation data is not. Both loaders retain the last
incomplete batch. Override dataloaders() if you need custom sampling or
collation.
Common options
Python parameter names use underscores; their CLI equivalents use hyphens,
for example batch_size becomes --batch-size. Boolean options also have a
negative form, such as --no-use-diffusion.
Parameter |
Default |
Purpose |
|---|---|---|
|
|
Spatial dimensionality. |
|
|
Conditional diffusion training or direct regression. |
|
|
Model input channels. |
|
Automatic |
Defaults to 2 for diffusion and 1 for regression. |
|
Automatic |
Defaults to the value of |
|
Empty |
Named WiDiT preset; inspect |
|
|
Transformer architecture when no preset is selected. |
|
|
Patch and attention window sizes. |
|
|
Select U-Net; configure with |
|
|
Data loader batch size and worker count. |
|
|
Training duration and AdamW learning rate. |
|
|
Log every this many successful training steps. |
|
|
Output root and experiment subdirectory name. |
|
|
Enable Weights & Biases logging; |
|
|
Load an existing model instead of constructing one. |
The default output-channel counts assume single-channel targets. For
multi-channel regression set out_channels to the target channel count;
diffusion with learned variance needs twice that count. A preset defines its
own spatial dimensionality and core architecture. Loading a checkpoint uses
the saved model configuration rather than the model-building options.
Custom losses
Override loss() with @method to configure the loss in a child app.
Cluey forwards its arguments from train() and exposes them as CLI options.
Without an override, regression uses MSE (or loss_fn="smoothl1"), and
diffusion retains its original noise-MSE plus variance objective.
For example, add a penalty for mismatched mean image intensity:
import torch
from cluey import method
from widitapp import WiDiTApp
class ImageLoss(torch.nn.Module):
def __init__(self, mean_weight, use_diffusion):
super().__init__()
self.mean_weight = mean_weight
self.use_diffusion = use_diffusion
def forward(self, prediction, target):
dims = tuple(range(1, prediction.ndim))
loss = self.mean_weight * (
prediction.mean(dims) - target.mean(dims)
).square().mean()
if not self.use_diffusion:
loss = loss + (prediction - target).square().mean()
return loss
class MyApp(WiDiTApp):
# Also implement datasets() for your application.
@method
def loss(self, mean_weight: float = 0.1, use_diffusion: bool = True):
return ImageLoss(mean_weight, use_diffusion)
Call MyApp().train(mean_weight=0.2, use_diffusion=False, ...) or use
--mean-weight 0.2 on the train command.
In regression, the returned module defines the entire objective and must return
a scalar. In diffusion, it receives the unclipped clean-image estimate and clean
target, and its result is added to the standard objective. Return a scalar batch
mean or a tensor of shape [N]; include component weights inside your module.
The estimate is computed from the sampled noisy target in the existing forward
pass, with gradients retained. It is not an image generated by a full sampling
loop. Estimates at high noise levels can have large errors, so choose weights
accordingly. Loss modules are moved to the training device; their parameters are
not optimized by the model optimizer.
The base loss() returns None for diffusion, adding no penalty. With this
default, validation still measures noise MSE only. With an extra diffusion loss,
validation and best-checkpoint selection use the full objective, including the
variance term and the extra image loss. Regression validation uses the selected
loss and also reports MSE and Smooth L1.
For finer control, call widitapp.training.train directly with a model and
data loaders. It additionally accepts precision (fp16 by default,
bf16, fp32, or no), loss_fn (mse, smoothl1/huber, or
a PyTorch loss module), diffusion_loss_fn (an optional extra image-loss
module), metrics (a dictionary of validation modules), wandb_config,
and wandb_log_artifacts. The app’s loss() result
is passed as loss_fn for regression or diffusion_loss_fn for diffusion.
Precision and the additional W&B configuration are not forwarded by the app.
Custom validation metrics
The app’s metrics() method returns a dictionary of names to PyTorch modules.
The default is {"mse": MSELoss(), "smoothl1": SmoothL1Loss()} for regression
and an empty dictionary for diffusion, preserving the existing defaults.
Extend it with Cluey’s @method("super") to inherit both metrics and options:
class MyApp(WiDiTApp):
@method("super")
def metrics(self, **kwargs):
metrics = super().metrics(**kwargs)
metrics["custom_error"] = MyExistingMetric()
return metrics
Each metric module accepts (prediction, target) and returns a scalar tensor
representing the batch mean. Modules run in evaluation mode without gradients
on the validation device. They should compute each batch independently; this
interface does not manage stateful update/compute/reset metrics.
For regression, prediction is the direct EMA model output. For diffusion,
it is the unclipped clean-image estimate at a sampled timestep, using the same
forward pass as the validation loss. It is not the result of a full sampling
loop. Adding metrics does not enable an auxiliary loss or change val/loss
or checkpoint selection. Non-finite metric values are reported without
discarding an otherwise finite validation loss.
Use names without the val/ prefix; loss is reserved. Values are averaged
across processes and then across validation batches, giving batches equal
weight, including an incomplete last batch. After each validation epoch,
metrics appear as val/custom_error in the training log and in W&B when
wandb=True. They are not evaluated during training steps.
Subclasses can add, replace, or remove entries in the inherited dictionary.
When calling the lower-level widitapp.training.train directly, metrics=None
selects the mode’s defaults and metrics={} disables extra metrics.
Checkpoints and prediction
Training maintains an exponential moving average (EMA) of model weights and evaluates that model after each epoch. With validation, outputs are written to:
results/<run_name>/
log.txt
checkpoints/
best.pt
best.pt stores the EMA model with the best validation loss. If the initial
validation loss is non-finite, a fallback checkpoint is saved until a finite
improvement is available. There is currently no separate final checkpoint or
optimizer-state resume: checkpoint=... initializes a new training run with
saved model weights.
Load a checkpoint with WiDiTApp().model(checkpoint=Path(".../best.pt")) or
widit.load_model(path). For regression, move the loaded model and inputs to
the same device, call model.eval(), and predict under torch.no_grad().
Diffusion inference requires an iterative sampling loop, rather than a single
forward pass; widitapp.diffusion.create_diffusion supplies the diffusion
process and sampling methods.
To expose prediction in your app, override predict() with @main from
Cluey. Implement input loading, preprocessing, checkpoint loading, regression
or diffusion inference, and output writing for your data format.
Package your command-line app
If your class lives in myapp/apps.py, add these entry points to your app’s
Poetry configuration and install that app with python -m pip install -e .:
[tool.poetry.scripts]
myapp = "myapp.apps:MyApp.main"
myapp-tools = "myapp.apps:MyApp.tools"
myapp invokes the method decorated with @main (your prediction method),
while myapp-tools train invokes inherited training. The training-only
example above can use just the myapp-tools entry point until prediction is
implemented. Decorate additional commands with Cluey’s @tool.
Example applications
The following projects demonstrate how to extend WiDiTApp:
Supercat (
supercat/apps.py): provides 2D and 3D super-resolution datasets, tiled prediction, and additional porosity tools. Itssupercat/pretrain.pyprovides further subclasses for image and video pretraining.Dosefusion (
dosefusion/apps.py): builds datasets from a CSV with partition selection and implements volume prediction with diffusion and fusion sampling options. Its data preparation and prediction code live indosefusion/data.pyanddosefusion/models.py.
In both projects, pyproject.toml registers the app’s .main and .tools
entry points. Use their dataset and prediction implementations as examples
when adapting the base class to your own files and task.