RuntimeError: expected scalar type Float but found Half in deform_conv2d

jsetty: <p>Hi, I have a huge network with some Deformable CNNs (torch.ops.torchvision.deform_conv2d) in it. I am using apex 0.1, Cuda 11, and torch 1.11. I have set the optimization level to O1 for mixed-precision training.</p>
<p>I have used nn.DataParallel and also cast the model to GPU properly. Unfortunately, I have this error. Could you please help me fix this?</p>
<pre><code class="lang-auto">File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/container.py", line 141, in forward
    input = module(input)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "&lt;path&gt;/model/dcn.py", line 56, in forward
    x = torchvision.ops.deform_conv2d(input=x,
  File "/opt/conda/lib/python3.8/site-packages/torchvision/ops/deform_conv.py", line 89, in deform_conv2d
    return torch.ops.torchvision.deform_conv2d(
RuntimeError: expected scalar type Float but found Half
</code></pre>
<p>Thanks!</p>

ptrblck: <p><code>apex.amp</code> is deprecated so please use the native mixed-precision training util. via <code>torch.cuda.amp</code>. You can find examples <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples">here</a>.</p>

jsetty: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Thanks for your inputs. I moved to torch.cuda.amp.</p>
<p>I am using <code>scaler = torch.cuda.amp.GradScaler()</code>.</p>
<p>I am doing the casting as following:</p>
<pre><code class="lang-auto">with torch.cuda.amp.autocast():
    preds = model(inputs)
    loss = criterion(preds, labels.float())
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
</code></pre>
<p>But I have an error in a different module for torch.einsum().</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "main.py", line 333, in &lt;module&gt;
    main()
  File "main.py", line 327, in main
    train(model, train_loader, test_loader, args, criterion, optimizer, device)
  File "main.py", line 229, in train
    scaler.scale(loss).backward()
  File "/opt/conda/lib/python3.8/site-packages/torch/_tensor.py", line 352, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/function.py", line 199, in apply
    return user_fn(self, *args)
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/function.py", line 340, in wrapper
    outputs = fn(ctx, *args)
  File "&lt;path&gt;/model/pac.py", line 179, in backward
    grad_in_mul_k = torch.einsum('iomn,ojkl-&gt;ijklmn', (grad_output, weight))
  File "/opt/conda/lib/python3.8/site-packages/torch/functional.py", line 325, in einsum
    return einsum(equation, *_operands)
  File "/opt/conda/lib/python3.8/site-packages/torch/functional.py", line 327, in einsum
    return _VF.einsum(equation, operands)  # type: ignore[attr-defined]
RuntimeError: expected scalar type Half but found Float
</code></pre>
<p>Please take a look.</p>
<p>Thanks!</p>

ptrblck: <p>I guess you might be using an older <code>torchvision</code> release, since the operation is available via <code>torchvision.ops.DeformConv2d</code> and <a href="https://github.com/pytorch/vision/pull/2898">this PR</a> added <code>autocast</code> support ~1 year ago.</p>
<p>This code snippet works fine for me:</p>
<pre><code class="lang-python">from torchvision import ops

deform_layer = ops.DeformConv2d(in_channels=3, out_channels=64, kernel_size=3).cuda()
x = torch.randn((1,3,9,9)).cuda()
offset = torch.randn(1,3*3*2,7,7).cuda()
with torch.cuda.amp.autocast():
    out = deform_layer(x, offset)
</code></pre>