In this tutorial, we demonstrate how to use the Lexer library with a PyTorch Super-resolution model used in PyTorch's tutorial's for running models using ONNXRuntime.

Introduction to ONNX and ONNXRuntime

Here we'll provide a quick summary of the Open Neural Network Exchange (ONNX).

ONNX, as described by the organization, is:

an open format built to represent machine learning models. ONNX defines a common set of operators - the building blocks of machine learning and deep learning models - and a common file format to enable AI developers to use models with a variety of frameworks, tools, runtimes, and compilers.

The sibling project ONNX Runtime "automatically parses through your model to identify optimization opportunities and provides access to the best hardware acceleration available."

If you'd like to learn more, we recommend:

Comparing Your Model's Results Between PyTorch and ONNXRuntime

You'll want to start your script with imports for:

PyTorch
ONNX to access the ONNX Runtime library
The Lexer Context and Decorator

# Third party imports
import torch.onnx
import torch.utils.model_zoo as model_zoo
from torch import nn
from torch.nn import init

# Lexer imports
from lexer.api.lexer_context import LexerContext
from lexer.decorators.lexer import Lexer

Apply the Lexer Decorator to a Model Class

If you've read our guide on the Lexer Decorator, you'll be ready to apply @Lexer decorator to a class. The example class below is a model example from PyTorch).

@Lexer(
    input_names=["input"],
    output_names=["output"],
)
class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor**2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

        self._initialize_weights()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x

    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv2.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv3.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv4.weight)

Initiate Your Model and Benchmark With Lexer

With the class defined and the Lexer decorator applied, you can now prep your model for execution, call the Lexer benchmark() API, and print out the benchmarking results.

if __name__ == "__main__":

    # Create the super-resolution model by using the above model definition.
    torch_model = SuperResolutionNet(upscale_factor=3)

    # Load pretrained model weights
    model_url = "https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth"
    batch_size = 1  # just a random number

    # Initialize model with the pretrained weights
    map_location = lambda storage, loc: storage
    if torch.cuda.is_available():
        map_location = None
    torch_model.load_state_dict(
        model_zoo.load_url(model_url, map_location=map_location)
    )

    # set the model to inference mode
    torch_model.eval()

    # Input to the model
    input_tensor = torch.randn(batch_size, 1, 224, 224, requires_grad=True)

    # ------------------ Model Prep Ends, Lexer Benchmarking Begins ----------------------------------------------

    lexer_context = LexerContext()

    lexer_context.benchmark(
        torch_model=torch_model,
        input=input_tensor,
        num_iterations=5,
        batch_size=batch_size,
    )
    lexer_context.flush(name="benchmark", target_directory=".")
  
    print(torch_model.lexer.schema_json(indent=2))

The Full SuperResolutionNet with Lexer Script

Piecing it all together gives you the below one file script that you can save to a .py file and execute.

# Third party imports
import torch.onnx
import torch.utils.model_zoo as model_zoo
from torch import nn
from torch.nn import init

# Lexer imports
from lexer.api.lexer_context import LexerContext
from lexer.decorators.lexer import Lexer


@Lexer(
    input_names=["input"],
    output_names=["output"],
)
class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor**2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

        self._initialize_weights()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x

    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv2.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv3.weight, init.calculate_gain("relu"))
        init.orthogonal_(self.conv4.weight)


if __name__ == "__main__":

    # Create the super-resolution model by using the above model definition.
    torch_model = SuperResolutionNet(upscale_factor=3)

    # Load pretrained model weights
    model_url = "https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth"
    batch_size = 1  # just a random number

    # Initialize model with the pretrained weights
    map_location = lambda storage, loc: storage
    if torch.cuda.is_available():
        map_location = None
    torch_model.load_state_dict(
        model_zoo.load_url(model_url, map_location=map_location)
    )

    # set the model to inference mode
    torch_model.eval()

    # Input to the model
    input_tensor = torch.randn(batch_size, 1, 224, 224, requires_grad=True)

    # ------------------ Model Prep Ends, Lexer Benchmarking Begins ----------------------------------------------

    lexer_context = LexerContext()

    lexer_context.benchmark(
        torch_model=torch_model,
        input=input_tensor,
        num_iterations=5,
        batch_size=batch_size,
    )
    lexer_context.flush(name="benchmark", target_directory=".")

    print(torch_model.lexer.schema_json(indent=2))

See Benchmark Results

By executing the script from your terminal, you should now see results like the below, which includes JSON of your benchmark metrics.

You'll notice that you'll get two objects returned in the results array.

The default Lexer Configuration enables two engines to be benchmarked: Torch and ONNXRuntime.

The key insight for you as a user will be able to compare how the same model performs when executed through two different engines, and how those engines impact performance!

Results generated and exported to SuperResolutionNet.csv!
Generated results:
-------------------
[{'num_samples': 5, 'variance': 6.403878733313712e-06, 'mean_ms': 27.737423396320082, 'p50_ms': 26.324664999265224, 'p90_ms': 30.89536899351515, 'p95_ms': 31.47643099364359, 'p99_ms': 31.941280593746338, 'throughput_per_sec': 36.0523753670887, 'engine': 'onnxruntime', 'version': '1.13.1', 'batch_size': 1, 'optimizer': 'N/A', 'providers': 'N/A', 'device': 'cpu', 'timestamp': '2023-01-24 01:48:15.215648'}, 
 {'num_samples': 5, 'variance': 5.308219160244051e-06, 'mean_ms': 36.46238420042209, 'p50_ms': 35.657344997162, 'p90_ms': 39.08862060052343, 'p95_ms': 39.98508080258034, 'p99_ms': 40.70224896422587, 'throughput_per_sec': 27.425524192365454, 'engine': 'torch', 'version': '1.13.1', 'batch_size': 1, 'optimizer': 'N/A', 'providers': 'N/A', 'device': 'cpu', 'timestamp': '2023-01-24 01:48:15.215648'}]

By including the flush function, you will also get same data in CSV format as part of the Lexer artifact folder.