TorchScript enables PyTorch models to be converted into a format that can run independently of Python. This allows models to be deployed in production environments, including servers without Python runtime. Two primary methods exist for converting PyTorch models to TorchScript: tracing and scripting.

Trracing

Tracing involves executing a model instance while recording all operation. This method is suitable for models with straightforward, linear data flow that lack control structures like loops or condisional statements.

import torch

class IncrementModel(torch.nn.Module):
    def forward(self, tensor_input):
        return tensor_input + 2

model_instance = IncrementModel()
sample_data = torch.rand(2, 4)

traced_model = torch.jit.trace(model_instance, sample_data)
traced_model.save("increment_model.pt")

In this example, torch.jit.trace captures the operations of the IncrementModel when provided with sample input data, then saves the traced model to disk.

Scripting

Scripting converts models by analyzing Python source code directly, without requiring model execution. This approach is particularly useful for models containing control flow elements such as conditional statements or loops.

import torch

class ConditionalModel(torch.nn.Module):
    def forward(self, data_tensor):
        if torch.mean(data_tensor) > 0.7:
            return data_tensor
        else:
            return data_tensor * 3

model_instance = ConditionalModel()
scripted_model = torch.jit.script(model_instance)
scripted_model.save("conditional_model.pt")

Here, torch.jit.script transforms the ConditionalModel into TorchScript format by examining its code structure, then saves the resulting model.

Using TorchScript Models

TorchScript models saved in this format can be utilized without Python environments. They can be loaded and executed through C++ APIs, facilitating deployement in production settings that may not support Python directly.

#include <torch/script.h>

int main() {
    torch::jit::script::Module loaded_module;
    loaded_module = torch::jit::load("increment_model.pt");

    std::vector<torch::jit::IValue> model_inputs;
    model_inputs.push_back(torch::ones({2, 4}));

    at::Tensor result = loaded_module.forward(model_inputs).toTensor();
    std::cout << result << std::endl;
}

This C++ code demonstrates loading and executing a TorchScript model, enabling PyTorch models to integrate into broader production ecosystems, encluding environments without Python support.