dynamic_shape_example.py

import numpy as np

import tvm
from tvm import relay
from tvm.runtime.vm import VirtualMachine

target = "cuda"
data_shape = (relay.Any(), 3, 224, 224)
weight_shape = (32, 3, 3, 3)

data = relay.var("data", shape=data_shape, dtype="float32")
weight = relay.var("weight", shape=weight_shape, dtype="float32")
bias = relay.var("bias", shape=(weight_shape[0],), dtype="float32")
conv = relay.nn.conv2d(data, weight, kernel_size=(3, 3))
out = relay.nn.bias_add(conv, bias)
func = relay.Function([data, weight, bias], out)
mod = tvm.IRModule.from_expr(func)

with tvm.transform.PassContext(opt_level=3):
    vm_exec = relay.vm.compile(mod, target)

vm = VirtualMachine(vm_exec, ctx=tvm.context(target))

data_np_shape = tuple([1] + list(data_shape[1:]))
data_np = np.random.uniform(0, 255, size=data_np_shape).astype("float32")
weight_np = np.random.uniform(0, 1, size=weight_shape).astype("float32")
bias_np = np.random.uniform(0, 1, size=(weight_shape[0],)).astype("float32")

vm.set_input("main", **{"data": data_np, "weight": weight_np, "bias": bias_np})
vm.run()