Amp not working for spectral_norm?

[TDeVries]

Is spectral_norm supported by amp? I've recently been trying to get amp working on a GAN, but so far no luck.

Traceback (most recent call last):
 File "train.py", line 14, in <module>
   trainer.train()
 File "/private/home/tldevries/projects/self-attention-GAN-pytorch/trainer_fp16.py", line 177, in train
   d_out_real = self.D(real_images + inst_noise, real_labels)
 File "/private/home/tldevries/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
   result = self.forward(*input, **kwargs)
 File "/private/home/tldevries/projects/self-attention-GAN-pytorch/sagan_models.py", line 240, in forward
   h0 = self.opt_block1(x) # n x d_conv_dim   x 64 x 64
 File "/private/home/tldevries/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 489, in __call__
   result = self.forward(*input, **kwargs)
 File "/private/home/tldevries/projects/self-attention-GAN-pytorch/sagan_models.py", line 173, in forward
   x = self.snconv2d1(x)
 File "/private/home/tldevries/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 485, in __call__
   hook(self, input)
 File "/private/home/tldevries/miniconda3/lib/python3.7/site-packages/torch/nn/utils/spectral_norm.py", line 100, in __call__
   setattr(module, self.name, self.compute_weight(module, do_power_iteration=module.training))
 File "/private/home/tldevries/miniconda3/lib/python3.7/site-packages/torch/nn/utils/spectral_norm.py", line 86, in compute_weight
   sigma = torch.dot(u, torch.mv(weight_mat, v))
RuntimeError: Expected object of scalar type Float but got scalar type Half for argument #2 'tensor'

[mcarilli]

GANs are tricky. I'm working on an example that uses the new unified mixed precision API.
https://nvidia.github.io/apex/amp.html
https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users
Your problem may or may not be GAN-specific. Give me a couple days to finish the new-API GAN example. We can try to update yours based on that, and if it still fails, dig into spectral_norm itself.

[mcarilli]

Deduplicating to #179. If the problem turns out to be spectral norm, rather than the GAN control flow, we can reopen this.

[shuuki4]

@mcarilli I think this problem comes from the torch.dot function in SpectralNorm, not the GAN control flow. I am not sure whether this is a right solution, but adding 'dot' to type promoting functions as in #302 solved this issue.

[mcarilli]

Yes, that's a good solution. d627172