With torch dynamo, we can dispatch a pytorch model to other awesome deep learning framework/compilers for acceleration. Hidet is one of such deep learning compilers that accelerates your model with a bunch of optimizations (e.g., subgraph fusion, rewriting and kernel tuning). To use hidet, please first install it via
$ pip install hidet
Then you can enable it via torch.compile(model, backend='hidet')
as shown in the code snippet below:
import torch
import hidet
# Define pytorch model
model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).cuda().eval()
x = torch.rand(1, 3, 224, 224).cuda()
# Compile the model through Hidet
hidet.torch.dynamo_config.search_space(2) # tune the kernel performance
model_opt = torch.compile(model, backend='hidet')
# Run the optimized model
y = model_opt(x)
Here are some benchmarks (Batch Size = 1, NVIDIA RTX 3090, Bert sequence length=128, with float32 data type)
Learn more about hidet and its optimization options in the tutorial and GitHub repository. Hidet originates our research work that tries to simplify writing tensor program with our proposed task-mapping programming paradigm. Please checkout our paper for more details.
In fact, I do not hope to expose too much these technical details to the audience as most of them do not need to know this to use hidet. The blog section does not surves as this purpose as well.
Another way to use hidet in torch is to use it to add custom kernel instead of writing pytorch extension. We can educate people how to write hidet script when we have more support on this part (i.e., use hidet script to implement torch.autograd.Function) and write a pytorch blog. We can do this in the future, but I do not suggest we do it now. The audiance would be ML people who want to write pytorch extension to implement their kernels. For those people, they will have interest to learn hidet script.