Running UniverSeg pretrained model in mixed precision

[alinaryabtsev]

Hello all,

Recently I have been playing around with this model while achieving nice results. However, I wanted to enhance the results, by enlarging the support set. To do so, I tried to run the model with fp16 data and run the model in a mixed precision setting with pytorch's amp, but the resulting output contained NaNs. This is caused by an overflow occurring in the first decoder layer. Might it be that the model is not compatible with fp16 at all?

Here is a dummy example to reproduce:
model = universeg(pretrained=True).to(device) #%% dummy_image = (torch.zeros((1, 1, 128, 128))).to(torch.float16).to(device) support_labels = (torch.ones((1, 10, 1, 128, 128))).to(torch.float16).to(device) support_images = (torch.ones((1, 10, 1, 128, 128))).to(torch.float16).to(device) with torch.cuda.amp.autocast(): with torch.no_grad(): result_dummy = model(dummy_image, support_images, support_labels)

The output is a tensor which mostly contains values of NaN.

[VictorButoi]

Greetings! Interesting enough, we also have been trying to train things with reducing precisions and have been running into nans.

If you have any leads on this direction we'd be super excited because I believe we are similarly stuck.