Based on TensorRT sample code(#https://tvm.apache.org/docs/deploy/tensorrt.htm), I write a sample with Resnet18 model, the code are list as follows:
‘’ Test resnet18 TVM+TensorRT directly ‘’’ import tvm from tvm import relay
import numpy as np
from tvm.contrib.download import download_testdata
#Pytorch import
import torch import torchvision
from PIL import Image
model_name = “resnet18” model = getattr(torchvision.models, model_name)(pretrained=True) model = model.eval()
#we grab the TorchScripted model via tracing input_shape = [1,3,224,224] input_data = torch.randn(input_shape) scripted_model = torch.jit.trace(model,input_data).eval()
img_name = “cat.png” img = Image.open(img_name).resize((224,224))
from torchvision import transforms
my_preprocess = transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224,0.225]), ] )
img = my_preprocess(img) #print(img.shape) img = np.expand_dims(img,0) #print(img.shape)
#import the graph to Relay input_name = “input0” shape_list = [(input_name, img.shape)] mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
#print(params)
from tvm.relay.op.contrib.tensorrt import partition_for_tensorrt mod, config = partition_for_tensorrt(mod, params) #print(mod) #print(config)
target = “cuda” with tvm.transform.PassContext(opt_level=3, config={‘relay.ext.tensorrt.options’: config}): lib = relay.build(mod, target=target, params=params)
#export the module lib.export_library(‘resnet18_trt.so’)
#ctx = tvm.gpu(0) #load module #dev = tvm.cuda(0) dev = tvm.gpu(0) loaded_lib = tvm.runtime.load_module(‘resnet18_trt.so’)
gen_module = tvm.contrib.graph_executor.GraphModule(loaded_lib’default’)
dtype = “float32”
input_data = np.random.uniform(0,1, input_shape).astype(dtype)
output_shape = (1,1000)
gen_module.run(data=img)
out = gen_module.get_output(0, tvm.nd.empty(output_shape)) top1_tvm = np.argmax(out.asnumpy()[0])
synset_url = “”.join( [ “https://raw.githubusercontent.com/Cadene/”, “pretrained-models.pytorch/master/data/”, “imagenet_synsets.txt”, ] )
synset_name = “imagenet_synsets.txt” synset_path = download_testdata(synset_url, synset_name, module=“data”)
with open(synset_path) as f: synsets = f.readlines()
synsets = [x.strip() for x in synsets]
splits = [line.split(" ") for line in synsets]
key_to_classname = {spl[0]:" ".join(spl[1:]) for spl in splits}
#print(key_to_classname)
class_url = “”.join( [ “https://raw.githubusercontent.com/Cadene/”, “pretrained-models.pytorch/master/data/”, “imagenet_classes.txt”, ] )
class_name = “imagenet_classes.txt”
class_path = download_testdata(class_url, class_name, module=“data”)
with open(class_path) as f: class_id_to_key = f.readlines()
class_id_to_key = [x.strip() for x in class_id_to_key]
tvm_class_key = class_id_to_key[top1_tvm]
print(“Relay top-1 id: {}, class name: {}”.format(top1_tvm, key_to_classname[tvm_class_key]))
However, I input a cat image, the output is Relay top-1 id: 623, class name: letter opener, paper knife, paperknife. I double check use pytorch code, the output is defintely CAT!