Hello everyone!
I have a task of running a Pytorch model in the iOS app and I would like to give TVM a shot. For the time being I’m at the stage of model compilation.
I have successfully compiled it for MacOS using
TVMC
(Compiling and Optimizing a Model with TVMC — tvm 0.8.dev0 documentation). But I had a trouble coming up with the target that will compile the model for iOS. I looked through the target list and didn’t find one that suits my needs.
So I decided to go a different way and I have found a python script for Pytorch model converting. I gave it slight modifications and now it looks like this:
model_folder_path = "path to my tvm model"
model_name = "model name"
model = FAN(2)
model_path = "path to pytorch model"
checkpoint = torch.load(model_path,map_location = 'cpu')['state_dict']
model = torch.nn.DataParallel(model)
model.load_state_dict(checkpoint)
model = model.eval()
# We grab the TorchScripted model via tracing
input_shape = [1, 3, 256, 256]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()
# Preprocess the image and convert to tensor
img = prepare_torch_input("path to image") #creating an input to the model
######################################################################
# Import the graph to Relay
# -------------------------
# Convert PyTorch graph to Relay graph. The input name can be arbitrary.
input_name = "input"
shape_list = [(input_name, img.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
######################################################################
# Relay Build
# -----------
# Compile the graph to llvm target with given input specification.
arch = "arm64"
target = tvm.target.Target("llvm -mtriple=%s-apple-darwin" % arch, host="llvm")
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, params=params)
if os.path.exists(model_folder_path) and os.path.isdir(model_folder_path):
shutil.rmtree(model_folder_path)
os.mkdir(model_folder_path)
libpath = model_folder_path + "/" + model_name + ".so"
lib.export_library(libpath, cc="/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/bin/clang")
graph = lib.graph_json
graph_json_path = model_folder_path + "/" + model_name + ".json"
with open(graph_json_path, 'w') as fo:
fo.write(graph)
param_path = model_folder_path + "/" + model_name + ".params"
with open(param_path, 'wb') as fo:
fo.write(relay.save_param_dict(params))
This script successfully builds the model, and if I’m putting it into my iOS Xcode project, I have an error:
Building for iOS, but the linked item ‘model.so’ was built for macOS.
I understand that I’m doing something wrong with the target and I went through all TVM docs, but I didn’t find any articles about how to form correct targets for models compilation. I’d be happy if someone can nudge me into the right direction on how to solve it.