Merge pull request #235 from xju2/fix-3d

add an option for 3d images

Author: lucidrains

tests/test_readme.py

@@ -387,7 +387,7 @@ def test_latent_q():
     quantizer = LatentQuantize(
         levels = [5, 5, 8],      # number of levels per codebook dimension
         dim = 16,                   # input dim
-        commitment_loss_weight=0.1,  
+        commitment_loss_weight=0.1,
         quantization_loss_weight=0.1,
     )
 
@@ -488,3 +488,21 @@ def test_accum_ema_update():
     _ = vq(x)
 
     assert not torch.allclose(codebook_before, vq.codebook, atol = 1e-6)
+
+def test_vq_3d():
+    from vector_quantize_pytorch import VectorQuantize
+
+    quantizer = VectorQuantize(
+        dim = 64,
+        codebook_size = 512,                # codebook size
+        accept_3d_fmap=True,
+    )
+
+    x = torch.randn(1, 64, 16, 16, 16)  # (B, C, D, H, W)
+    quantizer.eval()
+
+    quantized, indices, commit_loss = quantizer(x)
+
+    assert x.shape == quantized.shape
+    assert indices.shape == (1, 16, 16, 16)
+    assert torch.allclose(quantized, quantizer.get_output_from_indices(indices))

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -770,6 +770,7 @@ def __init__(
         threshold_ema_dead_code = 0,
         channel_last = True,
         accept_image_fmap = False,
+        accept_3d_fmap = False,
         commitment_weight = 1.,
         commitment_use_cross_entropy_loss = False,
         orthogonal_reg_weight = 0.,
@@ -906,6 +907,7 @@ def __init__(
         self.codebook_size = codebook_size
 
         self.accept_image_fmap = accept_image_fmap
+        self.accept_3d_fmap = accept_3d_fmap
         self.channel_last = channel_last
 
         self.register_buffer('zero', tensor(0.), persistent = False)
@@ -953,7 +955,7 @@ def get_codes_from_indices(self, indices):
             codes = rearrange(codes, 'b h n d -> b n (h d)')
             codes = unpack_one(codes, 'b * d')
 
-        if not self.channel_last:
+        if not self.channel_last or self.accept_image_fmap or self.accept_3d_fmap:
             codes = rearrange(codes, 'b ... d -> b d ...')
 
         return codes
@@ -1000,7 +1002,11 @@ def update_ema_indices(self, x, indices, mask = None):
             height, width = x.shape[-2:]
             x = rearrange(x, 'b c h w -> b (h w) c')
 
-        if not self.channel_last and not self.accept_image_fmap:
+        if self.accept_3d_fmap:
+            assert not exists(mask)
+            x = rearrange(x, 'b c d h w -> b (d h w) c')
+
+        if not self.channel_last and not self.accept_image_fmap and not self.accept_3d_fmap:
             x = rearrange(x, 'b d n -> b n d')
 
         x = self.project_in(x)
@@ -1016,6 +1022,9 @@ def update_ema_indices(self, x, indices, mask = None):
         if self.accept_image_fmap:
              indices = rearrange(indices, 'b h w ... -> b (h w) ...')
 
+        if self.accept_3d_fmap:
+             indices = rearrange(indices, 'b d h w ... -> b (d h w) ...')
+
         if x.ndim == 2: # only one token
              indices = rearrange(indices, 'b ... -> b 1 ...')
 
@@ -1059,7 +1068,7 @@ def forward(
 
         shape, dtype, device, heads, is_multiheaded, codebook_size, return_loss = x.shape, x.dtype, x.device, self.heads, self.heads > 1, self.codebook_size, exists(indices)
 
-        need_transpose = not self.channel_last and not self.accept_image_fmap
+        need_transpose = not self.channel_last and not self.accept_image_fmap and not self.accept_3d_fmap
         should_inplace_optimize = exists(self.in_place_codebook_optimizer)
 
         # rearrange inputs
@@ -1069,6 +1078,11 @@ def forward(
             height, width = x.shape[-2:]
             x = rearrange(x, 'b c h w -> b (h w) c')
 
+        if self.accept_3d_fmap:
+            assert not exists(mask)
+            depth, height, width = x.shape[-3:]
+            x = rearrange(x, 'b c d h w -> b (d h w) c')
+
         if need_transpose:
             x = rearrange(x, 'b d n -> b n d')
 
@@ -1196,6 +1210,9 @@ def calculate_ce_loss(codes):
         if self.accept_image_fmap:
             embed_ind = rearrange(embed_ind, 'b (h w) ... -> b h w ...', h = height, w = width)
 
+        if self.accept_3d_fmap:
+            embed_ind = rearrange(embed_ind, 'b (d h w) ... -> b d h w ...', d = depth, h = height, w = width)
+
         if only_one:
             embed_ind = rearrange(embed_ind, 'b 1 ... -> b ...')
 
@@ -1289,6 +1306,9 @@ def calculate_ce_loss(codes):
         if self.accept_image_fmap:
             quantize = rearrange(quantize, 'b (h w) c -> b c h w', h = height, w = width)
 
+        if self.accept_3d_fmap:
+            quantize = rearrange(quantize, "b (d h w) c -> b c d h w", d=depth, h=height, w=width)
+
         if only_one:
             quantize = rearrange(quantize, 'b 1 d -> b d')