Understanding Trace vs. Serialize when using a TorchScript Model

C++
1

I have two functions as part of a torch.Module object

def trace_model(model, sample_input, save_fp=None):
    model.eval()
    traced_model = torch.jit.trace(model, sample_input)
    if save_fp is not None:
        traced_model.save(save_fp)
    return traced_model


def serialize_model(model, sample_input, save_fp=None):
    model.eval()
    serialized_model = torch.jit.script(model, sample_input)
    if save_fp is not None:
        serialized_model.save(save_fp)
    return serialized_model

One traces the model and saves the result, while the other serializes it. Functionally what is the difference and benefits/draw backs of tracing versus serializing in the context of using the resulting .pt file in a C++ application ?

2

torch.jit.trace will record the used operations via executing the forward pass with the provided input. Conditional control flow and other data-dependent properties will be captured during this execution and saved into the graph.
torch.jit.script is able to understand loops, conditions, etc. and allows for a more flexible graph creation.

1 Like