Repository: espnet/espnet

Issue:
loss diverged to nan when training vits with fp16 option enabled.
I trained VITS using the original dataset with the fp16 option enabled (--use_amp true), but after turning the training around for a few days, the loss values diverged to NAN in the process.
I think the exception handling when values outside the range of fp16 is not working properly.
I did not see this error when the fp16 option was not enabled for the same data set.


Discussion:
Thank you for your report. Actually, I have not tried fp16 setting so much.
Do you have any comments about doubtful points?
To be honest, the learning process was going well until the middle of the program, so I am not sure what exactly is wrong.
https://github.com/jaywalnut310/vits/issues/15
https://github.com/jaywalnut310/vits/issues/48
These are the official implementations, and a bug was reported here as well, where learning was not possible or loss was NaN when fp16 was enabled.

In my case, I modified espnet's VITS implementation slightly based on the PR, to be able to learn... but I don't know how it effect to result.
https://github.com/jaywalnut310/vits/pull/34



Resolving patch:
diff --git a/espnet2/gan_tts/vits/vits.py b/espnet2/gan_tts/vits/vits.py
--- a/espnet2/gan_tts/vits/vits.py
+++ b/espnet2/gan_tts/vits/vits.py
@@ -3,6 +3,8 @@
 
 """VITS module for GAN-TTS task."""
 
+from contextlib import contextmanager
+from distutils.version import LooseVersion
 from typing import Any
 from typing import Dict
 from typing import Optional
@@ -37,6 +39,14 @@
     "hifigan_multi_scale_multi_period_discriminator": HiFiGANMultiScaleMultiPeriodDiscriminator,  # NOQA
 }
 
+if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
+    from torch.cuda.amp import autocast
+else:
+    # Nothing to do if torch<1.6.0
+    @contextmanager
+    def autocast(enabled=True):  # NOQA
+        yield
+
 
 class VITS(AbsGANTTS):
     """VITS module (generator + discriminator).
@@ -398,18 +408,19 @@ def _forward_generator(
             p = self.discriminator(speech_)
 
         # calculate losses
-        mel_loss = self.mel_loss(speech_hat_, speech_)
-        kl_loss = self.kl_loss(z_p, logs_q, m_p, logs_p, z_mask)
-        dur_loss = torch.sum(dur_nll.float())
-        adv_loss = self.generator_adv_loss(p_hat)
-        feat_match_loss = self.feat_match_loss(p_hat, p)
-
-        mel_loss = mel_loss * self.lambda_mel
-        kl_loss = kl_loss * self.lambda_kl
-        dur_loss = dur_loss * self.lambda_dur
-        adv_loss = adv_loss * self.lambda_adv
-        feat_match_loss = feat_match_loss * self.lambda_feat_match
-        loss = mel_loss + kl_loss + dur_loss + adv_loss + feat_match_loss
+        with autocast(enabled=False):
+            mel_loss = self.mel_loss(speech_hat_, speech_)
+            kl_loss = self.kl_loss(z_p, logs_q, m_p, logs_p, z_mask)
+            dur_loss = torch.sum(dur_nll.float())
+            adv_loss = self.generator_adv_loss(p_hat)
+            feat_match_loss = self.feat_match_loss(p_hat, p)
+
+            mel_loss = mel_loss * self.lambda_mel
+            kl_loss = kl_loss * self.lambda_kl
+            dur_loss = dur_loss * self.lambda_dur
+            adv_loss = adv_loss * self.lambda_adv
+            feat_match_loss = feat_match_loss * self.lambda_feat_match
+            loss = mel_loss + kl_loss + dur_loss + adv_loss + feat_match_loss
 
         stats = dict(
             generator_loss=loss.item(),
@@ -504,8 +515,9 @@ def _forward_discrminator(
         p = self.discriminator(speech_)
 
         # calculate losses
-        real_loss, fake_loss = self.discriminator_adv_loss(p_hat, p)
-        loss = real_loss + fake_loss
+        with autocast(enabled=False):
+            real_loss, fake_loss = self.discriminator_adv_loss(p_hat, p)
+            loss = real_loss + fake_loss
 
         stats = dict(
             discriminator_loss=loss.item(),


Test patch: