readme

Author: pesser

README.md

@@ -15,6 +15,10 @@
 
 
 ### News
+- Included a bugfix for the quantizer. For backward compatibility it is
+  disabled by default (which corresponds to always training with `beta=1.0`).
+  Use `legacy=False` in the quantizer config to enable it.
+  Thanks [richcmwang](https://github.com/richcmwang) and [wcshin-git](https://github.com/wcshin-git)!
 - Our paper received an update: See https://arxiv.org/abs/2012.09841v3 and the corresponding changelog.
 - Added a pretrained, [1.4B transformer model](https://k00.fr/s511rwcv) trained for class-conditional ImageNet synthesis, which obtains state-of-the-art FID scores among autoregressive approaches and outperforms BigGAN.
 - Added pretrained, unconditional models on [FFHQ](https://k00.fr/yndvfu95) and [CelebA-HQ](https://k00.fr/2xkmielf).

taming/modules/vqvae/quantize.py

@@ -18,6 +18,10 @@ class VectorQuantizer(nn.Module):
     _____________________________________________
     """
 
+    # NOTE: this class contains a bug regarding beta; see VectorQuantizer2 for
+    # a fix and use legacy=False to apply that fix. VectorQuantizer2 can be
+    # used wherever VectorQuantizer has been used before and is additionally
+    # more efficient.
     def __init__(self, n_e, e_dim, beta):
         super(VectorQuantizer, self).__init__()
         self.n_e = n_e
@@ -211,7 +215,9 @@ class VectorQuantizer2(nn.Module):
     Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly
     avoids costly matrix multiplications and allows for post-hoc remapping of indices.
     """
-    # TODO: check beta fix, maybe include 'legacy' version?
+    # NOTE: due to a bug the beta term was applied to the wrong term. for
+    # backwards compatibility we use the buggy version by default, but you can
+    # specify legacy=False to fix it.
     def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random",
                  sane_index_shape=False, legacy=True):
         super().__init__()
@@ -320,4 +326,4 @@ def get_codebook_entry(self, indices, shape):
             # reshape back to match original input shape
             z_q = z_q.permute(0, 3, 1, 2).contiguous()
 
-        return z_q
+        return z_q