Autocast with BCELoss() on CPU

HaoY: <p>It is possible use <code>bfloat16</code> training for a model with binary cross entropy loss?</p>
<p>The following code will throw out runtime error: “Found dtype Float but expected BFloat16” because my y_true from my dataloader is torch.float32 type.</p>
<pre><code class="lang-auto">loss_fun = torch.nn.BCELoss()
with torch.autocast(device_type='cpu'):
            y_pred = model(x)
            loss = loss_fun(y_pred, y_true)
</code></pre>
<p>If I manually do <code>y_true = y_true.to(torch.bfloat16)</code> before the <code>with</code> statement with the following codes:</p>
<pre><code class="lang-auto">loss_fun = torch.nn.BCELoss()
y_true = y_true.to(torch.bfloat16)
with torch.autocast(device_type='cpu'):
            y_pred = model(x)
            loss = loss_fun(y_pred, y_true)
</code></pre>
<p>I got RuntimeError: “binary_cross_entropy” not implemented for ‘BFloat16’.</p>
<p>From the official doc <a href="https://pytorch.org/docs/master/amp.html#cpu-op-specific-behavior" class="inline-onebox" rel="noopener nofollow ugc">Automatic Mixed Precision package - torch.amp — PyTorch master documentation</a>, <code>binary_cross_entropy</code> is listed under <code>float32</code>  CPU Ops.   I also read from the same doc that</p>
<blockquote>
<p>When entering an autocast-enabled region, Tensors may be any type. You should not call <code>half()</code> or <code>bfloat16()</code> on your model(s) or inputs when using autocasting.</p>
</blockquote>
<p>But how should I code such that <code>model(x)</code> uses <code>bfloat16</code> and then later <code>loss_fun(y_pred, y_true)</code> uses <code>float32</code>.</p>

ptrblck: <p>You haven’t posted a code snippet reproducing the issue, but this small example works for me:</p>
<pre data-code-wrap="python"><code class="lang-python">device = "cpu"
x = torch.randn(1, 10, device=device)
y_true = torch.empty(1, 10, device=device).uniform_(0, 1)
model = nn.Linear(10, 10).to(device)
loss_fun = torch.nn.BCELoss()

with torch.autocast(device_type=device, dtype=torch.bfloat16):
    y_pred = model(x).sigmoid()
    loss = loss_fun(y_pred, y_true)
loss.backward()

print(model.weight.grad.abs().sum())
# tensor(1.3052)
</code></pre>

HaoY: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thanks for your prompt response.</p>
<p>I just realized I made a silly mistake:  I forgot my training codes have a utility module that use <code>y_pred</code> and defined <code>loss_fun</code> to re-calculate the loss together with other metrics for tracking/logging.   I did not wrap the loss computation logic in that module under <code>autocast</code> though that module also consumes the same <code>y_pred</code> of type <code>bfloat16</code>.</p>
<p>The errors that I reported in my original post were actually cast from that utility module. If I don’t explicitly convert y_true to <code>bfloat16</code>, <code>loss_fun</code> got two input arguments of different dtypes.  If I explicitly convert y_true to <code>bfloat16</code>, both input arguments are <code>bfloat16</code> but <code>binary_cross_entropy</code> can not process <code>bfloat16</code> because it is not under <code>autocast</code>.</p>
<p>In conclusion,  autocast should work well with BCELoss() on CPUs.</p>