jwfromm opened a new pull request, #11199:
URL: https://github.com/apache/tvm/pull/11199
This PR adds a new pass to `relay.transform` that creates a dispatcher
around an input module to handle multiple input shapes. For example consider a
case where I'd like to optimize my model to handle both `batch_size=1` and
`batch_size=4`. I can now do so elegantly as follows:
```
shape_dict = {'data': [1, 3, 224, 224]}
model_bs1 = tvmc.load('my_model.onnx', shape_dict=shape_dict)
tvmc.tune(model, log_file='batch_1.logs')
shape_dict = {'data': [4, 3, 224, 224]}
model_bs4 = tvmc.load('my_model.onnx', shape_dict=shape_dict)
tvmc.tune(model, log_file='batch_4.logs')
# Create dispatcher for multiple batch sizes.
flex_mod = relay.transform.FlexibleShapeDispatch(buckets=[1,
4])(model_bs1.mod)
with ApplyHistoryBest(['batch_1.logs', 'batch_4.logs']):
exe = relay.vm.compile(flex_mod, "llvm")
# Now we can run inputs with either batch 1 or batch 4 and get the tuned
performance!
batch_1 = np.random.rand(1, 3, 224, 224).astype("float32")
vm.benchmark(tvm.cpu(), batch_2, func_name="main")
batch_4 = np.random.rand(4, 3, 224, 224).astype("float32")
vm.benchmark(tvm.cpu(), batch_4, func_name="main")
```
As seen above `FlexibleShapeDispatch` is a simple halfway point between
fully static and fully dynamic graphs that allows us to leverage TVM tuning. If
an input shape is not provided in `buckets`, it will either run fully
dynamically using `relay.Any`, or if the `auto_pad` argument is set for
`FlexibleShapeDispatch`, padding will be applied to match the closest bucket.
To make applying tuning logs more convenient, I also added the ability to
load and merge multiple files to both autotvm and autoscheduler.
Thanks @jroesch for providing the backbone of this implementation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
