use default epoch_start_ema=20 and update_ema_freq=10

agent/pretrain/train_agent.py

@@ -83,7 +83,8 @@ class PreTrainAgent:
         # Training params
         self.n_epochs = cfg.train.n_epochs
         self.batch_size = cfg.train.batch_size
-        self.epoch_start_ema = cfg.train.epoch_start_ema
+        self.epoch_start_ema = cfg.train.get("epoch_start_ema", 20)
+        self.update_ema_freq = cfg.train.get("update_ema_freq", 10)
         self.val_freq = cfg.train.get("val_freq", 100)
 
         # Logging, checkpoints

agent/pretrain/train_diffusion_agent.py

@@ -21,6 +21,7 @@ class TrainDiffusionAgent(PreTrainAgent):
 
         timer = Timer()
         self.epoch = 1
+        cnt_batch = 0
         for _ in range(self.n_epochs):
 
             # train
@@ -38,7 +39,9 @@ class TrainDiffusionAgent(PreTrainAgent):
                 self.optimizer.zero_grad()
 
                 # update ema
-                self.step_ema()
+                if cnt_batch % self.update_ema_freq == 0:
+                    self.step_ema()
+                cnt_batch += 1
             loss_train = np.mean(loss_train_epoch)
 
             # validate

agent/pretrain/train_gaussian_agent.py

@@ -24,6 +24,7 @@ class TrainGaussianAgent(PreTrainAgent):
 
         timer = Timer()
         self.epoch = 1
+        cnt_batch = 0
         for _ in range(self.n_epochs):
 
             # train
@@ -46,7 +47,9 @@ class TrainGaussianAgent(PreTrainAgent):
                 self.optimizer.zero_grad()
 
                 # update ema
-                self.step_ema()
+                if cnt_batch % self.update_ema_freq == 0:
+                    self.step_ema()
+                cnt_batch += 1
             loss_train = np.mean(loss_train_epoch)
             ent_train = np.mean(ent_train_epoch)