Sigmoid vs Binary Cross Entropy Loss

Question

In my torch model, the last layer is a torch.nn.Sigmoid() and the loss is the torch.nn.BCELoss. In the training step, the following error has occurred:

RuntimeError: torch.nn.functional.binary_cross_entropy and torch.nn.BCELoss are unsafe to autocast.
Many models use a sigmoid layer right before the binary cross entropy layer.
In this case, combine the two layers using torch.nn.functional.binary_cross_entropy_with_logits
or torch.nn.BCEWithLogitsLoss.  binary_cross_entropy_with_logits and BCEWithLogits are
safe to autocast.

However, when trying to reproduce this error while computing the loss and backpropagation, everything goes correctly:

import torch
from torch import nn

# last layer
sigmoid = nn.Sigmoid()
# loss
bce_loss = nn.BCELoss()


# the true classes
true_cls = torch.tensor([
            [0.],
            [1.]])

# model prediction classes
pred_cls = sigmoid(
    torch.tensor([
           [0.4949],
           [0.4824]],requires_grad=True)
)
pred_cls
# tensor([[0.6213],
#         [0.6183]], grad_fn=<SigmoidBackward>)

out = bce_loss(pred_cls, true_cls)
out
# tensor(0.7258, grad_fn=<BinaryCrossEntropyBackward>)

out.backward()

What am i missing? I appreciate any help you can provide.

Answer 1

You have to move it to cuda first and enable the autocast, like this:

import torch
from torch import nn
from torch.cuda.amp import autocast

# last layer
sigmoid = nn.Sigmoid().cuda()
# loss
bce_loss = nn.BCELoss().cuda()


# the true classes
true_cls = torch.tensor([
            [0.],
            [1.]]).cuda()

with autocast():

    # model prediction classes
    pred_cls = sigmoid(
        torch.tensor([
               [0.4949],
               [0.4824]], requires_grad=True

        ).cuda()
    )

    pred_cls
    # tensor([[0.6213],
    #         [0.6183]], grad_fn=<SigmoidBackward>)

    out = bce_loss(pred_cls, true_cls)
    out
    # tensor(0.7258, grad_fn=<BinaryCrossEntropyBackward>)

    out.backward()

RuntimeError: torch.nn.functional.binary_cross_entropy and torch.nn.BCELoss are unsafe to autocast.
Many models use a sigmoid layer right before the binary cross entropy layer.
In this case, combine the two layers using torch.nn.functional.binary_cross_entropy_with_logits
or torch.nn.BCEWithLogitsLoss.  binary_cross_entropy_with_logits and BCEWithLogits are
safe to autocast.