Mixed precision inference on ARM servers

anijain2305: <p>Hi,</p>
<p>My usecase is to take a FP32 pre-trained PyTorch model, convert it to FP16 (both weights and computation that is amenable to Fp16 computation) and then trace the model. Later, I will read this model in TVM (a deep learning compiler) and use it to generate code for ARM servers. ARM servers have instructions to speed up FP16 computation. Please let me know if this is possible today.</p>
<p>Note that I need to use Torchscript (as TVM input is traced model). Also, the goal here is get speedup so my hope is that the resulting traced model has some operators (like conv2d, dense) whose inputs dtypes are in FP16.</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a></p>

ptrblck: <p><code>autocast</code> is not supported yet in scripted models, but in case your model doesn’t use any data dependent control flow, you could try to trace the model inside the <code>autocast</code> context.<br>
I don’t know, how the TorchScript model would be provided to TVM and what kind of instructions ARM servers are using.</p>

anijain2305: <p>I see. Thanks for the reply <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
I will read more about <code>autocast</code> tomorrow and will try tracing it.</p>
<p>Meanwhile, I also used the following script (using <code>half</code>) which gave me a <code>fp16</code> graph</p>
<pre><code class="lang-auto">import torchvision
import torch
import tvm
from tvm import relay

is_fp16 = True

model = torchvision.models.resnet18(pretrained=True)
x = torch.rand(1, 3, 224, 224)
if is_fp16:
    model = model.cuda().half()
    x = x.cuda().half()
    scripted_model = torch.jit.trace(model, (x))
    scripted_model.save("rn50_fp16.pt")
else:
    scripted_model = torch.jit.trace(model, (x))
    scripted_model.save("rn50_fp32.pt")

input_name = "input0"
input_shape = (1, 3, 224, 224)
shape_list = [(input_name, input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
mod = relay.transform.InferType()(mod)
print(mod)
</code></pre>
<p>My usecase is inference. The above code converted the whole model to <code>fp16</code>. But I am worried that it might be doing to much of <code>fp16</code> and may need to go back to <code>fp32</code> for some operators like <code>sum</code> or <code>batch_norm</code>.</p>

ptrblck: <p>That’s a reasonable concern and is the reason for the implementation of the <code>amp.autocast</code>, which would use FP32 precision, where necessary. For your use case you could certainly try to call <code>.half()</code> on the entire model and check the results manually.</p>

anijain2305: <p>Makes sense. So, I followed up on <code>autocast</code>. This is what my code looks like</p>
<pre><code class="lang-auto">model = torchvision.models.resnet18()

x = torch.rand(1, 3, 224, 224).cuda()
model = model.cuda()
with torch.cuda.amp.autocast(enabled=True):
    scripted_model = torch.jit.trace(model, (x))
    scripted_model.save("rn50_amp.pt")
</code></pre>
<p>It failed with the following error - <code>RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient</code></p>
<p>I then added no_grad</p>
<pre><code class="lang-auto">model = torchvision.models.resnet18()

x = torch.rand(1, 3, 224, 224).cuda()
model = model.cuda()
with torch.cuda.amp.autocast(enabled=True):
    with torch.no_grad():
        scripted_model = torch.jit.trace(model, (x))
        scripted_model.save("rn50_amp.pt")
</code></pre>
<p>But this also failed with some different error</p>
<pre><code class="lang-auto">torch.jit._trace.TracingCheckError: Tracing failed sanity checks!
ERROR: Graphs differed across invocations!
</code></pre>
<p>I am not sure if my usage is correct.</p>

ptrblck: <p>Thanks for the update. I guess that some (new) graph passes might be running into these issues, but I’m unsure which optimizations are performed. Until scripting + amp is fully supported, you could manually cast the model and check the outputs.</p>

anijain2305: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> for the prompt responses. You have been really helpful. I will manually cast the models till then.</p>