default the device to the device that the diffusion prior parameters are on, if the trainer was never given the accelerator nor device

README.md

@@ -20,18 +20,20 @@ As of 5/23/22, it is no longer SOTA. SOTA will be <a href="https://github.com/lu
 
 - Decoder is now verified working for unconditional generation on my experimental setup for Oxford flowers. 2 researchers have also confirmed Decoder is working for them.
 
-<img src="./samples/oxford.png" width="600px" />
+<img src="./samples/oxford.png" width="450px" />
 
 *ongoing at 21k steps*
 
 - <a href="https://twitter.com/Buntworthy/status/1529475416775434240?t=0GEge3Kr9I36cjcUVCQUTg">Justin Pinkney</a> successfully trained the diffusion prior in the repository for his CLIP to Stylegan2 text-to-image application
 
+- <a href="https://github.com/rom1504">Romain</a> has scaled up training to 800 GPUs with the available scripts without any issues
+
 ## Pre-Trained Models
 
 - LAION is training prior models. Checkpoints are available on <a href="https://huggingface.co/zenglishuci/conditioned-prior">🤗huggingface</a> and the training statistics are available on <a href="https://wandb.ai/nousr_laion/conditioned-prior/reports/LAION-DALLE2-PyTorch-Prior--VmlldzoyMDI2OTIx">🐝WANDB</a>.
 - Decoder - <a href="https://wandb.ai/veldrovive/dalle2_train_decoder/runs/jkrtg0so?workspace=user-veldrovive">In-progress test run</a> 🚧
 - Decoder - <a href="https://wandb.ai/veldrovive/dalle2_train_decoder/runs/3d5rytsa?workspace=">Another test run with sparse attention</a>
-- DALL-E 2 🚧
+- DALL-E 2 🚧 - <a href="https://github.com/LAION-AI/dalle2-laion">DALL-E 2 Laion repository</a>
 
 ## Appreciation
 

dalle2_pytorch/dalle2_pytorch.py

@@ -63,11 +63,16 @@ def default(val, d):
         return val
     return d() if callable(d) else d
 
-def cast_tuple(val, length = 1):
+def cast_tuple(val, length = None):
     if isinstance(val, list):
         val = tuple(val)
 
-    return val if isinstance(val, tuple) else ((val,) * length)
+    out = val if isinstance(val, tuple) else ((val,) * default(length, 1))
+
+    if exists(length):
+        assert len(out) == length
+
+    return out
 
 def module_device(module):
     return next(module.parameters()).device
@@ -330,6 +335,10 @@ def approx_standard_normal_cdf(x):
 def discretized_gaussian_log_likelihood(x, *, means, log_scales, thres = 0.999):
     assert x.shape == means.shape == log_scales.shape
 
+    # attempting to correct nan gradients when learned variance is turned on
+    # in the setting of deepspeed fp16
+    eps = 1e-12 if x.dtype == torch.float32 else 1e-5
+
     centered_x = x - means
     inv_stdv = torch.exp(-log_scales)
     plus_in = inv_stdv * (centered_x + 1. / 255.)
@@ -344,7 +353,7 @@ def discretized_gaussian_log_likelihood(x, *, means, log_scales, thres = 0.999):
         log_cdf_plus,
         torch.where(x > thres,
             log_one_minus_cdf_min,
-            log(cdf_delta)))
+            log(cdf_delta, eps = eps)))
 
     return log_probs
 
@@ -485,14 +494,16 @@ class NoiseScheduler(nn.Module):
 # diffusion prior
 
 class LayerNorm(nn.Module):
-    def __init__(self, dim):
+    def __init__(self, dim, eps = 1e-5):
         super().__init__()
-        self.gamma = nn.Parameter(torch.ones(dim))
-        self.register_buffer("beta", torch.zeros(dim))
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(dim))
 
     def forward(self, x):
-        return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)
-
+        x = x / x.amax(dim = -1, keepdim = True).detach()
+        var = torch.var(x, dim = -1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = -1, keepdim = True)
+        return (x - mean) * (var + self.eps).rsqrt() * self.g
 
 class ChanLayerNorm(nn.Module):
     def __init__(self, dim, eps = 1e-5):
@@ -501,10 +512,10 @@ class ChanLayerNorm(nn.Module):
         self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
 
     def forward(self, x):
+        x = x / x.amax(dim = 1, keepdim = True).detach()
         var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = 1, keepdim = True)
-        return (x - mean) / (var + self.eps).sqrt() * self.g
-
+        return (x - mean) * (var + self.eps).rsqrt() * self.g
 
 class Residual(nn.Module):
     def __init__(self, fn):
@@ -624,10 +635,13 @@ class Attention(nn.Module):
         heads = 8,
         dropout = 0.,
         causal = False,
-        rotary_emb = None
+        rotary_emb = None,
+        pb_relax_alpha = 32 ** 2
     ):
         super().__init__()
-        self.scale = dim_head ** -0.5
+        self.pb_relax_alpha = pb_relax_alpha
+        self.scale = dim_head ** -0.5 * (pb_relax_alpha ** -1)
+
         self.heads = heads
         inner_dim = dim_head * heads
 
@@ -691,7 +705,10 @@ class Attention(nn.Module):
 
         # attention
 
-        attn = sim.softmax(dim = -1, dtype = torch.float32)
+        sim = sim - sim.amax(dim = -1, keepdim = True).detach()
+        sim = sim * self.pb_relax_alpha
+
+        attn = sim.softmax(dim = -1)
         attn = self.dropout(attn)
 
         # aggregate values
@@ -1114,11 +1131,12 @@ class SinusoidalPosEmb(nn.Module):
         self.dim = dim
 
     def forward(self, x):
+        dtype, device = x.dtype, x.device
         half_dim = self.dim // 2
         emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, device = x.device) * -emb)
+        emb = torch.exp(torch.arange(half_dim, device = device, dtype = dtype) * -emb)
         emb = rearrange(x, 'i -> i 1') * rearrange(emb, 'j -> 1 j')
-        return torch.cat((emb.sin(), emb.cos()), dim = -1)
+        return torch.cat((emb.sin(), emb.cos()), dim = -1).type(dtype)
 
 class Block(nn.Module):
     def __init__(
@@ -1205,10 +1223,12 @@ class CrossAttention(nn.Module):
         dim_head = 64,
         heads = 8,
         dropout = 0.,
-        norm_context = False
+        norm_context = False,
+        pb_relax_alpha = 32 ** 2
     ):
         super().__init__()
-        self.scale = dim_head ** -0.5
+        self.pb_relax_alpha = pb_relax_alpha
+        self.scale = dim_head ** -0.5 * (pb_relax_alpha ** -1)
         self.heads = heads
         inner_dim = dim_head * heads
 
@@ -1254,7 +1274,10 @@ class CrossAttention(nn.Module):
             mask = rearrange(mask, 'b j -> b 1 1 j')
             sim = sim.masked_fill(~mask, max_neg_value)
 
-        attn = sim.softmax(dim = -1, dtype = torch.float32)
+        sim = sim - sim.amax(dim = -1, keepdim = True).detach()
+        sim = sim * self.pb_relax_alpha
+
+        attn = sim.softmax(dim = -1)
 
         out = einsum('b h i j, b h j d -> b h i d', attn, v)
         out = rearrange(out, 'b h n d -> b n (h d)')
@@ -1341,6 +1364,7 @@ class Unet(nn.Module):
         dim_mults=(1, 2, 4, 8),
         channels = 3,
         channels_out = None,
+        self_attn = False,
         attn_dim_head = 32,
         attn_heads = 16,
         lowres_cond = False, # for cascading diffusion - https://cascaded-diffusion.github.io/
@@ -1387,6 +1411,8 @@ class Unet(nn.Module):
         dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
         in_out = list(zip(dims[:-1], dims[1:]))
 
+        num_stages = len(in_out)
+
         # time, image embeddings, and optional text encoding
 
         cond_dim = default(cond_dim, dim)
@@ -1450,14 +1476,16 @@ class Unet(nn.Module):
 
         attn_kwargs = dict(heads = attn_heads, dim_head = attn_dim_head)
 
+        self_attn = cast_tuple(self_attn, num_stages)
+
+        create_self_attn = lambda dim: EinopsToAndFrom('b c h w', 'b (h w) c', Residual(Attention(dim, **attn_kwargs)))
+
         # resnet block klass
 
-        resnet_groups = cast_tuple(resnet_groups, len(in_out))
+        resnet_groups = cast_tuple(resnet_groups, num_stages)
         top_level_resnet_group = first(resnet_groups)
 
-        num_resnet_blocks = cast_tuple(num_resnet_blocks, len(in_out))
-
-        assert len(resnet_groups) == len(in_out)
+        num_resnet_blocks = cast_tuple(num_resnet_blocks, num_stages)
 
         # downsample klass
 
@@ -1479,9 +1507,9 @@ class Unet(nn.Module):
         self.ups = nn.ModuleList([])
         num_resolutions = len(in_out)
 
-        skip_connect_dims = []          # keeping track of skip connection dimensions
+        skip_connect_dims = [] # keeping track of skip connection dimensions
 
-        for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks) in enumerate(zip(in_out, resnet_groups, num_resnet_blocks)):
+        for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks, layer_self_attn) in enumerate(zip(in_out, resnet_groups, num_resnet_blocks, self_attn)):
             is_first = ind == 0
             is_last = ind >= (num_resolutions - 1)
             layer_cond_dim = cond_dim if not is_first else None
@@ -1489,30 +1517,42 @@ class Unet(nn.Module):
             dim_layer = dim_out if memory_efficient else dim_in
             skip_connect_dims.append(dim_layer)
 
+            attention = nn.Identity()
+            if layer_self_attn:
+                attention = create_self_attn(dim_layer)
+            elif sparse_attn:
+                attention = Residual(LinearAttention(dim_layer, **attn_kwargs))
+
             self.downs.append(nn.ModuleList([
                 downsample_klass(dim_in, dim_out = dim_out) if memory_efficient else None,
                 ResnetBlock(dim_layer, dim_layer, time_cond_dim = time_cond_dim, groups = groups),
-                Residual(LinearAttention(dim_layer, **attn_kwargs)) if sparse_attn else nn.Identity(),
                 nn.ModuleList([ResnetBlock(dim_layer, dim_layer, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups) for _ in range(layer_num_resnet_blocks)]),
+                attention,
                 downsample_klass(dim_layer, dim_out = dim_out) if not is_last and not memory_efficient else nn.Conv2d(dim_layer, dim_out, 1)
             ]))
 
         mid_dim = dims[-1]
 
         self.mid_block1 = ResnetBlock(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
-        self.mid_attn = EinopsToAndFrom('b c h w', 'b (h w) c', Residual(Attention(mid_dim, **attn_kwargs))) if attend_at_middle else None
+        self.mid_attn = create_self_attn(mid_dim)
         self.mid_block2 = ResnetBlock(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
 
-        for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks) in enumerate(zip(reversed(in_out), reversed(resnet_groups), reversed(num_resnet_blocks))):
+        for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks, layer_self_attn) in enumerate(zip(reversed(in_out), reversed(resnet_groups), reversed(num_resnet_blocks), reversed(self_attn))):
             is_last = ind >= (len(in_out) - 1)
             layer_cond_dim = cond_dim if not is_last else None
 
             skip_connect_dim = skip_connect_dims.pop()
 
+            attention = nn.Identity()
+            if layer_self_attn:
+                attention = create_self_attn(dim_out)
+            elif sparse_attn:
+                attention = Residual(LinearAttention(dim_out, **attn_kwargs))
+
             self.ups.append(nn.ModuleList([
                 ResnetBlock(dim_out + skip_connect_dim, dim_out, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups),
-                Residual(LinearAttention(dim_out, **attn_kwargs)) if sparse_attn else nn.Identity(),
                 nn.ModuleList([ResnetBlock(dim_out + skip_connect_dim, dim_out, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups)  for _ in range(layer_num_resnet_blocks)]),
+                attention,
                 upsample_klass(dim_out, dim_in) if not is_last or memory_efficient else nn.Identity()
             ]))
 
@@ -1591,6 +1631,7 @@ class Unet(nn.Module):
 
         # time conditioning
 
+        time = time.type_as(x)
         time_hiddens = self.to_time_hiddens(time)
 
         time_tokens = self.to_time_tokens(time_hiddens)
@@ -1690,18 +1731,19 @@ class Unet(nn.Module):
 
         hiddens = []
 
-        for pre_downsample, init_block, sparse_attn, resnet_blocks, post_downsample in self.downs:
+        for pre_downsample, init_block, resnet_blocks, attn, post_downsample in self.downs:
             if exists(pre_downsample):
                 x = pre_downsample(x)
 
             x = init_block(x, t, c)
-            x = sparse_attn(x)
-            hiddens.append(x)
 
             for resnet_block in resnet_blocks:
                 x = resnet_block(x, t, c)
                 hiddens.append(x)
 
+            x = attn(x)
+            hiddens.append(x)
+
             if exists(post_downsample):
                 x = post_downsample(x)
 
@@ -1714,15 +1756,15 @@ class Unet(nn.Module):
 
         connect_skip = lambda fmap: torch.cat((fmap, hiddens.pop() * self.skip_connect_scale), dim = 1)
 
-        for init_block, sparse_attn, resnet_blocks, upsample in self.ups:
+        for init_block, resnet_blocks, attn, upsample in self.ups:
             x = connect_skip(x)
             x = init_block(x, t, c)
-            x = sparse_attn(x)
 
             for resnet_block in resnet_blocks:
                 x = connect_skip(x)
                 x = resnet_block(x, t, c)
 
+            x = attn(x)
             x = upsample(x)
 
         x = torch.cat((x, r), dim = 1)

dalle2_pytorch/train_configs.py

@@ -216,6 +216,7 @@ class UnetConfig(BaseModel):
     cond_on_text_encodings: bool = None
     cond_dim: int = None
     channels: int = 3
+    self_attn: ListOrTuple(int)
     attn_dim_head: int = 32
     attn_heads: int = 16
 
@@ -292,6 +293,7 @@ class DecoderTrainConfig(BaseModel):
     epochs: int = 20
     lr: SingularOrIterable(float) = 1e-4
     wd: SingularOrIterable(float) = 0.01
+    warmup_steps: Optional[SingularOrIterable(int)] = None
     find_unused_parameters: bool = True
     max_grad_norm: SingularOrIterable(float) = 0.5
     save_every_n_samples: int = 100000
@@ -345,17 +347,17 @@ class TrainDecoderConfig(BaseModel):
         img_emb_url = data_config.img_embeddings_url
         text_emb_url = data_config.text_embeddings_url
 
-        if using_text_encodings:
+        if using_text_embeddings:
             # Then we need some way to get the embeddings
             assert using_clip or exists(text_emb_url), 'If text conditioning, either clip or text_embeddings_url must be provided'
 
         if using_clip:
-            if using_text_encodings:
+            if using_text_embeddings:
                 assert not exists(text_emb_url) or not exists(img_emb_url), 'Loaded clip, but also provided text_embeddings_url and img_embeddings_url. This is redundant. Remove the clip model or the text embeddings'
             else:
                 assert not exists(img_emb_url), 'Loaded clip, but also provided img_embeddings_url. This is redundant. Remove the clip model or the embeddings'
 
         if text_emb_url:
-            assert using_text_encodings, "Text embeddings are being loaded, but text embeddings are not being conditioned on. This will slow down the dataloader for no reason."
+            assert using_text_embeddings, "Text embeddings are being loaded, but text embeddings are not being conditioned on. This will slow down the dataloader for no reason."
 
         return values

dalle2_pytorch/trainer.py

@@ -3,10 +3,13 @@ import copy
 from pathlib import Path
 from math import ceil
 from functools import partial, wraps
+from contextlib import nullcontext
 from collections.abc import Iterable
 
 import torch
+import torch.nn.functional as F
 from torch import nn
+from torch.optim.lr_scheduler import LambdaLR
 from torch.cuda.amp import autocast, GradScaler
 
 from dalle2_pytorch.dalle2_pytorch import Decoder, DiffusionPrior
@@ -14,6 +17,8 @@ from dalle2_pytorch.optimizer import get_optimizer
 from dalle2_pytorch.version import __version__
 from packaging import version
 
+import pytorch_warmup as warmup
+
 from ema_pytorch import EMA
 
 from accelerate import Accelerator
@@ -162,19 +167,32 @@ class DiffusionPriorTrainer(nn.Module):
         group_wd_params = True,
         device = None,
         accelerator = None,
+        verbose = True,
         **kwargs
     ):
         super().__init__()
         assert isinstance(diffusion_prior, DiffusionPrior)
         assert not exists(accelerator) or isinstance(accelerator, Accelerator)
-        assert exists(accelerator) or exists(device), "You must supply some method of obtaining a device."
         ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
 
+        # verbosity
+
+        self.verbose = verbose
+
         # assign some helpful member vars
+
         self.accelerator = accelerator
-        self.device = accelerator.device if exists(accelerator) else device
         self.text_conditioned = diffusion_prior.condition_on_text_encodings
 
+        # setting the device
+
+        if not exists(accelerator) and not exists(device):
+            diffusion_prior_device = next(diffusion_prior.parameters()).device
+            self.print(f'accelerator not given, and device not specified: defaulting to device of diffusion prior parameters - {diffusion_prior_device}')
+            self.device = diffusion_prior_device
+        else:
+            self.device = accelerator.device if exists(accelerator) else device
+
         # save model
 
         self.diffusion_prior = diffusion_prior
@@ -210,11 +228,14 @@ class DiffusionPriorTrainer(nn.Module):
 
         # track steps internally
 
-        self.register_buffer('step', torch.tensor([0]))
+        self.register_buffer('step', torch.tensor([0], device = self.device))
 
     # accelerator wrappers
 
     def print(self, msg):
+        if not self.verbose:
+            return
+
         if exists(self.accelerator):
             self.accelerator.print(msg)
         else:
@@ -428,6 +449,7 @@ class DecoderTrainer(nn.Module):
         lr = 1e-4,
         wd = 1e-2,
         eps = 1e-8,
+        warmup_steps = None,
         max_grad_norm = 0.5,
         amp = False,
         group_wd_params = True,
@@ -449,13 +471,15 @@ class DecoderTrainer(nn.Module):
         # be able to finely customize learning rate, weight decay
         # per unet
 
-        lr, wd, eps = map(partial(cast_tuple, length = self.num_unets), (lr, wd, eps))
+        lr, wd, eps, warmup_steps = map(partial(cast_tuple, length = self.num_unets), (lr, wd, eps, warmup_steps))
 
         assert all([unet_lr < 1e-3 for unet_lr in lr]), 'your learning rate is too high, recommend sticking with 1e-4, at most 5e-4'
 
         optimizers = []
+        schedulers = []
+        warmup_schedulers = []
 
-        for unet, unet_lr, unet_wd, unet_eps in zip(decoder.unets, lr, wd, eps):
+        for unet, unet_lr, unet_wd, unet_eps, unet_warmup_steps in zip(decoder.unets, lr, wd, eps, warmup_steps):
             optimizer = get_optimizer(
                 unet.parameters(),
                 lr = unet_lr,
@@ -467,6 +491,13 @@ class DecoderTrainer(nn.Module):
 
             optimizers.append(optimizer)
 
+            scheduler = LambdaLR(optimizer, lr_lambda = lambda step: 1.0)
+
+            warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period = unet_warmup_steps) if exists(unet_warmup_steps) else None
+            warmup_schedulers.append(warmup_scheduler)
+
+            schedulers.append(scheduler)
+
             if self.use_ema:
                 self.ema_unets.append(EMA(unet, **ema_kwargs))
 
@@ -474,15 +505,27 @@ class DecoderTrainer(nn.Module):
 
         self.max_grad_norm = max_grad_norm
 
-        self.register_buffer('step', torch.tensor([0.]))
+        self.register_buffer('steps', torch.tensor([0] * self.num_unets))
 
         decoder, *optimizers = list(self.accelerator.prepare(decoder, *optimizers))
+        schedulers = list(self.accelerator.prepare(*schedulers))
 
         self.decoder = decoder
 
+        # store optimizers
+
         for opt_ind, optimizer in zip(range(len(optimizers)), optimizers):
             setattr(self, f'optim{opt_ind}', optimizer)
 
+        # store schedulers
+
+        for sched_ind, scheduler in zip(range(len(schedulers)), schedulers):
+            setattr(self, f'sched{sched_ind}', scheduler)
+
+        # store warmup schedulers
+
+        self.warmup_schedulers = warmup_schedulers
+
     def save(self, path, overwrite = True, **kwargs):
         path = Path(path)
         assert not (path.exists() and not overwrite)
@@ -491,7 +534,7 @@ class DecoderTrainer(nn.Module):
         save_obj = dict(
             model = self.accelerator.unwrap_model(self.decoder).state_dict(),
             version = __version__,
-            step = self.step.item(),
+            steps = self.steps.cpu(),
             **kwargs
         )
 
@@ -510,17 +553,22 @@ class DecoderTrainer(nn.Module):
             self.accelerator.print(f'loading saved decoder at version {loaded_obj["version"]}, but current package version is {__version__}')
 
         self.accelerator.unwrap_model(self.decoder).load_state_dict(loaded_obj['model'], strict = strict)
-        self.step.copy_(torch.ones_like(self.step) * loaded_obj['step'])
+        self.steps.copy_(loaded_obj['steps'])
 
         if only_model:
             return loaded_obj
 
-        for ind in range(0, self.num_unets):
+        for ind, last_step in zip(range(0, self.num_unets), self.steps.tolist()):
+
             optimizer_key = f'optim{ind}'
             optimizer = getattr(self, optimizer_key)
+            warmup_scheduler = self.warmup_schedulers[ind]
 
             self.accelerator.unwrap_model(optimizer).load_state_dict(loaded_obj[optimizer_key])
 
+            if exists(warmup_scheduler):
+                warmup_scheduler.last_step = last_step
+
         if self.use_ema:
             assert 'ema' in loaded_obj
             self.ema_unets.load_state_dict(loaded_obj['ema'], strict = strict)
@@ -539,6 +587,12 @@ class DecoderTrainer(nn.Module):
     def unets(self):
         return nn.ModuleList([ema.ema_model for ema in self.ema_unets])
 
+    def increment_step(self, unet_number):
+        assert 1 <= unet_number <= self.num_unets
+
+        unet_index_tensor = torch.tensor(unet_number - 1, device = self.steps.device)
+        self.steps += F.one_hot(unet_index_tensor, num_classes = len(self.steps))
+
     def update(self, unet_number = None):
         if self.num_unets == 1:
             unet_number = default(unet_number, 1)
@@ -547,17 +601,25 @@ class DecoderTrainer(nn.Module):
         index = unet_number - 1
 
         optimizer = getattr(self, f'optim{index}')
+        scheduler = getattr(self, f'sched{index}')
 
         if exists(self.max_grad_norm):
             self.accelerator.clip_grad_norm_(self.decoder.parameters(), self.max_grad_norm)  # Automatically unscales gradients
+
         optimizer.step()
         optimizer.zero_grad()
 
+        warmup_scheduler = self.warmup_schedulers[index]
+        scheduler_context = warmup_scheduler.dampening if exists(warmup_scheduler) else nullcontext
+
+        with scheduler_context():
+            scheduler.step()
+
         if self.use_ema:
             ema_unet = self.ema_unets[index]
             ema_unet.update()
 
-        self.step += 1
+        self.increment_step(unet_number)
 
     @torch.no_grad()
     @cast_torch_tensor
@@ -607,7 +669,6 @@ class DecoderTrainer(nn.Module):
         total_loss = 0.
 
         for chunk_size_frac, (chunked_args, chunked_kwargs) in split_args_and_kwargs(*args, split_size = max_batch_size, **kwargs):
-            # with autocast(enabled = self.amp):
             with self.accelerator.autocast():
                 loss = self.decoder(*chunked_args, unet_number = unet_number, **chunked_kwargs)
                 loss = loss * chunk_size_frac

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.15.2'
+__version__ = '0.16.12'

setup.py

@@ -37,6 +37,7 @@ setup(
     'packaging',
     'pillow',
     'pydantic',
+    'pytorch-warmup',
     'resize-right>=0.0.2',
     'rotary-embedding-torch',
     'torch>=1.10',