Deploy mixed precision model in libtorch

volpato30: <p>Hi,</p>
<p>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 <a href="https://pytorch.org/cppdocs/api/library_root.html" rel="nofollow noopener">libtorch library API</a>.</p>
<p>May I ask what is the proper way to deploy a mixed precision model in libtorch?</p>
<p>Thanks,</p>
<p>Rui</p>

ptrblck: <p>I’m not sure, if <code>libtorch</code> supports amp fully, but <a class="mention" href="/u/mcarilli">@mcarilli</a> would know. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

mcarilli: <p>libtorch doesn’t officially support autocast in C++ yet.  You can try to rig it by imitating <a href="https://github.com/pytorch/pytorch/blob/e7ed0b3fae816b5a39d94c4f67e345d7a784a7ef/torch/cuda/amp/autocast_mode.py#L119-L128" rel="nofollow noopener">the C++ calls made by <code>torch.cuda.amp.autocast</code></a> during <code>__enter__</code> at the beginning of your C++ autocasted region and imitating the C++ calls made by <code>__exit__</code> at the end of your C++ autocasted region.  The mapping between <code>torch._C</code> functions and C++ -side function calls can be seen <a href="https://github.com/pytorch/pytorch/blob/e7ed0b3fae816b5a39d94c4f67e345d7a784a7ef/torch/csrc/autograd/init.cpp#L84-L121" rel="nofollow noopener">here</a>.</p>
<p>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.</p>

Robin_Lobel: <p>Hi <a class="mention" href="/u/mcarilli">@mcarilli</a> could you provide a very simple example on how to do it “as is” with libtorch ?</p>
<p>Here’s my training loop right now:</p>
<pre><code class="lang-auto">    torch::optim::Adam optim(model-&gt;parameters(), torch::optim::AdamOptions(1e-3));
    for(int b=0; b&lt;num_batch; b++)
    {
        optim.zero_grad();
        torch::Tensor result = model(source[b]);
        torch::Tensor loss = torch::l1_loss(result, target);
        loss.backward();
        optim.step();
    }
</code></pre>
<p>I’m not sure where to insert the (unofficial) amp calls and how… thanks !</p>
<p>BTW when can we expect a proper libtorch implementation (RAII guard as you mentionned) ?</p>

mcarilli: <p>For non-nested autocast usage it’s not too complex:</p>
<pre><code class="lang-auto">    torch::optim::Adam optim(model-&gt;parameters(), torch::optim::AdamOptions(1e-3));
    for(int b=0; b&lt;num_batch; b++)
    {
        optim.zero_grad();

        at::autocast::set_enabled(true);
        torch::Tensor result = model(source[b]);
        at::autocast::clear_cache();
        at::autocast::set_enabled(false);

        torch::Tensor loss = torch::l1_loss(result, target);
        loss.backward();
        optim.step();
    }
</code></pre>
<p>For nested usage, ie nesting autocast-disabled regions in autocast-enabled regions, manually imitating the behavior of the Python context manager for nested usage would be harder, but the planned RAII guard can make this easy on the user side.</p>
<p>I’m not sure when I’ll PR the RAII guard, I have a backlog of other bugs to work through (mostly adding coverage for ops that currently dont work with autocast).</p>
<p>However, there’s a bigger problem here:  your code is a training loop.  When training with autocast, you should also use gradient scaling.  If your network happens to converge without gradient scaling, great, but I expect you will need it.</p>
<p>Imitating <code>torch.cuda.amp.GradScaler</code> on the C++ side is not as easy as imitating <code>torch.cuda.amp.autocast</code>. <code>autocast</code>'s Python source is (mostly) a thin wrapper around the calls above, but <code>GradScaler</code>'s Python source contains a fair amount of important logic.  I do plan to create an equivalent C++ <code>GradScaler</code>, but that’s a substantial effort.  In the meantime, you can try to manually duplicate <a href="https://github.com/pytorch/pytorch/blob/e7ed0b3fae816b5a39d94c4f67e345d7a784a7ef/torch/cuda/amp/grad_scaler.py#L40" rel="nofollow noopener">GradScaler’s logic</a>.  Much of GradScaler’s complexity comes from handling the possibility of different layers on different devices.  If you can rule that out in your case you can write the logic more simply, but it’s still not trivial.</p>
<p>Using <code>autocast</code> as shown above for inference (not training) should be fine without GradScaler.</p>

Robin_Lobel: <p>Thanks for the autocast code.</p>
<p>Indeed, the scaler doesn’t seem trivial. I guess I’ll have to wait for a C++ implementation of it, as the tensor core speedup is specially interesting for training more than inference. Hopefully PyTorch 1.7 ?</p>

mcarilli: <p>Not 1.7.  The branch cut + code freeze for each release precedes the actual release by a month or more (as it must, to iron out bugs) and it’s coming up too soon.  For 1.7 my priority is autocast coverage for all high-demand ops that 1.6 missed (RNNs most notably).</p>
<p>1.8 is plausible though.  File an issue on Pytorch github and tag <a class="mention" href="/u/mcarilli">@mcarilli</a> <a class="mention" href="/u/gchanan">@gchanan</a> <a class="mention" href="/u/ngimel">@ngimel</a> so we know people want mixed precision training with the C++ API.</p>

Robin_Lobel: <p><a class="mention" href="/u/mcarilli">@mcarilli</a> with the upcoming RTX 3080/3090 introducing BF16 support (correct if I’m wrong), maybe I can skip gradient scaling and use your code as-is ?<br>
However, how do I tell AMP that I want to use BF16 and not FP16 ?</p>

mcarilli: <p>Hardware support for bfloat16 is there, but software support in pytorch’s cuda backend and cuda libraries is WIP.  Once it’s present, I’ll add a bfloat16 option to autocast.</p>
<p>Why is it essential that you train in C++?  Python training usually isn’t a problem and there’s lot of near-term work underway (and done) to reduce eager mode cpu overheads and make the difference between C++ and Python less important.</p>

Robin_Lobel: <p>Simply because I’m much more used to C++ than Python <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
And PyTorch is the only big framework to provide that possibility (TensorFlow doesn’t offer C++ model generation+training for instance).</p>
<p>Why full C++ support is relevant : python is the de-facto language for data scientists, but more and more desktop app developers are coming to deep learning, including for training models, and those desktop app developers are usually more familiar with C++ than python (you don’t code software in python).<br>
For instance <a href="https://new.steinberg.net/spectralayers/new-features/" rel="nofollow noopener">my desktop app</a> is entirely in C++ and I know exactly how to format data in C++, but I have zero clue in python, and I’m not a big fan of type-less languages too. Not judging python here, but it’s really not my cup of tea and I know I’m not the only one <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Another argument for that: I’m also generating and managing my training data in C++. So it’s a consistent pipeline from end to end, from data generation and management, to training, to implementation in the end-user software.</p>

davidw82: <p>Hi guys, and more specifically <a class="mention" href="/u/mcarilli">@mcarilli</a><br>
What is the current status of AMP using libtorch for inference?</p>
<p>I tried a similar approach to your suggestion in this thread:<br>
at::autocast::set_enabled(true);<br>
torch::Tensor result = model(source[b]);<br>
at::autocast::clear_cache();<br>
at::autocast::set_enabled(false);</p>
<p>This get me no speed difference at all, so I’m wondering if it’s actually used or not. Any other pointers here?</p>
<p>Cheers,<br>
David</p>

davidw82: <p>Just to follow up on my own post from yesterday. I posted a bit too early, and had a classic coding mistake present that made the provided code not run. I’ve got it running correctly now, and I do experience around a 15% increase in performance compared to running my model in full float (an a RTX2080Ti card). Thanks all!</p>

JamesDickens: <p>Have you tried this with torchscript models as well?</p>