Deploy mixed precision model in libtorch


I tried the torch.cuda.amp.autocast in PyTorch and it works well for my model. Now I want to deploy my trained model in C++ with the nightly built libtorch (version 1.7.0). However, I cannot find a corresponding function for autocast in the libtorch library API.

May I ask what is the proper way to deploy a mixed precision model in libtorch?



1 Like

I’m not sure, if libtorch supports amp fully, but @mcarilli would know. :slight_smile:


libtorch doesn’t officially support autocast in C++ yet. You can try to rig it by imitating the C++ calls made by torch.cuda.amp.autocast during __enter__ at the beginning of your C++ autocasted region and imitating the C++ calls made by __exit__ at the end of your C++ autocasted region. The mapping between torch._C functions and C++ -side function calls can be seen here.

Right now doing the above is definitely “off menu.” I plan to support autocast in libtorch by adding an RAII guard that does all those things in an exception-safe way.