Hi, I modified the original PyTorch densenet121 to extract features at different scales. Here 1 is the modified densenet; it still can load the weights provided by PyTorch. I cannot send you our exact decoder, but I guess any SegNet-like architecture for pixel-wise segmentation should show the same behavior.
To check if the error was caused by my code, I used this repo, which is a deeplab v3 implementation with pretrained weights for cityscapes. I tested torch.cuda.amp, O1 from apex.amp, and no mixed precision (the last one was tested with half of the size of the input image used by the mixed precision tests).
I’ve run some epochs and found that torch.cuda.amp is twice as slow as apex.amp, and no mixed precision. In numbers, 1 iteration (forward and backward) using torch.cuda.amp takes ~2.50 seconds, the other two take around ~1.20 seconds.
Configuration:
batch size: 2 per GPU
Image size: 512x512 (mixed precision), 256x512 (no mixed precision)
Batch sync: Yes
Dataset: Cityscapes semantic segmentation
Code: https://github.com/nyoki-mtl/pytorch-segmentation
optimizer: adam
loss: cross entropy
System:
Ubuntu 18.04
Cuda: 10.2
CudNN: 8.0.3
Pytorch: 1.6
GPU: 2 GTX 1080Ti
If you need more details please let me know. Thanks for the help .

p.s.
Just to update more information about Apex. I’m still getting gradient overflow and scale_loss being reduce to zero when using apex.parallel.convert_syncbn_model() after few epochs.