capture module inputs in PyTorch

capture.py

import torch
from contextlib import contextmanager


_capture_inputs_dict = {}
_capture_inputs_active = False


@contextmanager 
def capture_inputs():
    global _capture_inputs_dict
    global _capture_inputs_active
    assert not _capture_inputs_active, "only one at a time allowed"
    prev_dict = _capture_inputs_dict
    _capture_inputs_active = True
    _capture_inputs_dict = {}
    try:
        yield _capture_inputs_dict
    finally:
        _capture_inputs_active = False
        _capture_inputs_dict = prev_dict


class CaptureModule(torch.nn.Module):
    def __init__(self, module, name):
        super().__init__()
        self.module = module
        self.name = name
    def forward(self, *args, **kwargs):
        if _capture_inputs_active:
            _capture_inputs_dict[self.name] = [a.detach() for a in args]
        return self.module(*args, **kwargs)