another fix for training under amp

Author: lucidrains

denoising_diffusion_pytorch/classifier_free_guidance.py

@@ -9,6 +9,7 @@
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
+from torch.cuda.amp import autocast
 
 from einops import rearrange, reduce, repeat
 from einops.layers.torch import Rearrange
@@ -711,6 +712,7 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
 
         return img
 
+    @autocast(enabled = False)
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -3,6 +3,7 @@
 from torch import sqrt
 from torch import nn, einsum
 import torch.nn.functional as F
+from torch.cuda.amp import autocast
 from torch.special import expm1
 
 from tqdm import tqdm
@@ -233,6 +234,7 @@ def sample(self, batch_size = 16):
 
     # training related functions - noise prediction
 
+    @autocast(enabled = False)
     def q_sample(self, x_start, times, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -706,6 +706,7 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
 
         return img
 
+    @autocast(enabled = False)
     def q_sample(self, x_start, t, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
@@ -823,23 +824,22 @@ def __init__(
         num_samples = 25,
         results_folder = './results',
         amp = False,
-        fp16 = False,
+        mixed_precision_type = 'fp16',
         split_batches = True,
         convert_image_to = None,
         calculate_fid = True,
-        inception_block_idx = 2048
+        inception_block_idx = 2048,
+        max_grad_norm = 1.
     ):
         super().__init__()
 
         # accelerator
 
         self.accelerator = Accelerator(
             split_batches = split_batches,
-            mixed_precision = 'fp16' if fp16 else 'no'
+            mixed_precision = mixed_precision_type if amp else 'no'
         )
 
-        self.accelerator.native_amp = amp
-
         # model
 
         self.model = diffusion_model
@@ -867,6 +867,8 @@ def __init__(
         self.train_num_steps = train_num_steps
         self.image_size = diffusion_model.image_size
 
+        self.max_grad_norm = max_grad_norm
+
         # dataset and dataloader
 
         self.ds = Dataset(folder, self.image_size, augment_horizontal_flip = augment_horizontal_flip, convert_image_to = convert_image_to)
@@ -980,7 +982,7 @@ def train(self):
 
                     self.accelerator.backward(loss)
 
-                accelerator.clip_grad_norm_(self.model.parameters(), 1.0)
+                accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
                 pbar.set_description(f'loss: {total_loss:.4f}')
 
                 accelerator.wait_for_everyone()

denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -10,6 +10,7 @@
 from ema_pytorch import EMA
 from torch import nn, einsum
 import torch.nn.functional as F
+from torch.cuda.amp import autocast
 
 from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
@@ -639,6 +640,7 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
 
         return img
 
+    @autocast(enabled = False)
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 

denoising_diffusion_pytorch/guided_diffusion.py

@@ -9,6 +9,7 @@
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
+from torch.cuda.amp import autocast
 from torch.utils.data import Dataset, DataLoader
 
 from torch.optim import Adam
@@ -708,6 +709,7 @@ def interpolate(self, x1, x2, t = None, lam = 0.5):
 
         return img
 
+    @autocast(enabled = False)
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 

denoising_diffusion_pytorch/simple_diffusion.py

@@ -6,6 +6,7 @@
 from torch import nn, einsum
 import torch.nn.functional as F
 from torch.special import expm1
+from torch.cuda.amp import autocast
 
 from tqdm import tqdm
 from einops import rearrange, repeat, reduce, pack, unpack
@@ -653,6 +654,7 @@ def sample(self, batch_size = 16):
 
     # training related functions - noise prediction
 
+    @autocast(enabled = False)
     def q_sample(self, x_start, times, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 

denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py

@@ -4,6 +4,7 @@
 from torch import nn, einsum
 import torch.nn.functional as F
 from torch.special import expm1
+from torch.cuda.amp import autocast
 
 from tqdm import tqdm
 from einops import rearrange, repeat, reduce
@@ -149,6 +150,7 @@ def sample(self, batch_size = 16):
 
     # training related functions - noise prediction
 
+    @autocast(enabled = False)
     def q_sample(self, x_start, times, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.7.4'
+__version__ = '1.7.6'