add reset+ema for vq-vae

.vscode/launch.json

@@ -0,0 +1,16 @@
+{
+    // 使用 IntelliSense 了解相关属性。 
+    // 悬停以查看现有属性的描述。
+    // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: 当前文件",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": false
+        }
+    ]
+}

configs/vq_vae.yaml

@@ -1,21 +1,22 @@
 model_params:
-  name: 'VQVAE'
+  name: "VQVAE"
   in_channels: 3
+  quantizer: "ema_reset"
   embedding_dim: 64
   num_embeddings: 512
   img_size: 64
-  beta: 0.25
+  beta: 1.0
+  mu: 0.99
 
 data_params:
   data_path: "Data/"
   train_batch_size: 64
-  val_batch_size:  64
+  val_batch_size: 64
   patch_size: 64
   num_workers: 4
 
-
 exp_params:
-  LR: 0.005
+  LR: 0.0003
   weight_decay: 0.0
   scheduler_gamma: 0.0
   kld_weight: 0.00025
@@ -27,4 +28,4 @@ trainer_params:
 
 logging_params:
   save_dir: "logs/"
-  name: 'VQVAE'
+  name: "VQVAE"

experiment.py

@@ -66,7 +66,12 @@ def sample_images(self):
         test_input, test_label = next(iter(self.trainer.datamodule.test_dataloader()))
         test_input = test_input.to(self.curr_device)
         test_label = test_label.to(self.curr_device)
-
+        vutils.save_image(test_input.data,
+                          os.path.join(self.logger.log_dir , 
+                                       "Origin", 
+                                       f"origin_{self.logger.name}_Epoch_{self.current_epoch}.png"),
+                          normalize=True,
+                          nrow=12)
 #         test_input, test_label = batch
         recons = self.model.generate(test_input, labels = test_label)
         vutils.save_image(recons.data,

models/__init__.py

@@ -21,6 +21,7 @@
 from .vq_vae import *
 from .betatc_vae import *
 from .dip_vae import *
+from .quantizer import *
 
 
 # Aliases