Compiling a PyTorch DNN with HPVM Tensor Extension

We will now compile a DNN model defined in PyTorch, VGG16-CIFAR10, with HPVM Tensor Extension. This will result in a binary that directly runs on your system. The Python package torch2hpvm that is installed along with HPVM serves this purpose.

  • This package lives under hpvm/projects/torch2hpvm and should have been installed by the installer.

  • The Keras frontend similarly compiles Keras model through HPVM down to binary, and its usage can be found in the documentation.

Note that we’ll be working under directory hpvm/benchmarks/dnn_benchmarks. When you work in a different directory, the procedure below will apply similarly as long as HPVM is built and installed. “””

First, download our pre-trained model and slice of dataset:

curl -L https://databank.illinois.edu/datafiles/o3izd/download -o ./model_params.tar.gz
tar xzf model_params.tar.gz

This should create a directory model_params with the model parameters.

Now, we define 2 datasets for autotuning and testing for VGG16. These datasets are provided as ./model_params/vgg16_cifar10/{tune|test}_{input|labels}.bin, where tune and test prefixes signify tuning and testing set.

from torch2hpvm import BinDataset
from pathlib import Path

data_dir = Path("./model_params/vgg16_cifar10")
dataset_shape = 5000, 3, 32, 32  # NCHW format.
tuneset = BinDataset(data_dir / "tune_input.bin", data_dir / "tune_labels.bin", dataset_shape)
testset = BinDataset(data_dir / "test_input.bin", data_dir / "test_labels.bin", dataset_shape)

BinDataset is a utility torch2hpvm provides for creating dataset over binary files. Any instance of torch.utils.data.Dataset can be used here.

Note that each module is bound to 2 datasets: a “tune” and a “test” set. The generated binary accepts an argument to be either the string “tune” or “test”, and performs inference over a dataset accordingly.

Create a VGG16 and load its pretrained checkpoint:

import torch
from torch.nn import Module
from pytorch import dnn  # Defined at `hpvm/benchmarks/dnn_benchmarks/pytorch/dnn`

model: Module = dnn.VGG16Cifar10()
checkpoint = "model_params/pytorch/vgg16_cifar10.pth.tar"
model.load_state_dict(torch.load(checkpoint))

Any torch.nn.Module can be similarly used, as long as they only contain the tensor operators supported in HPVM. See “Supported Operators” in PyTorch frontend and Keras frontend.

Now we are ready to export the model. The main functioning class of torch2hpvm is ModelExporter:

from torch2hpvm import ModelExporter

output_dir = Path("./vgg16_cifar10")
build_dir = output_dir / "build"
target_binary = build_dir / "vgg16_cifar10"
batch_size = 500
conf_file = "../../test/hpvm_c_dnn/vgg16_cifar10/tuner_confs.txt"
exporter = ModelExporter(model, tuneset, testset, output_dir, config_file=conf_file)
exporter.generate(batch_size=batch_size).compile(target_binary, build_dir)

output_dir, build_dir, and target_binary define the folder for code generation, compilation, and path to the compiled binary respectively. batch_size is the batch size the binary uses during inference.

  • Note that conf_file is the path to an HPVM approximation configuration file. This file decides what approximation the binary will use during inference. This path is hardcoded into the binary and is only read when the binary starts, so it’s fine to have conf_file point to a non-existing path. An example can be found at hpvm/test/hpvm_c_dnn/vgg16_cifar10/tuner_confs.txt.

  • exporter.generate generates the HPVM-C code while exporter.compile is a helper that invokes the HPVM compiler for you.

Now there should be a binary at ./vgg16_cifar10/build/vgg16_cifar10. Running it without argument will perform an inference over the test set. (The accuracy of inference is written to the file ./final_accuracy.)