fix some issues with ddpm 1d

Author: lucidrains

README.md

@@ -45,6 +45,7 @@ diffusion = GaussianDiffusion(
 training_images = torch.rand(8, 3, 128, 128) # images are normalized from 0 to 1
 loss = diffusion(training_images)
 loss.backward()
+
 # after a lot of training
 
 sampled_images = diffusion.sample(batch_size = 4)

denoising_diffusion_pytorch/__init__.py

@@ -6,5 +6,4 @@
 from denoising_diffusion_pytorch.elucidated_diffusion import ElucidatedDiffusion
 from denoising_diffusion_pytorch.v_param_continuous_time_gaussian_diffusion import VParamContinuousTimeGaussianDiffusion
 
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch_1d import GaussianDiffusion1D, Unet1D, Trainer1D
-
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch_1d import GaussianDiffusion1D, Unet1D, Trainer1D, Dataset1D

denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -1,21 +1,22 @@
 import math
-from multiprocessing import cpu_count
 from pathlib import Path
 from random import random
 from functools import partial
 from collections import namedtuple
+from multiprocessing import cpu_count
 
 import torch
-from accelerate import Accelerator
-from ema_pytorch import EMA
-from torch import nn, einsum
+from torch import nn, einsum, Tensor
 import torch.nn.functional as F
 from torch.cuda.amp import autocast
+from torch.optim import Adam
+from torch.utils.data import Dataset, DataLoader
 
 from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
-from torch.optim import Adam
-from torch.utils.data import Dataset, DataLoader
+
+from accelerate import Accelerator
+from ema_pytorch import EMA
 
 from tqdm.auto import tqdm
 
@@ -67,6 +68,19 @@ def normalize_to_neg_one_to_one(img):
 def unnormalize_to_zero_to_one(t):
     return (t + 1) * 0.5
 
+# data
+
+class Dataset1D(Dataset):
+    def __init__(self, tensor: Tensor):
+        super().__init__()
+        self.tensor = tensor.clone()
+
+    def __len__(self):
+        return len(self.tensor)
+
+    def __getitem__(self, idx):
+        return self.tensor[idx].clone()
+
 # small helper modules
 
 class Residual(nn.Module):
@@ -714,7 +728,7 @@ def __init__(
         num_samples = 25,
         results_folder = './results',
         amp = False,
-        fp16 = False,
+        mixed_precision_type = 'fp16',
         split_batches = True,
     ):
         super().__init__()
@@ -723,11 +737,9 @@ def __init__(
 
         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

denoising_diffusion_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.8.2'
+__version__ = '1.8.3'