take care of the all gather for contrastive loss

Author: lucidrains

gigagan_pytorch/distributed.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn.functional as F
+from torch.autograd import Function
+import torch.distributed as dist
+
+from einops import rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def pad_dim_to(t, length, dim = 0):
+    pad_length = length - t.shape[dim]
+    zero_pairs = (-dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    return F.pad(t, (*((0, 0) * zero_pairs), 0, pad_length))
+
+# distributed helpers
+
+def all_gather_variable_dim(t, dim = 0, sizes = None):
+    device, world_size = t.device, dist.get_world_size()
+
+    if not exists(sizes):
+        size = torch.tensor(t.shape[dim], device = device, dtype = torch.long)
+        sizes = [torch.empty_like(size, device = device, dtype = torch.long) for i in range(world_size)]
+        dist.all_gather(sizes, size)
+        sizes = torch.stack(sizes)
+
+    max_size = sizes.amax().item()
+    padded_t = pad_dim_to(t, max_size, dim = dim)
+
+    gathered_tensors = [torch.empty(padded_t.shape, device = device, dtype = padded_t.dtype) for i in range(world_size)]
+    dist.all_gather(gathered_tensors, padded_t)
+
+    gathered_tensor = torch.cat(gathered_tensors, dim = dim)
+    seq = torch.arange(max_size, device = device)
+
+    mask = rearrange(seq, 'j -> 1 j') < rearrange(sizes, 'i -> i 1')
+    mask = rearrange(mask, 'i j -> (i j)')
+    seq = torch.arange(mask.shape[-1], device = device)
+    indices = seq[mask]
+
+    gathered_tensor = gathered_tensor.index_select(dim, indices)
+
+    return gathered_tensor, sizes
+
+class AllGather(Function):
+    @staticmethod
+    def forward(ctx, x, dim, sizes):
+        is_dist = dist.is_initialized() and dist.get_world_size() > 1
+        ctx.is_dist = is_dist
+
+        if not is_dist:
+            return x, None
+
+        x, batch_sizes = all_gather_variable_dim(x, dim = dim, sizes = sizes)
+        ctx.batch_sizes = batch_sizes.tolist()
+        ctx.dim = dim
+        return x, batch_sizes
+
+    @staticmethod
+    def backward(ctx, grads, _):
+        if not ctx.is_dist:
+            return grads, None, None
+
+        batch_sizes, rank = ctx.batch_sizes, dist.get_rank()
+        grads_by_rank = grads.split(batch_sizes, dim = ctx.dim)
+        return grads_by_rank[rank], None, None
+
+all_gather = AllGather.apply

gigagan_pytorch/gigagan_pytorch.py

@@ -26,6 +26,7 @@
 from gigagan_pytorch.version import __version__
 from gigagan_pytorch.open_clip import OpenClipAdapter
 from gigagan_pytorch.optimizer import get_optimizer
+from gigagan_pytorch.distributed import all_gather
 
 from tqdm import tqdm
 
@@ -175,8 +176,11 @@ def aux_clip_loss(
 ):
     assert exists(texts) ^ exists(text_embeds)
 
+    images, batch_sizes = all_gather(images, 0, None)
+
     if exists(texts):
         text_embeds, _ = clip.embed_texts(texts)
+        text_embeds, _ = all_gather(text_embeds, 0, batch_sizes)
 
     return clip.contrastive_loss(images = images, text_embeds = text_embeds)
 
@@ -1572,6 +1576,9 @@ def init_(self, m):
     def resize_image_to(self, images, resolution):
         return F.interpolate(images, resolution, mode = self.resize_mode)
 
+    def real_images_to_rgbs(self, images):
+        return [self.resize_image_to(images, resolution) for resolution in self.multiscale_input_resolutions]
+
     @property
     def total_params(self):
         return sum([p.numel() for p in self.parameters()])
@@ -2160,7 +2167,7 @@ def train_discriminator_step(
             real_images = real_images.to(self.device)
             real_images.requires_grad_()
 
-            real_images_rgbs = [self.resize_image_to(real_images, resolution) for resolution in self.unwrapped_D.multiscale_input_resolutions]
+            real_images_rgbs = self.unwrapped_D.real_images_to_rgbs(real_images)
 
             # diff augment real images
 
@@ -2331,8 +2338,11 @@ def train_discriminator_step(
                         calc_aux_loss = False
                     )
 
+                    real_images_rgbs = self.D.real_images_to_rgbs(real_images)
+
                     real_logits, *_ = self.D(
                         real_images,
+                        real_images_rgbs,
                         texts = texts,
                         return_multiscale_outputs = False,
                         calc_aux_loss = False

gigagan_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.2.16'
+__version__ = '0.2.17'