Merge pull request #61 from theAdamColton/fix_not_save_affine

fix sd loading, don't update affine during eval

Author: lucidrains

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -304,8 +304,10 @@ def __init__(
         self.register_buffer('batch_mean', None)
         self.register_buffer('batch_variance', None)
 
-        self.register_buffer('codebook_mean', None)
-        self.register_buffer('codebook_variance', None)
+        self.register_buffer('codebook_mean_needs_init', torch.Tensor([True]))
+        self.register_buffer('codebook_mean', torch.empty(num_codebooks, 1, dim))
+        self.register_buffer('codebook_variance_needs_init', torch.Tensor([True]))
+        self.register_buffer('codebook_variance', torch.empty(num_codebooks, 1, dim))
 
     @torch.jit.ignore
     def init_embed_(self, data):
@@ -329,8 +331,14 @@ def init_embed_(self, data):
     def update_with_decay(self, buffer_name, new_value, decay):
         old_value = getattr(self, buffer_name)
 
-        if not exists(old_value):
+        needs_init = getattr(self, buffer_name + "_needs_init", False)
+
+        if needs_init:
+            self.register_buffer(buffer_name + "_needs_init", torch.Tensor([False]))
+
+        if not exists(old_value) or needs_init:
             self.register_buffer(buffer_name, new_value.detach())
+
             return
 
         value = old_value * decay + new_value.detach() * (1 - decay)
@@ -344,8 +352,9 @@ def update_affine(self, data, embed):
 
         data, embed = map(lambda t: rearrange(t, 'h ... d -> h (...) d'), (data, embed))
 
-        self.update_with_decay('codebook_mean', reduce(embed, 'h n d -> h 1 d', 'mean'), self.affine_param_codebook_decay)
-        self.update_with_decay('codebook_variance', reduce(embed, 'h n d -> h 1 d', var_fn), self.affine_param_codebook_decay)
+        if self.training:
+            self.update_with_decay('codebook_mean', reduce(embed, 'h n d -> h 1 d', 'mean'), self.affine_param_codebook_decay)
+            self.update_with_decay('codebook_variance', reduce(embed, 'h n d -> h 1 d', var_fn), self.affine_param_codebook_decay)
 
         if not self.sync_affine_param:
             self.update_with_decay('batch_mean', reduce(data, 'h n d -> h 1 d', 'mean'), self.affine_param_batch_decay)