TorchExtraContext

English | 中文

A PyTorch context manager for collecting auxiliary losses, metrics, intermediate outputs, hooks, and shared state from nested torch.nn.Module objects without changing forward() signatures.

What It Does

In PyTorch, auxiliary losses and intermediate values are often produced deep inside nested torch.nn.Module objects. The traditional approach is to return them layer by layer:

def forward(self, x):
    x, loss_a = self.block_a(x)
    x, loss_b = self.block_b(x)
    logits = self.head(x)
    return logits, {"loss_a": loss_a, "loss_b": loss_b}

This approach couples auxiliary data collection to the return values of every module on the call path. Modules that do not create the auxiliary data still need to receive, merge, and return it.

Instead of modifying forward() signatures layer by layer, we prefer to register the data in the module where it is produced:

import torchextractx.keras_style  # Enable Keras-style API

def forward(self, x):
    x = self.layer1(x)
    self.add_loss("loss_a", compute_loss(x))
    return x

Then collect all losses in one line:

with ExtraContext(model) as ctx:
    output = model(x)
    all_losses = ctx.get_losses()  # Done

🚀 Installation

From PyPI

pip install torchextractx

From source

git clone https://github.com/5o1/TorchExtraContext.git
cd TorchExtraContext
pip install -e .

With dev tools

git clone https://github.com/5o1/TorchExtraContext.git
cd TorchExtraContext
pip install -e ".[dev]"  # pytest, black, mypy

📖 Quick Start

Usage

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchextractx import ExtraContext, get_context  # Step #1: Import the context helpers

class FeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 10)
    
    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        get_context(self).add_loss("feature_mean_loss", x.mean())  # Step #2: Register the loss where it is created
        
        x = self.fc2(x)
        return x


class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.extractor = FeatureExtractor()
        self.classifier = nn.Linear(10, 2)
    
    def forward(self, x):
        x = self.extractor(x)
        x = self.classifier(x)
        return x


model = Classifier()
optimizer = torch.optim.Adam(model.parameters())
x = torch.randn(32, 10)
targets = torch.randint(0, 2, (32,))

with ExtraContext(model) as ctx:  # Step #3: Create the collection context in the training step
    logits = model(x)
    main_loss = F.cross_entropy(logits, targets)
    extra_losses = ctx.get_losses()  # Step #4: Collect losses registered inside the model

    total_loss = main_loss
    for loss in extra_losses.values():
        total_loss = total_loss + 0.1 * loss

    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

Enable Keras Style

import torchextractx.keras_style  # Step #1: Enable torch.nn.Module.add_loss
from torchextractx import ExtraContext

class FeatureExtractor(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        self.add_loss("feature_mean_loss", x.mean())  # Step #2: Register the loss where it is created

        x = self.fc2(x)
        return x


...

with ExtraContext(model) as ctx:  # Step #3: Create the collection context in the training step
    logits = model(x)
    main_loss = F.cross_entropy(logits, targets)
    extra_losses = ctx.get_losses()  # Step #4: Collect losses registered inside the model

    ...

Use with Lightning

Lightning's LightningModule is a subclass of torch.nn.Module, so the same APIs can be used in Lightning training_step, validation_step, or inside module forward() methods. Wrap one model call with ExtraContext(self) in the step, then collect losses, metrics, hooks, shared objects, and logs from the context.

📚 Other Usage

Export Metrics

loss usually goes into total_loss and participates in backward(). A metric is only for monitoring, such as accuracy. Keeping them separate avoids mixing observation-only values into the optimized loss.

class ClassifierHead(nn.Module):
    def __init__(self):
        ...

    def forward(self, x, targets):
        logits = self.classifier(x)

        with torch.no_grad():
            preds = logits.argmax(dim=-1)
            acc = (preds == targets).float().mean()
        get_context(self).add_metric("train/accuracy", acc)  # For monitoring only

        return logits


head = ClassifierHead()
with ExtraContext(head) as ctx:
    logits = head(x, targets)
    main_loss = F.cross_entropy(logits, targets)
    metrics = ctx.get_metrics()

Hooks

A hook can store a callback that should run after forward(). ExtraContext only collects these functions; the training step decides when to execute them.

class FeatureBlock(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        y = self.block(x)
        get_context(self).add_hook("feature_norm", lambda: y.detach().norm().item())
        return y


block = FeatureBlock()
with ExtraContext(block) as ctx:
    output = block(x)
    hook_results = {
        name: [hook() for hook in hooks]
        for name, hooks in ctx.hooks.items()
    }

Shared Objects

put/get/pop can pass values out to the training step, but it can also share objects between modules in the same forward pass. This bypasses forward() arguments and return values. Note that this method is not very safe: it relies on key names and module call order, so it is best kept to small experimental uses.

class Encoder(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        latent = self.backbone(x)
        get_context(self).put("latent", latent)
        return latent


class Decoder(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        ctx = get_context(self)
        latent = ctx.pop("latent")
        return self.head(x + latent)


class AutoEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

    def forward(self, x):
        encoded = self.encoder(x)
        return self.decoder(encoded)


model = AutoEncoder()
with ExtraContext(model):
    output = model(x)

Shared Logger

Pass a logger when creating ExtraContext if several modules need to write logs. Modules still call get_context(self).log(...); Note: Keras style does not monkey patch log onto torch.nn.Module.

records = []


def logger(event, **fields):
    records.append((event, fields))


class FeatureBlock(nn.Module):
    def __init__(self):
        ...

    def forward(self, x):
        y = self.block(x)
        get_context(self).log("feature_block", mean=y.detach().mean().item())
        return y


block = FeatureBlock()
with ExtraContext(block, logger=logger):
    output = block(x)

🔧 API

Package Map

torchextractx.context.ExtraContext

torchextractx.null_context.NullContext

torchextractx.keras_style

Note: Use get_context(self).log(...) for logging; Keras style does not monkey patch log onto torch.nn.Module.

⚠️ Notes

Thread Safety

Not thread-safe. Don't use from multiple threads simultaneously.

Nested Contexts

Nested ExtraContext on the same model is not allowed and will raise an error:

with ExtraContext(model):
    with ExtraContext(model):  # Raises ValueError
        pass

Memory Management

All data is cleared after exiting the with block. Don't access outside:

with ExtraContext(model) as ctx:
    output = model(x)
    losses = ctx.get_losses()  # OK

losses = ctx.get_losses()  # Error

🧪 Running Tests

pytest tests/ -v

Or with unittest:

python -m unittest discover tests/ -v

📋 Requirements

• Python ≥ 3.10
• PyTorch ≥ 2.0.0

📝 License

MIT - see LICENSE

🚢 Release Process

• The canonical version lives in pyproject.toml under [project].version.
• Release tags must use vX.Y.Z and match that version exactly.
• GitHub Release notes are sourced from CHANGELOG.md.
• Pushing a matching tag runs tests, builds distributions, creates a GitHub Release, and publishes to PyPI through PyPI Trusted Publishing.
• PyPI publishing does not require a PyPI token; authentication is handled by GitHub Actions OpenID Connect.

🤝 Contributing

Welcome to submit issues and PRs.