Using half precision

s_n: <p>Hi,</p>
<p>I am new to using the half-precision for tensors in PyTorch. So I had a very basic question if it’s possible that in my neural network model I can have some variables as half tensors and some as normal full precision tensors?</p>
<p>Basically my model is taking too much memory so instead of decreasing the batch size, I wanted to check if it’s possible to make some variables as half-precision.</p>
<p>Thanks!</p>

ptrblck: <p>Yes, you can change the <code>dtype</code> of a tensor manually and make sure that the performed operations use expected input types.<br>
While this manual approach would be feasible, I would recommend to have a look at the <a href="https://pytorch.org/docs/stable/amp.html">automatic mixed-precision training</a>, which provides utility functions such as automatic casting and gradient scaling.</p>

s_n: <p>Hi,</p>
<p>While using amp with my custom forward function using autograd.function. I am getting below error:</p>
<p>NameError: name ‘custom_fwd’ is not defined</p>
<p>I have added <span class="mention">@custom_fwd</span> decorator above my forward function.</p>

ptrblck: <p>What does <code>print(torch.cuda.amp.custom_fwd)</code> return?</p>

ericlormul: <p>I encounter similar problem, need to save memory.</p>
<p>I tried automatic mixed-precision training with autocast, but by directly plugging in the autocast, it doesn’t significantly reduce memory usage. I  think it is due to that autocast works better with certain operations.</p>
<p>I also tried model.half().to(‘cuda’) which also doesn’t show significant memory save. Is it supposed to significantly reduce memory used by model?</p>

ptrblck: <p>Yes, calling <code>.half()</code> on the model as well as the inputs would give you an estimate of how much memory would be saved in the extreme case of using <code>float16</code> for all operations. Make sure to check the memory usage via <code>torch.cuda.memory_summary()</code>, since <code>nvidia-smi</code> will also show the cache not only the allocated memory.</p>

ericlormul: <p>I tried with and without .half() on my model, the outputs of torch.cuda.memory_summary() seem to have no difference on Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed. Could you provide an example model that shows significant memory reduce after set data type to half?</p>
<p>Also, is there an documents explaining the meaning of Allocated memory,  Active memory, GPU reserved memory and  Non-releasable memory? Thanks.</p>

ptrblck: <p>The <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management">memory management docs</a> would give you some more information.</p>
<p>I’m not able to reproduce the same memory usage and see the expected reduction by 2 using <code>float16</code> for a single tensor:</p>
<pre><code class="lang-python">x = torch.randn(1024, 1024, dtype=torch.float32, device='cuda') # switch the dtype to torch.float16
</code></pre>
<p>Output:</p>
<pre><code class="lang-python"># float16:
|---------------------------------------------------------------------------|
| Allocated memory      |    2048 KB |    2048 KB |    2048 KB |       0 B  |
|       from large pool |    2048 KB |    2048 KB |    2048 KB |       0 B  |
|       from small pool |       0 KB |       0 KB |       0 KB |       0 B  |
|---------------------------------------------------------------------------|
# float32:
|---------------------------------------------------------------------------|
| Allocated memory      |    4096 KB |    4096 KB |    4096 KB |       0 B  |
|       from large pool |    4096 KB |    4096 KB |    4096 KB |       0 B  |
|       from small pool |       0 KB |       0 KB |       0 KB |       0 B  |
|---------------------------------------------------------------------------|
</code></pre>
<p>as well as <code>torchvision.models.resnet18</code>:</p>
<pre><code class="lang-python"># float16:
|---------------------------------------------------------------------------|
| Allocated memory      |   23768 KB |   23768 KB |   23768 KB |       0 B  |
|       from large pool |   20480 KB |   20480 KB |   20480 KB |       0 B  |
|       from small pool |    3288 KB |    3288 KB |    3288 KB |       0 B  |
|---------------------------------------------------------------------------|
# float32:
|---------------------------------------------------------------------------|
| Allocated memory      |   45763 KB |   45763 KB |   45763 KB |       0 B  |
|       from large pool |   42368 KB |   42368 KB |   42368 KB |       0 B  |
|       from small pool |    3395 KB |    3395 KB |    3395 KB |       0 B  |
|---------------------------------------------------------------------------|
</code></pre>