DecoderTrainer sample method uses the exponentially moving averaged

README.md

@@ -760,7 +760,7 @@ decoder = Decoder(
     unet = (unet1, unet2),
     image_sizes = (128, 256),
     clip = clip,
-    timesteps = 1,
+    timesteps = 1000,
     condition_on_text_encodings = True
 ).cuda()
 
@@ -778,6 +778,12 @@ for unet_number in (1, 2):
     loss.backward()
 
     decoder_trainer.update(unet_number) # update the specific unet as well as its exponential moving average
+
+# after much training
+# you can sample from the exponentially moving averaged unets as so
+
+mock_image_embed = torch.randn(4, 512).cuda()
+images = decoder.sample(mock_image_embed, text = text) # (4, 3, 256, 256)
 ```
 
 ## CLI (wip)

dalle2_pytorch/train.py

@@ -144,6 +144,10 @@ class DecoderTrainer(nn.Module):
 
         self.max_grad_norm = max_grad_norm
 
+    @property
+    def unets(self):
+        return nn.ModuleList([ema.ema_model for ema in self.ema_unets])
+
     def scale(self, loss, *, unet_number):
         assert 1 <= unet_number <= self.num_unets
         index = unet_number - 1
@@ -169,7 +173,26 @@ class DecoderTrainer(nn.Module):
             ema_unet = self.ema_unets[index]
             ema_unet.update()
 
-    def forward(self, x, *, unet_number, **kwargs):
+    @torch.no_grad()
+    def sample(self, *args, **kwargs):
+        if self.use_ema:
+            trainable_unets = self.decoder.unets
+            self.decoder.unets = self.unets                  # swap in exponential moving averaged unets for sampling
+
+        output = self.decoder.sample(*args, **kwargs)
+
+        if self.use_ema:
+            self.decoder.unets = trainable_unets             # restore original training unets
+        return output
+
+    def forward(
+        self,
+        x,
+        *,
+        unet_number,
+        divisor = 1,
+        **kwargs
+    ):
         with autocast(enabled = self.amp):
             loss = self.decoder(x, unet_number = unet_number, **kwargs)
-        return self.scale(loss, unet_number = unet_number)
+        return self.scale(loss / divisor, unet_number = unet_number)

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.0.79',
+  version = '0.0.81',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',