Source code for catalyst.utils.tracing

from typing import Tuple, Union

import torch
from torch import jit

from catalyst.extras.forward_wrapper import ModelForwardWrapper
from catalyst.typing import TorchModel
from catalyst.utils.distributed import get_nn_from_ddp_module


def trace_model(
    model: TorchModel,
    batch: Union[Tuple[torch.Tensor], torch.Tensor],
    method_name: str = "forward",
) -> jit.ScriptModule:
    """Traces model using runner and batch.

    Args:
        model: Model to trace
        batch: Batch to trace the model
        method_name: Model's method name that will be
            used as entrypoint during tracing

    Example:
        .. code-block:: python

           import torch

           from catalyst.utils import trace_model

           class LinModel(torch.nn.Module):
               def __init__(self):
                   super().__init__()
                   self.lin1 = torch.nn.Linear(10, 10)
                   self.lin2 = torch.nn.Linear(2, 10)

               def forward(self, inp_1, inp_2):
                   return self.lin1(inp_1), self.lin2(inp_2)

               def first_only(self, inp_1):
                   return self.lin1(inp_1)

           lin_model = LinModel()
           traced_model = trace_model(
               lin_model, batch=torch.randn(1, 10), method_name="first_only"
           )

    Returns:
        jit.ScriptModule: Traced model
    """
    nn_model = get_nn_from_ddp_module(model)
    wrapped_model = ModelForwardWrapper(model=nn_model, method_name=method_name)
    traced = jit.trace(wrapped_model, example_inputs=batch)
    return traced


__all__ = ["trace_model"]