Hi @ptrblck,
Thanks so much for your response. I decreased the learning rate by an order of magnitude as per someone’s suggestion on a related GitHub Issue, and this seemed to resolve the problem! And thank you for your suggestion – I will take a look into the native amp implementation.
I do have a quick follow-up though: Given that the original loss trained well with a learning rate of 0.01, and the new loss trains well with a different learning rate of 0.001, I thought it would make sense to use two separate optimizers in the following manner:
optimizer = torch.optim.SGD(model.parameters(), 
                            0.01, 
                            momentum=args.momentum, 
                            weight_decay=args.weight_decay)
parameters_contrastive = list(model.conv2.parameters()) + 
                         list(model.maxpool2.parameters()) + 
                         list(model.conv3.parameters())
optimizer_contrastive = torch.optim.SGD(parameters_contrastive, 
                                        0.001,
                                        momentum=args.momentum,
                                        weight_decay=args.weight_decay)

model, [optimizer, optimizer_contrastive] = amp.initialize(model, 
                                                           [optimizer, optimizer_contrastive],
                                                           opt_level=args.opt_level,
                                                           keep_batchnorm_fp32=args.keep_batchnorm_fp32,
                                                           loss_scale=args.loss_scale,
                                                           num_losses=2)

and in the train loop:
output, rep3, rep4, rep5 = model(input)

loss = criterion(output, target)

# compute contrastive loss
loss_contrastive = getContrastiveLoss(target, rep3, rep4, rep5, contrastive_idxs)
loss_contrastive *= 0.1

# compute gradient and do SGD step
optimizer.zero_grad()
optimizer_contrastive.zero_grad()

# Mixed-precision training requires that the loss is scaled in order
# to prevent the gradients from underflow
with amp.scale_loss(loss, optimizer, loss_id=0) as scaled_loss:
    scaled_loss.backward(retain_graph=True)
with amp.scale_loss(loss_contrastive, optimizer_contrastive, loss_id=1) as scaled_loss:
    scaled_loss.backward()

optimizer.step()
optimizer_contrastive.step()

While both losses (loss and loss_contrastive) start off at the same order of magnitude (after loss_contrastive is scaled by 0.1), loss_contrastive is reduced very smoothly throughout training while loss remains pretty much constant. Again, just to reiterate from my original post, when this exact model is trained without the additional loss_contrastive, the loss is reduced very smoothly and the AlexNet is able to achieve a final top-5 accuracy of 80% on ImageNet.
I am unsure how the addition of my loss_contrastive is discouraging the network from optimizing the original, main loss. Do you have any idea why this might be the case, or what I can do to combat this?
Thanks once again for your help; You have helped me numerous times in the past as well and it always stuns me to see how people of your calibre are so generous and accessible even to beginners. I truly appreciate all your effort.