diff options
Diffstat (limited to 'train_ti.py')
-rw-r--r-- | train_ti.py | 19 |
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 | |||
7 | from contextlib import contextmanager, nullcontext | 7 | from contextlib import contextmanager, nullcontext |
8 | 8 | ||
9 | import torch | 9 | import torch |
10 | import torch.nn.functional as F | ||
10 | import torch.utils.checkpoint | 11 | import torch.utils.checkpoint |
11 | 12 | ||
12 | from accelerate import Accelerator | 13 | from accelerate import Accelerator |
@@ -22,7 +23,7 @@ from slugify import slugify | |||
22 | from util import load_config, load_embeddings_from_dir | 23 | from util import load_config, load_embeddings_from_dir |
23 | from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion | 24 | from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion |
24 | from data.csv import VlpnDataModule, VlpnDataItem | 25 | from data.csv import VlpnDataModule, VlpnDataItem |
25 | from training.common import run_model, generate_class_images | 26 | from training.common import loss_step, generate_class_images |
26 | from training.optimization import get_one_cycle_schedule | 27 | from training.optimization import get_one_cycle_schedule |
27 | from training.lr import LRFinder | 28 | from training.lr import LRFinder |
28 | from training.util import AverageMeter, CheckpointerBase, EMAModel, save_args | 29 | from training.util import AverageMeter, CheckpointerBase, EMAModel, save_args |
@@ -165,7 +166,7 @@ def parse_args(): | |||
165 | parser.add_argument( | 166 | parser.add_argument( |
166 | "--tag_dropout", | 167 | "--tag_dropout", |
167 | type=float, | 168 | type=float, |
168 | default=0, | 169 | default=0.1, |
169 | help="Tag dropout probability.", | 170 | help="Tag dropout probability.", |
170 | ) | 171 | ) |
171 | parser.add_argument( | 172 | parser.add_argument( |
@@ -866,14 +867,16 @@ def main(): | |||
866 | finally: | 867 | finally: |
867 | pass | 868 | pass |
868 | 869 | ||
870 | @torch.no_grad() | ||
869 | def on_clip(): | 871 | def on_clip(): |
870 | accelerator.clip_grad_norm_( | 872 | embeddings = text_encoder.text_model.embeddings.temp_token_embedding |
871 | text_encoder.text_model.embeddings.temp_token_embedding.parameters(), | 873 | |
872 | args.max_grad_norm | 874 | pre_norm = embeddings.weight.norm(dim=-1, keepdim=True) |
873 | ) | 875 | lambda_ = min(1.0, 100 * lr_scheduler.get_last_lr()[0]) |
876 | embeddings.weight[:] = F.normalize(embeddings.weight, dim=-1) * (pre_norm + lambda_ * (0.4 - pre_norm)) | ||
874 | 877 | ||
875 | loop = partial( | 878 | loop = partial( |
876 | run_model, | 879 | loss_step, |
877 | vae, | 880 | vae, |
878 | noise_scheduler, | 881 | noise_scheduler, |
879 | unet, | 882 | unet, |
@@ -971,7 +974,7 @@ def main(): | |||
971 | 974 | ||
972 | try: | 975 | try: |
973 | for epoch in range(num_epochs): | 976 | for epoch in range(num_epochs): |
974 | if accelerator.is_main_process: | 977 | if accelerator.is_main_process and False: |
975 | if epoch % args.sample_frequency == 0: | 978 | if epoch % args.sample_frequency == 0: |
976 | checkpointer.save_samples(global_step + global_step_offset, args.sample_steps) | 979 | checkpointer.save_samples(global_step + global_step_offset, args.sample_steps) |
977 | 980 | ||