TorchExtraContext: throw losses, metrics, logs, or whatever from anywhere!

1

Hello everyone!

I would like to introduce my self-made little tool:

TorchExtraContext is a small utility for collecting losses, metrics, logs, intermediate outputs, or other extra per-forward-pass information from anywhere inside a nested PyTorch model.

My original need was to throw auxiliary losses from the gate modules of a Mixture-of-Experts model, but I did not want to pass those losses layer by layer through every forward() return value. So I made this tool!

The basic idea is that you register a context object on the root nn.Module, usually inside something like the Lightning training_step. This context binds the root module and all child modules triggered from that root. When the with context exits, it is cleaned up.

Then any module inside the model can use the context to throw out a loss, metric, log, or whatever else you want to collect.

You can install it with:

pip install torchextractx

For example, you can use the Keras-style API like this:

import torch
import torch.nn as nn
import torch.nn.functional as F

import torchextractx.keras_style  # enable torch.nn.Module.add_loss
from torchextractx import ExtraContext


class GateModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(128, 16)

    def forward(self, x):
        gate_logits = self.proj(x)

        aux_loss = gate_logits.float().mean()
        self.add_loss("gate_aux_loss", aux_loss)

        return gate_logits


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.gate = GateModule()
        self.head = nn.Linear(16, 10)

    def forward(self, x):
        x = self.gate(x)
        return self.head(x)


model = Model()

with ExtraContext(model) as ctx:
    logits = model(x)

    main_loss = F.cross_entropy(logits, targets)
    extra_losses = ctx.get_losses()

    total_loss = main_loss + sum(extra_losses.values())

with ExtraContext(model) as ctx:
    output = model(x)
    losses = ctx.get_losses()
    metrics = ctx.get_metrics()
    logs = ctx.get_logs()

This tool has been very useful for my own research work. I published it about a year ago, but almost nobody noticed it. I would love to share it with more people, so now I am taking the initiative!

I recently refactored my ugly hand-written code with the help of Codex. The AI-assisted core code has been reviewed by my own human eyes. I also added a proper test and release workflow. It is only a small tool, but I really want more people to know about it.

One important note: it probably does not support torch.compile at the moment. If anyone needs that, issues and PRs are very welcome.

Thank you everyone!

I want attention, I want users, and I want to be loved! :wink:

1 Like