first make contrived unet work with modality only training

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.4.12"
+version = "0.4.14"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

train_image_only_with_unet.py

@@ -0,0 +1,112 @@
+from shutil import rmtree
+from pathlib import Path
+
+import torch
+from torch import tensor, nn
+from torch.nn import Module
+from torch.utils.data import Dataset, DataLoader
+from torch.optim import Adam
+
+from einops import rearrange
+
+import torchvision
+import torchvision.transforms as T
+from torchvision.utils import save_image
+
+from transfusion_pytorch import Transfusion, print_modality_sample
+
+rmtree('./results', ignore_errors = True)
+results_folder = Path('./results')
+results_folder.mkdir(exist_ok = True, parents = True)
+
+# functions
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+# encoder / decoder
+
+class Encoder(Module):
+    def forward(self, x):
+        x = rearrange(x, '... 1 (h p1) (w p2) -> ... (p1 p2) h w', p1 = 2, p2 = 2)
+        return x * 2 - 1
+
+class Decoder(Module):
+    def forward(self, x):
+        x = rearrange(x, '... (p1 p2) h w -> ... 1 (h p1) (w p2)', p1 = 2, p2 = 2, h = 14)
+        return ((x + 1) * 0.5).clamp(min = 0., max = 1.)
+
+model = Transfusion(
+    num_text_tokens = 10,
+    dim_latent = 4,
+    channel_first_latent = True,
+    modality_default_shape = (14, 14),
+    modality_encoder = Encoder(),
+    modality_decoder = Decoder(),
+    pre_post_transformer_enc_dec = (
+        nn.Conv2d(4, 64, 3, 2, 1),
+        nn.ConvTranspose2d(64, 4, 3, 2, 1, output_padding = 1),
+    ),
+    add_pos_emb = True,
+    modality_num_dim = 2,
+    velocity_consistency_loss_weight = 0.1,
+    transformer = dict(
+        dim = 64,
+        depth = 4,
+        dim_head = 32,
+        heads = 8
+    )
+).cuda()
+
+ema_model = model.create_ema()
+
+class MnistDataset(Dataset):
+    def __init__(self):
+        self.mnist = torchvision.datasets.MNIST(
+            './data',
+            download = True
+        )
+
+    def __len__(self):
+        return len(self.mnist)
+
+    def __getitem__(self, idx):
+        pil, labels = self.mnist[idx]
+        digit_tensor = T.PILToTensor()(pil)
+        return (digit_tensor / 255).float()
+
+def cycle(iter_dl):
+    while True:
+        for batch in iter_dl:
+            yield batch
+
+dataset = MnistDataset()
+
+dataloader = DataLoader(dataset, batch_size = 32, shuffle = True)
+iter_dl = cycle(dataloader)
+
+optimizer = Adam(model.parameters(), lr = 8e-4)
+
+# train loop
+
+for step in range(1, 100_000 + 1):
+
+    loss = model(next(iter_dl), velocity_consistency_ema_model = ema_model)
+    loss.backward()
+
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+
+    optimizer.step()
+    optimizer.zero_grad()
+
+    ema_model.update()
+
+    print(f'{step}: {loss.item():.3f}')
+
+    if divisible_by(step, 500):
+        image = ema_model.generate_modality_only(batch_size = 64)
+
+        save_image(
+            rearrange(image, '(gh gw) 1 h w -> 1 (gh h) (gw w)', gh = 8).detach().cpu(),
+            str(results_folder / f'{step}.png')
+        )

transfusion_pytorch/transfusion.py

@@ -1050,6 +1050,7 @@ def __init__(
         add_pos_emb: bool | tuple[bool, ...] = False,
         modality_encoder: Module | tuple[Module, ...] | None = None,
         modality_decoder: Module | tuple[Module, ...] | None = None,
+        pre_post_transformer_enc_dec: tuple[Module, Module] | tuple[tuple[Module, Module], ...] | None = None,
         modality_token_transform: tuple[ModalityTokenTransform, ...] | ModalityTokenTransform | None = None,
         modality_default_shape: tuple[int, ...] | tuple[tuple[int, ...], ...] | None = None,
         fallback_to_default_shape_if_invalid = False,
@@ -1188,14 +1189,32 @@ def __init__(
 
         assert len(self.modality_token_transform) == self.num_modalities
 
+        # prepare pre-post transformer encoder / decoder, for the learnable unets as in paper
+
+        if is_bearable(pre_post_transformer_enc_dec, tuple[Module, Module]):
+            pre_post_transformer_enc_dec = (pre_post_transformer_enc_dec,)
+
+        pre_post_transformer_enc_dec = cast_tuple(pre_post_transformer_enc_dec, self.num_modalities)
+        assert len(pre_post_transformer_enc_dec) == self.num_modalities
+
         # latent to model and back
         # by default will be Linear, with or without rearranges depending on channel_first_latent setting
-        # can also be overridden for the unet down/up as in the paper
+        # can also be overridden for the unet down/up as in the paper with `pre_post_transformer_enc_dec: tuple[Module, Module]`
 
         latent_to_model_projs = []
         model_to_latent_projs = []
 
-        for dim_latent, one_channel_first_latent in zip(self.dim_latents, self.channel_first_latent):
+        for (
+            dim_latent,
+            one_channel_first_latent,
+            enc_dec,
+         ) in zip(
+            self.dim_latents,
+            self.channel_first_latent,
+            pre_post_transformer_enc_dec
+        ):
+
+            pre_attend_enc, post_attend_dec = default(enc_dec, (None, None))
 
             latent_to_model_proj = Linear(dim_latent, dim) if dim_latent != dim else nn.Identity()
             model_to_latent_proj = Linear(dim, dim_latent, bias = False)
@@ -1204,8 +1223,8 @@ def __init__(
                 latent_to_model_proj = nn.Sequential(Rearrange('b d ... -> b ... d'), latent_to_model_proj, Rearrange('b ... d -> b d ...'))
                 model_to_latent_proj = nn.Sequential(Rearrange('b d ... -> b ... d'), model_to_latent_proj, Rearrange('b ... d -> b d ...'))
 
-            latent_to_model_projs.append(latent_to_model_proj)
-            model_to_latent_projs.append(model_to_latent_proj)
+            latent_to_model_projs.append(default(pre_attend_enc, latent_to_model_proj))
+            model_to_latent_projs.append(default(post_attend_dec, model_to_latent_proj))
 
         self.latent_to_model_projs = ModuleList(latent_to_model_projs)
         self.model_to_latent_projs = ModuleList(model_to_latent_projs)
@@ -1706,18 +1725,6 @@ def forward_modality(
                 mod.encoder.eval()
                 modalities = self.maybe_add_temp_batch_dim(mod.encoder)(modalities).detach()
 
-        # axial positions
-
-        if mod.add_pos_emb:
-            assert exists(mod.num_dim), f'modality_num_dim must be set for modality {modality_type} if further injecting axial positional embedding'
-
-            if mod.channel_first_latent:
-                _, _, *axial_dims = modalities.shape
-            else:
-                _, *axial_dims, _ = modalities.shape
-
-            assert len(axial_dims) == mod.num_dim, f'received modalities of ndim {len(axial_dims)} but expected {modality_num_dim}'
-
         # shapes and device
 
         tokens = modalities
@@ -1753,6 +1760,15 @@ def forward_modality(
         if mod.channel_first_latent:
             noised_tokens = rearrange(noised_tokens, 'b d ... -> b ... d')
 
+        # axial positions
+
+        if mod.add_pos_emb:
+            assert exists(mod.num_dim), f'modality_num_dim must be set for modality {modality_type} if further injecting axial positional embedding'
+
+            _, *axial_dims, _ = noised_tokens.shape
+
+            assert len(axial_dims) == mod.num_dim, f'received modalities of ndim {len(axial_dims)} but expected {modality_num_dim}'
+
         # maybe transform
 
         noised_tokens = mod.token_transform(noised_tokens)