use default epoch_start_ema=20 and update_ema_freq=10

This commit is contained in:
allenzren 2024-11-07 10:55:16 -05:00
parent c0921a1fb5
commit 0bdae945e9
82 changed files with 11 additions and 82 deletions

View File

@ -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

View File

@ -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

View File

@ -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)

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 15000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 15000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 15000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -34,7 +34,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -33,7 +33,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 3000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 100
model:

View File

@ -31,7 +31,7 @@ train:
first_cycle_steps: 1000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 100
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 3000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 100
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 1000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 100
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 8000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 8000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 8000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 3000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 100
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 3000
warmup_steps: 1
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 100
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -42,7 +42,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -41,7 +41,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -42,7 +42,6 @@ train:
first_cycle_steps: 8000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -41,7 +41,6 @@ train:
first_cycle_steps: 3000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -42,7 +42,6 @@ train:
first_cycle_steps: 8000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-4
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -41,7 +41,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -42,7 +42,6 @@ train:
first_cycle_steps: 8000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 10000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -41,7 +41,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -31,7 +31,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model:

View File

@ -32,7 +32,6 @@ train:
first_cycle_steps: 5000
warmup_steps: 100
min_lr: 1e-5
epoch_start_ema: 20
save_model_freq: 500
model: