Compile gluoncv ssd mobilenet model for ARM Mali GPU

compile-gluoncv-ssd.py

import tvm
from tvm import relay
from tvm.relay.expr_functor import ExprMutator
from tvm.contrib import graph_runtime
from gluoncv import model_zoo
from tvm import autotvm
 
class ScheduleConv2d(ExprMutator):
    def __init__(self, device):
        self.device = device
        super().__init__()
 
    def visit_call(self, expr):
        visit = super().visit_call(expr)
        if expr.op == tvm.relay.op.get("vision.get_valid_counts"):
           return relay.annotation.on_device(visit, self.device)
        return visit
 
def schedule_conv2d_on_gpu(expr):
      sched = ScheduleConv2d(tvm.cpu(0))
      return sched.visit(expr)
 
 
model_name = "ssd_512_mobilenet1.0_voc"
 
image_size = 512
dshape = (1, 3, image_size, image_size)
 
 
######################################################################
# Convert and compile model for ARM Mali GPU
 
block = model_zoo.get_model(model_name, pretrained=True)
 
def compile(target, target_host):
    net, params = relay.frontend.from_mxnet(block, {"data": dshape})
    net = schedule_conv2d_on_gpu(net)
    net = relay.ir_pass.infer_type(net)
    with relay.build_config(opt_level=3, fallback_device=tvm.opencl(0)):
        # apply autotvm optimal logs if you have them
        #with autotvm.apply_graph_best("ssd_512_resnet50_v1_voc_opt.log"):
        #with autotvm.apply_graph_best("yolo3_darknet53_voc_opt.log"):
            graph, lib, params = relay.build(net, target, params=params, target_host=target_host)
    return graph, lib, params
 
 
target = {"cpu": tvm.target.arm_cpu('rk3399'), "opencl": tvm.target.mali('rk3399')}
target_host = 'llvm -target=aarch64-linux-gnu'
graph, lib, params = compile(target, target_host)

########
# To run compiled model on edge device we need to specify both cpu and opencl in the context
# ctx = [tvm.cpu(0), tvm.opencl(0)]
# m = graph_runtime.create(graph, lib, ctx)