TI: Use grad clipping from LoRA #104

1 files changed, 11 insertions, 8 deletions

diff --git a/train_ti.py b/train_ti.py
index 4e2c3c5..1054a5d 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -7,6 +7,7 @@ from pathlib import Path
 from contextlib import contextmanager, nullcontext
 
 import torch
+import torch.nn.functional as F
 import torch.utils.checkpoint
 
 from accelerate import Accelerator
@@ -22,7 +23,7 @@ from slugify import slugify
 from util import load_config, load_embeddings_from_dir
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import VlpnDataModule, VlpnDataItem
-from training.common import run_model, generate_class_images
+from training.common import loss_step, generate_class_images
 from training.optimization import get_one_cycle_schedule
 from training.lr import LRFinder
 from training.util import AverageMeter, CheckpointerBase, EMAModel, save_args
@@ -165,7 +166,7 @@ def parse_args():
     parser.add_argument(
         "--tag_dropout",
         type=float,
-        default=0,
+        default=0.1,
         help="Tag dropout probability.",
     )
     parser.add_argument(
@@ -866,14 +867,16 @@ def main():
         finally:
             pass
 
+    @torch.no_grad()
     def on_clip():
-        accelerator.clip_grad_norm_(
-            text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
-            args.max_grad_norm
-        )
+        embeddings = text_encoder.text_model.embeddings.temp_token_embedding
+
+        pre_norm = embeddings.weight.norm(dim=-1, keepdim=True)
+        lambda_ = min(1.0, 100 * lr_scheduler.get_last_lr()[0])
+        embeddings.weight[:] = F.normalize(embeddings.weight, dim=-1) * (pre_norm + lambda_ * (0.4 - pre_norm))
 
     loop = partial(
-        run_model,
+        loss_step,
         vae,
         noise_scheduler,
         unet,
@@ -971,7 +974,7 @@ def main():
 
     try:
         for epoch in range(num_epochs):
-            if accelerator.is_main_process:
+            if accelerator.is_main_process and False:
                 if epoch % args.sample_frequency == 0:
                     checkpointer.save_samples(global_step + global_step_offset, args.sample_steps)