address first two points #78

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.9.7',
+  version = '1.9.9',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/lookup_free_quantization.py

@@ -11,6 +11,7 @@
 
 import torch
 from torch import nn, Tensor
+import torch.nn.functional as F
 from torch.nn import Module
 
 from einops import rearrange, reduce, pack, unpack
@@ -53,8 +54,9 @@ def __init__(
         dim = None,
         codebook_size = None,
         entropy_loss_weight = 0.1,
+        commitment_loss_weight = 1.,
         diversity_gamma = 2.5,
-        straight_through_activation = nn.Tanh(),
+        straight_through_activation = nn.Identity(),
         num_codebooks = 1,
         keep_num_codebooks_dim = None
     ):
@@ -91,6 +93,10 @@ def __init__(
         self.diversity_gamma = diversity_gamma
         self.entropy_loss_weight = entropy_loss_weight
 
+        # commitment loss
+
+        self.commitment_loss_weight = commitment_loss_weight
+
         # for no auxiliary loss, during inference
 
         self.register_buffer('mask', 2 ** torch.arange(codebook_dim - 1, -1, -1))
@@ -157,6 +163,8 @@ def forward(
 
         # quantize by eq 3.
 
+        original_input = x
+
         ones = torch.ones_like(x)
         quantized = torch.where(x > 0, ones, -ones)
 
@@ -190,7 +198,12 @@ def forward(
             # if not training, just return dummy 0
             entropy_aux_loss = self.zero
 
-        entropy_aux_loss = entropy_aux_loss * self.entropy_loss_weight
+        # commit loss
+
+        if self.training:
+            commit_loss = F.mse_loss(original_input, quantized.detach())
+        else:
+            commit_loss = self.zero
 
         # merge back codebook dim
 
@@ -213,4 +226,8 @@ def forward(
         if not self.keep_num_codebooks_dim:
             indices = rearrange(indices, '... 1 -> ...')
 
+        # complete aux loss
+
+        aux_loss = entropy_aux_loss * self.entropy_loss_weight + commit_loss * self.commitment_loss_weight
+
         return Return(x, indices, entropy_aux_loss)