return loss breakdown for LFQ for logging purposes

Author: lucidrains

setup.py

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

vector_quantize_pytorch/lookup_free_quantization.py

@@ -20,6 +20,8 @@
 
 Return = namedtuple('Return', ['quantized', 'indices', 'entropy_aux_loss'])
 
+LossBreakdown = namedtuple('LossBreakdown', ['per_sample_entropy', 'batch_entropy', 'commitment'])
+
 # helper functions
 
 def exists(v):
@@ -151,7 +153,8 @@ def indices_to_codes(
     def forward(
         self,
         x,
-        inv_temperature = 1.
+        inv_temperature = 1.,
+        return_loss_breakdown = False
     ):
         """
         einstein notation
@@ -214,7 +217,7 @@ def forward(
             entropy_aux_loss = per_sample_entropy - self.diversity_gamma * codebook_entropy
         else:
             # if not training, just return dummy 0
-            entropy_aux_loss = self.zero
+            entropy_aux_loss = per_sample_entropy = codebook_entropy = self.zero
 
         # commit loss
 
@@ -248,4 +251,9 @@ def forward(
 
         aux_loss = entropy_aux_loss * self.entropy_loss_weight + commit_loss * self.commitment_loss_weight
 
-        return Return(x, indices, aux_loss)
+        ret = Return(x, indices, aux_loss)
+
+        if not return_loss_breakdown:
+            return ret
+
+        return ret, LossBreakdown(per_sample_entropy, codebook_entropy, commit_loss)