4.compile_save_lib.py

#!/bin/env python

import tvm
from tvm import relay
from module import *
import numpy as np
import torch

import shutil
import torchvision
from torchvision import transforms as tfs
from torch.utils.tensorboard import SummaryWriter

from tvm.contrib import graph_executor
from tvm.contrib import utils

dev = tvm.cpu(0)

my_module = torch.load("my_module_0.pth")
my_module.eval()


random_input = torch.randn([1, 1, 29, 29])
scripted_module = torch.jit.trace(my_module, random_input).eval()


input_name = "input0"

img = torchvision.io.read_image("img_fortest/2.jpg", mode = torchvision.io.image.ImageReadMode.GRAY)
img = tfs.functional.resize(img, (29, 29))
img = img.reshape((-1, 1, 29, 29))
img = tfs.functional.autocontrast(img)
img = tfs.functional.convert_image_dtype(img, dtype = torch.float32)

shape_list = [(input_name, img.shape)]
mod, params = relay.frontend.from_pytorch(scripted_module, shape_list)

print(mod.astext(show_meta_data = False))

#mod.summary()
#exit()

target = tvm.target.Target("llvm", host = "llvm")
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level = 3):
	lib = relay.build(mod, target = target, params = params)


m = graph_executor.GraphModule(lib["default"](dev))
# Set inputs
m.set_input(input_name, img)
# Execute
m.run()
# Get outputs
tvm_output = m.get_output(0)
print(tvm_output)
print(np.argmax(tvm_output.numpy()[0]))

temp = utils.tempdir()
path_lib = temp.relpath("my_module.tar")
#print(path_lib)
lib.export_library("my_module.tar")
#print(temp.listdir())