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 haveconf_file
point to a non-existing path. An example can be found athpvm/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
.)