NaNs in torch.nn.functional.grid_sample with Mixed Precision

tyb_10: <p>Hi,<br>
I get NaNs when I run my training in Mixed precision training mode.<br>
I have pointed the error at the following place</p>
<pre><code class="lang-auto">torch.nn.functional.grid_sample(input, grid, padding_mode="zeros",
                    mode='bilinear',
                    align_corners=True,)
</code></pre>
<p>when I change the sampling mode to <code>nearest</code>, it does not happen anymore.</p>
<pre><code class="lang-auto">torch.nn.functional.grid_sample(input, grid, padding_mode="zeros",
                    mode='nearest',
                    align_corners=True,)
</code></pre>
<p>both sampling modes work fine in 32 bit training.<br>
My inputs and the grid does not contain NaNs or Infs , what could be the possible reason for this and the solution.</p>

ptrblck: <p>Could you post the shapes as well as the min. and max. values for <code>input</code> and <code>grid</code> so that we can reproduce this issue, as this code works fine in <code>float16</code>:</p>
<pre><code class="lang-python">input = torch.arange(4*4).view(1, 1, 4, 4).half().cuda()
print(input)

# Create grid to upsample input
d = torch.linspace(-1, 1, 8)
meshx, meshy = torch.meshgrid((d, d))
grid = torch.stack((meshy, meshx), 2)
grid = grid.unsqueeze(0) # add batch dim
grid = grid.half().cuda()

output = torch.nn.functional.grid_sample(input, grid, mode='bilinear')
print(torch.isfinite(output).all())
</code></pre>

tyb_10: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>  thanks for the kind feedback.<br>
Actually on investigation Grid infact does have the <code>inf</code> as maximum value, sorry for the incorrect information above.<br>
Then does it have undefined behavior or the following still applies according to the documentation ?</p>
<pre><code class="lang-auto">If grid has values outside the range of [-1, 1], the corresponding outputs are handled as defined by padding_mode. Options are .... 
</code></pre>
<pre><code class="lang-auto">input size: torch.Size([1, 32, 124, 384])
input min , max: 0.0,15.419422149658203
Grid size: torch.Size([1, 16, 200, 2])
Grid min, max: -0.9350162744522095,inf
</code></pre>
<p>and what could be done in this case</p>
<p>will it be valid to do following<br>
<code>grid[grid[..., :] &gt; 1.0] = 2.0</code></p>

ptrblck: <p>It depends on your use case what the valid workaround would be.<br>
I.e. why are <code>Inf</code> values in the grid and what would they mean?<br>
If you think these grid values should use the padding, your workaround might work, on the other hand you might want to investigate why these <code>Inf</code> values are created and avoid them.</p>

tyb_10: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Any pointers as to why grid would have <code>Inf</code> during mixed precision training only ?</p>
<p>following is the exact line where <code>Infs</code>  happen in half precision</p>
<p><code>pts_2d = torch.mm(pts_3d_rect, torch.transpose(proj_matrix, 0, 1))</code></p>

ptrblck: <p><code>float16</code> can easily overflow if you are using values with a value close to the min. and max. values:</p>
<pre><code class="lang-python">torch.finfo(torch.float16).max
&gt; 65504.0
</code></pre>
<p>E.g. this code snippet overflows in the second approach and yields <code>Infs</code> in the result after applying <code>torch.mm</code>:</p>
<pre><code class="lang-python">x = torch.randn(1024, 1024, device='cuda').half()
y = torch.mm(x, x)
print(torch.isfinite(y).all())
&gt; tensor(True, device='cuda:0')

x = torch.randn(1024, 1024, device='cuda').half() * 2**13
print(torch.isfinite(x).all())
&gt; tensor(True, device='cuda:0')

y = torch.mm(x, x)
print(torch.isfinite(y).all())
&gt; tensor(False, device='cuda:0')
</code></pre>
<p>Since the <code>grid</code> would contain valid values in <code>[-1, 1]</code>, this would usually not be a problem.<br>
However, it seems you are expecting to work with large values (which would then use the padding values), so you might disable <code>autocast</code> for the <code>grid</code> sampling operation and the <code>grid</code> creation.</p>

tyb_10: <p>thanks so much for the suggestion of disabling <code>autocast</code> , that indeed is nice solution.<br>
and also i found out that it depends on the input data , it overflows for certain datasets (with larger values for camera intrinsics values for example.)</p>

tyb_10: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
actually coming back to this, because explicit autocasting does not sit well while converting the model to torchscript and complains</p>
<pre><code class="lang-auto">torch.jit.frontend.UnsupportedNodeError: function definitions aren't supported:
</code></pre>
<pre><code class="lang-auto">def foo(
    x: torch.Tensor, y: torch.Tensor
) -&gt; torch.Tensor:
    with autocast(enabled=False):
       ..........
    return z
</code></pre>
<p>how can we make autocast context manager and torchscript conversion co-exist ?</p>

ptrblck: <p>Automatic mixed-precision is not supported yet in torchscript and we are working on it.<br>
If your model doesn’t use data dependent conditions etc. you might be able to use tracing (although you should verify the correctness).</p>