Use context manager for EMA, on_train/eval hooks

1 files changed, 66 insertions, 54 deletions

diff --git a/train_ti.py b/train_ti.py
index aa2bf02..f622299 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -2,10 +2,9 @@ import argparse
 import math
 import datetime
 import logging
-import copy
-from pathlib import Path
 from functools import partial
-from contextlib import nullcontext
+from pathlib import Path
+from contextlib import contextmanager, nullcontext
 
 import torch
 import torch.utils.checkpoint
@@ -849,11 +848,24 @@ def main():
     num_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
     val_steps = num_val_steps_per_epoch * num_epochs
 
+    @contextmanager
     def on_train():
-        tokenizer.train()
+        try:
+            tokenizer.train()
+            yield
+        finally:
+            tokenizer.eval()
 
+    @contextmanager
     def on_eval():
-        tokenizer.eval()
+        try:
+            ema_context = ema_embeddings.apply_temporary(
+                text_encoder.text_model.embeddings.temp_token_embedding.parameters()) if args.use_ema is not None and eval else nullcontext()
+
+            with ema_context:
+                yield
+        finally:
+            pass
 
     loop = partial(
         run_model,
@@ -961,80 +973,80 @@ def main():
             local_progress_bar.reset()
 
             text_encoder.train()
-            on_train()
 
-            for step, batch in enumerate(train_dataloader):
-                with accelerator.accumulate(text_encoder):
-                    loss, acc, bsz = loop(step, batch)
+            with on_train():
+                for step, batch in enumerate(train_dataloader):
+                    with accelerator.accumulate(text_encoder):
+                        loss, acc, bsz = loop(step, batch)
 
-                    accelerator.backward(loss)
+                        accelerator.backward(loss)
 
-                    optimizer.step()
-                    if not accelerator.optimizer_step_was_skipped:
-                        lr_scheduler.step()
-                    optimizer.zero_grad(set_to_none=True)
+                        optimizer.step()
+                        if not accelerator.optimizer_step_was_skipped:
+                            lr_scheduler.step()
+                        optimizer.zero_grad(set_to_none=True)
 
-                    avg_loss.update(loss.detach_(), bsz)
-                    avg_acc.update(acc.detach_(), bsz)
+                        avg_loss.update(loss.detach_(), bsz)
+                        avg_acc.update(acc.detach_(), bsz)
 
-                # Checks if the accelerator has performed an optimization step behind the scenes
-                if accelerator.sync_gradients:
-                    if args.use_ema:
-                        ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters())
+                    # Checks if the accelerator has performed an optimization step behind the scenes
+                    if accelerator.sync_gradients:
+                        if args.use_ema:
+                            ema_embeddings.step(text_encoder.text_model.embeddings.temp_token_embedding.parameters())
 
-                    local_progress_bar.update(1)
-                    global_progress_bar.update(1)
+                        local_progress_bar.update(1)
+                        global_progress_bar.update(1)
 
-                    global_step += 1
+                        global_step += 1
 
-                logs = {
-                    "train/loss": avg_loss.avg.item(),
-                    "train/acc": avg_acc.avg.item(),
-                    "train/cur_loss": loss.item(),
-                    "train/cur_acc": acc.item(),
-                    "lr": lr_scheduler.get_last_lr()[0],
-                }
-                if args.use_ema:
-                    logs["ema_decay"] = ema_embeddings.decay
+                    logs = {
+                        "train/loss": avg_loss.avg.item(),
+                        "train/acc": avg_acc.avg.item(),
+                        "train/cur_loss": loss.item(),
+                        "train/cur_acc": acc.item(),
+                        "lr": lr_scheduler.get_last_lr()[0],
+                    }
+                    if args.use_ema:
+                        logs["ema_decay"] = ema_embeddings.decay
 
-                accelerator.log(logs, step=global_step)
+                    accelerator.log(logs, step=global_step)
 
-                local_progress_bar.set_postfix(**logs)
+                    local_progress_bar.set_postfix(**logs)
 
-                if global_step >= args.max_train_steps:
-                    break
+                    if global_step >= args.max_train_steps:
+                        break
 
             accelerator.wait_for_everyone()
 
             text_encoder.eval()
-            on_eval()
 
             cur_loss_val = AverageMeter()
             cur_acc_val = AverageMeter()
 
             with torch.inference_mode():
-                for step, batch in enumerate(val_dataloader):
-                    loss, acc, bsz = loop(step, batch, True)
+                with on_eval():
+                    for step, batch in enumerate(val_dataloader):
+                        loss, acc, bsz = loop(step, batch, True)
 
-                    loss = loss.detach_()
-                    acc = acc.detach_()
+                        loss = loss.detach_()
+                        acc = acc.detach_()
 
-                    cur_loss_val.update(loss, bsz)
-                    cur_acc_val.update(acc, bsz)
+                        cur_loss_val.update(loss, bsz)
+                        cur_acc_val.update(acc, bsz)
 
-                    avg_loss_val.update(loss, bsz)
-                    avg_acc_val.update(acc, bsz)
+                        avg_loss_val.update(loss, bsz)
+                        avg_acc_val.update(acc, bsz)
 
-                    local_progress_bar.update(1)
-                    global_progress_bar.update(1)
+                        local_progress_bar.update(1)
+                        global_progress_bar.update(1)
 
-                    logs = {
-                        "val/loss": avg_loss_val.avg.item(),
-                        "val/acc": avg_acc_val.avg.item(),
-                        "val/cur_loss": loss.item(),
-                        "val/cur_acc": acc.item(),
-                    }
-                    local_progress_bar.set_postfix(**logs)
+                        logs = {
+                            "val/loss": avg_loss_val.avg.item(),
+                            "val/acc": avg_acc_val.avg.item(),
+                            "val/cur_loss": loss.item(),
+                            "val/cur_acc": acc.item(),
+                        }
+                        local_progress_bar.set_postfix(**logs)
 
             logs["val/cur_loss"] = cur_loss_val.avg.item()
             logs["val/cur_acc"] = cur_acc_val.avg.item()