Delete train_network_opt.py

train_network_opt.py

@@ -1,832 +0,0 @@
-from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
-from torch.optim import Optimizer
-from torch.cuda.amp import autocast
-from torch.nn.parallel import DistributedDataParallel as DDP
-from typing import Optional, Union
-import importlib
-import argparse
-import gc
-import math
-import os
-import random
-import time
-import json
-
-from tqdm import tqdm
-import torch
-from accelerate.utils import set_seed
-import diffusers
-from diffusers import DDPMScheduler
-print("**********************************")
-#先に
-#pip install torch_optimizer
-#が必要
-try:
-  import torch_optimizer as optim
-except:
-  print("torch_optimizerがインストールされていないためAdafactorとAdastand以外の追加optimzierは使えません。\noptimizerの変更をしたい場合先にpip install torch_optimizerでライブラリを追加してください")
-try:
-  import adastand
-except:
-  print("※Adastandが使えません")
-
-from transformers.optimization import Adafactor, AdafactorSchedule
-print("**********************************")
-##### バケット拡張のためのモジュール
-import append_module
-######
-import library.train_util as train_util
-from library.train_util import DreamBoothDataset, FineTuningDataset
-
-
-def collate_fn(examples):
-  return examples[0]
-
-
-def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
-  logs = {"loss/current": current_loss, "loss/average": avr_loss}
-
-  if args.network_train_unet_only:
-    logs["lr/unet"] = lr_scheduler.get_last_lr()[0]
-  elif args.network_train_text_encoder_only:
-    logs["lr/textencoder"] = lr_scheduler.get_last_lr()[0]
-  else:
-    last_lrs = lr_scheduler.get_last_lr()
-    if len(last_lrs) == 2:
-      logs["lr/textencoder"] = float(last_lrs[0])
-      logs["lr/unet"] = float(last_lrs[-1])          # may be same to textencoder
-    else:
-      if len(last_lrs) == 4:
-        logs_names = ["textencoder", "lora_unet_mid_block", "unet_down_blocks", "unet_up_blocks"]
-      elif len(last_lrs) == 8:
-        logs_names = ["textencoder", "unet_midblock"]
-        for i in range(3):
-          logs_names.append(f"unet_down_blocks_{i}")
-          logs_names.append(f"unet_up_blocks_{i+1}")
-      else:
-        logs_names = []
-        for i in range(12):
-          logs_names.append(f"text_model_encoder_layers_{i}_")
-        logs_names.append("unet_midblock")
-        for i in range(3):
-          logs_names.append(f"unet_down_blocks_{i}")
-          logs_names.append(f"unet_up_blocks_{i+1}")
-      
-      for last_lr, logs_name in zip(last_lrs, logs_names):
-        logs[f"lr/{logs_name}"] = float(last_lr)
-
-  return logs
-
-
-# Monkeypatch newer get_scheduler() function overridng current version of diffusers.optimizer.get_scheduler
-# code is taken from https://github.com/huggingface/diffusers diffusers.optimizer, commit d87cc15977b87160c30abaace3894e802ad9e1e6
-# Which is a newer release of diffusers than currently packaged with sd-scripts
-# This code can be removed when newer diffusers version (v0.12.1 or greater) is tested and implemented to sd-scripts
-
-
-def get_scheduler_fix(
-    name: Union[str, SchedulerType],
-    optimizer: Optimizer,
-    num_warmup_steps: Optional[int] = None,
-    num_training_steps: Optional[int] = None,
-    num_cycles: float = 1.,
-    power: float = 1.0,
-):
-  """
-  Unified API to get any scheduler from its name.
-  Args:
-      name (`str` or `SchedulerType`):
-          The name of the scheduler to use.
-      optimizer (`torch.optim.Optimizer`):
-          The optimizer that will be used during training.
-      num_warmup_steps (`int`, *optional*):
-          The number of warmup steps to do. This is not required by all schedulers (hence the argument being
-          optional), the function will raise an error if it's unset and the scheduler type requires it.
-      num_training_steps (`int``, *optional*):
-          The number of training steps to do. This is not required by all schedulers (hence the argument being
-          optional), the function will raise an error if it's unset and the scheduler type requires it.
-      num_cycles (`int`, *optional*):
-          The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
-      power (`float`, *optional*, defaults to 1.0):
-          Power factor. See `POLYNOMIAL` scheduler
-      last_epoch (`int`, *optional*, defaults to -1):
-          The index of the last epoch when resuming training.
-  """
-  name = SchedulerType(name)
-  schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
-  if name == SchedulerType.CONSTANT:
-    return schedule_func(optimizer)
-
-  # All other schedulers require `num_warmup_steps`
-  if num_warmup_steps is None:
-    raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
-
-  if name == SchedulerType.CONSTANT_WITH_WARMUP:
-    return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
-
-  # All other schedulers require `num_training_steps`
-  if num_training_steps is None:
-    raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
-
-  if name == SchedulerType.COSINE:
-      print(f"{name} num_cycles: {num_cycles}")
-      return schedule_func(
-          optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles
-      )
-  if name == SchedulerType.COSINE_WITH_RESTARTS:
-      print(f"{name} num_cycles: {int(num_cycles)}")
-      return schedule_func(
-          optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=int(num_cycles)
-      )
-
-  if name == SchedulerType.POLYNOMIAL:
-    return schedule_func(
-        optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power
-    )
-
-  return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
-
-
-def train(args):
-  session_id = random.randint(0, 2**32)
-  training_started_at = time.time()
-  train_util.verify_training_args(args)
-  train_util.prepare_dataset_args(args, True)
-
-  cache_latents = args.cache_latents
-  use_dreambooth_method = args.in_json is None
-
-  if args.seed is not None:
-    set_seed(args.seed)
-
-  tokenizer = train_util.load_tokenizer(args)
-
-  # データセットを準備する
-  if use_dreambooth_method:
-    if args.min_resolution:
-      args.min_resolution = tuple([int(r) for r in args.min_resolution.split(',')])
-      if len(args.min_resolution) == 1:
-        args.min_resolution = (args.min_resolution[0], args.min_resolution[0])
-
-    print("Use DreamBooth method.")
-    train_dataset = append_module.DreamBoothDataset(args.train_batch_size, args.train_data_dir, args.reg_data_dir,
-                                      tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
-                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso, 
-                                      args.bucket_reso_steps, args.bucket_no_upscale, 
-                                      args.prior_loss_weight, args.flip_aug, args.color_aug, args.face_crop_aug_range, 
-                                      args.random_crop, args.debug_dataset, args.min_resolution, args.area_step)
-  else:
-    print("Train with captions.")
-    train_dataset = FineTuningDataset(args.in_json, args.train_batch_size, args.train_data_dir,
-                                      tokenizer, args.max_token_length, args.shuffle_caption, args.keep_tokens,
-                                      args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
-                                      args.bucket_reso_steps, args.bucket_no_upscale, 
-                                      args.flip_aug, args.color_aug, args.face_crop_aug_range, args.random_crop,
-                                      args.dataset_repeats, args.debug_dataset)
-
-  # 学習データのdropout率を設定する
-  train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
-
-  train_dataset.make_buckets()
-
-  if args.debug_dataset:
-    train_util.debug_dataset(train_dataset)
-    return
-  if len(train_dataset) == 0:
-    print("No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）")
-    return
-
-  # acceleratorを準備する
-  print("prepare accelerator")
-  accelerator, unwrap_model = train_util.prepare_accelerator(args)
-
-  # mixed precisionに対応した型を用意しておき適宜castする
-  weight_dtype, save_dtype = train_util.prepare_dtype(args)
-
-  # モデルを読み込む
-  text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype)
-  # unnecessary, but work on low-ram device
-  text_encoder.to("cuda")
-  unet.to("cuda")
-  # モデルに xformers とか memory efficient attention を組み込む
-  train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
-
-  # 学習を準備する
-  if cache_latents:
-    vae.to(accelerator.device, dtype=weight_dtype)
-    vae.requires_grad_(False)
-    vae.eval()
-    with torch.no_grad():
-      train_dataset.cache_latents(vae)
-    vae.to("cpu")
-    if torch.cuda.is_available():
-      torch.cuda.empty_cache()
-    gc.collect()
-
-  # prepare network
-  print("import network module:", args.network_module)
-  network_module = importlib.import_module(args.network_module)
-
-  net_kwargs = {}
-  if args.network_args is not None:
-    for net_arg in args.network_args:
-      key, value = net_arg.split('=')
-      net_kwargs[key] = value
-
-  # if a new network is added in future, add if ~ then blocks for each network (;'∀')
-  network = network_module.create_network(1.0, args.network_dim, args.network_alpha, vae, text_encoder, unet, **net_kwargs)
-  if network is None:
-    return
-
-  if args.network_weights is not None:
-    print("load network weights from:", args.network_weights)
-    network.load_weights(args.network_weights)
-
-  train_unet = not args.network_train_text_encoder_only
-  train_text_encoder = not args.network_train_unet_only
-  network.apply_to(text_encoder, unet, train_text_encoder, train_unet)
-
-  if args.gradient_checkpointing:
-    unet.enable_gradient_checkpointing()
-    text_encoder.gradient_checkpointing_enable()
-    network.enable_gradient_checkpointing()                   # may have no effect
-
-  # 学習に必要なクラスを準備する
-  print("prepare optimizer, data loader etc.")
-  try:
-    print(f"torch_optimzier version is {optim.__version__}")
-    not_torch_optimizer_flag = False
-  except:
-    not_torch_optimizer_flag = True
-  try:
-    print(f"adastand version is {adastand.__version__()}")
-    not_adasatand_optimzier_flag = False
-  except:
-    not_adasatand_optimzier_flag = True
-
-  # 8-bit Adamを使う
-  if args.optimizer=="Adafactor" or args.optimizer=="Adastand" or args.optimizer=="Adastand_belief":
-    not_torch_optimizer_flag = False
-    if args.optimizer=="Adafactor":
-      not_adasatand_optimzier_flag = False
-  if not_torch_optimizer_flag or not_adasatand_optimzier_flag:
-    print(f"==========================\n必要なライブラリがないため {args.optimizer} の使用ができません。optimizerを AdamW に変更して実行します\n==========================")
-    args.optimizer="AdamW"
-  if args.use_8bit_adam:
-    if not args.optimizer=="AdamW" and not args.optimizer=="Lamb":
-      print(f"\n==========================\n{args.optimizer} は8bitAdamに実装されていないので8bitAdamをオフにします\n==========================\n")
-      args.use_8bit_adam=False
-  if args.use_8bit_adam:
-    try:
-      import bitsandbytes as bnb
-    except ImportError:
-      raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
-    print("use 8-bit Adam optimizer")
-    args.training_comment=f"{args.training_comment} use_8bit_adam=True"
-    if args.optimizer=="Lamb":
-      optimizer_class = bnb.optim.LAMB8bit
-    else:
-      args.optimizer="AdamW"
-      optimizer_class = bnb.optim.AdamW8bit
-  else:
-    print(f"use {args.optimizer}")
-    if args.optimizer=="RAdam":
-      optimizer_class = torch.optim.RAdam
-    elif args.optimizer=="AdaBound":
-      optimizer_class = optim.AdaBound
-    elif args.optimizer=="AdaBelief":
-      optimizer_class = optim.AdaBelief
-    elif args.optimizer=="AdamP":
-      optimizer_class = optim.AdamP
-    elif args.optimizer=="Adafactor":
-      optimizer_class = Adafactor
-    elif args.optimizer=="Adastand":
-      optimizer_class = adastand.Adastand
-    elif args.optimizer=="Adastand_belief":
-      optimizer_class = adastand.Adastand_b
-    elif args.optimizer=="AggMo":
-      optimizer_class = optim.AggMo
-    elif args.optimizer=="Apollo":
-      optimizer_class = optim.Apollo
-    elif args.optimizer=="Lamb":
-      optimizer_class = optim.Lamb
-    elif args.optimizer=="Ranger":
-      optimizer_class = optim.Ranger
-    elif args.optimizer=="RangerVA":
-      optimizer_class = optim.RangerVA
-    elif args.optimizer=="Yogi":
-      optimizer_class = optim.Yogi
-    elif args.optimizer=="Shampoo":
-      optimizer_class = optim.Shampoo
-    elif args.optimizer=="NovoGrad":
-      optimizer_class = optim.NovoGrad
-    elif args.optimizer=="QHAdam":
-      optimizer_class = optim.QHAdam
-    elif args.optimizer=="DiffGrad" or args.optimizer=="Lookahead_DiffGrad":
-      optimizer_class = optim.DiffGrad
-    elif args.optimizer=="MADGRAD":
-      optimizer_class = optim.MADGRAD
-    else:
-      optimizer_class = torch.optim.AdamW
-
-  # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
-  #optimizerデフォ設定
-  if args.optimizer_arg==None:
-    if args.optimizer=="AdaBelief":
-      args.optimizer_arg = ["eps=1e-16","betas=0.9,0.999","weight_decouple=True","rectify=False","fixed_decay=False"]
-    elif args.optimizer=="DiffGrad":
-      args.optimizer_arg = ["eps=1e-16"]
-  optimizer_arg = {}
-  lookahed_arg = {"k": 5, "alpha": 0.5}
-  adafactor_scheduler_arg = {"initial_lr": 0.}
-  int_args = ["k","n_sma_threshold","warmup"]
-  str_args = ["transformer","grad_transformer"]
-  if not args.optimizer_arg==None and len(args.optimizer_arg)>0:
-    for _opt_arg in args.optimizer_arg:
-      key, value = _opt_arg.split("=")
-      if value=="True" or value=="False":
-        optimizer_arg[key]=bool((value=="True"))
-      elif key=="betas" or key=="nus" or key=="eps2" or (key=="eps" and "," in value):
-        _value = value.split(",")
-        optimizer_arg[key] = (float(_value[0]),float(_value[1]))
-        del _value
-      elif key in int_args:
-        if "Lookahead" in args.optimizer:
-          lookahed_arg[key] = int(value)
-        else:
-          optimizer_arg[key] = int(value)
-      elif key in str_args:
-        optimizer_arg[key] = value
-      else:
-        if key=="alpha" and "Lookahead" in args.optimizer:
-          lookahed_arg[key] = int(value)
-        elif key=="initial_lr" and args.optimizer == "Adafactor":
-          adafactor_scheduler_arg[key] = float(value)
-        else:
-          optimizer_arg[key] = float(value)
-    del _opt_arg
-  AdafactorScheduler_Flag = False
-  list_of_init_lr = []
-  if args.optimizer=="Adafactor":
-    if not "relative_step" in optimizer_arg:
-      optimizer_arg["relative_step"] = True
-    if "warmup_init" in optimizer_arg:
-      if optimizer_arg["warmup_init"]==True and optimizer_arg["relative_step"]==False:
-        print("**************\nwarmup_initはrelative_stepがオンである必要があるためrelative_stepをオンにします\n**************")
-        optimizer_arg["relative_step"] = True
-    if optimizer_arg["relative_step"] == True:
-      AdafactorScheduler_Flag = True
-      list_of_init_lr = [0.,0.]
-      if args.text_encoder_lr is not None: list_of_init_lr[0] = float(args.text_encoder_lr)
-      if args.unet_lr is not None: list_of_init_lr[1] = float(args.unet_lr)
-      #if not "initial_lr" in adafactor_scheduler_arg:
-      #  adafactor_scheduler_arg = args.learning_rate
-      args.learning_rate = None
-      args.text_encoder_lr = None
-      args.unet_lr = None
-  print(f"optimizer arg: {optimizer_arg}")
-  print("=-----------------------------------=")
-  if not AdafactorScheduler_Flag: args.split_lora_networks = False
-  if args.split_lora_networks:
-    lora_names = append_module.create_split_names(args.split_lora_networks, args.split_lora_level)
-    append_module.replace_prepare_optimizer_params(network)
-    trainable_params, _list_of_init_lr = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, list_of_init_lr, lora_names)
-  else:
-    trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
-    _list_of_init_lr = []
-  print(f"trainable_params_len: {len(trainable_params)}")
-  if len(_list_of_init_lr)>0:
-    list_of_init_lr = _list_of_init_lr
-    print(f"split loras network is {len(list_of_init_lr)}")
-  if len(list_of_init_lr) > 0:
-    adafactor_scheduler_arg["initial_lr"] = list_of_init_lr
-
-  optimizer = optimizer_class(trainable_params, lr=args.learning_rate, **optimizer_arg)
-
-  if args.optimizer=="Lookahead_DiffGrad" or args.optimizer=="Lookahedad_Adam":
-    optimizer = optim.Lookahead(optimizer, **lookahed_arg)
-    print(f"lookahed_arg: {lookahed_arg}")
-
-  # dataloaderを準備する
-  # DataLoaderのプロセス数：0はメインプロセスになる
-  n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1)      # cpu_count-1 ただし最大で指定された数まで
-  train_dataloader = torch.utils.data.DataLoader(
-      train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
-
-  # 学習ステップ数を計算する
-  if args.max_train_epochs is not None:
-    args.max_train_steps = args.max_train_epochs * len(train_dataloader)
-    print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
-
-  # lr schedulerを用意する
-  # lr_scheduler = diffusers.optimization.get_scheduler(
-  if AdafactorScheduler_Flag:
-    print("===================================\nAdafactorはデフォルトでrelative_stepがオンになっているので lrは自動算出されるためLrScheculerの指定も無効になります\nもし任意のLrやLr_Schedulerを使いたい場合は --optimizer_arg relative_ste=False を指定してください\nまた任意のLrを使う場合は scale_parameter=False も併せて指定するのが推奨です\n===================================")
-    lr_scheduler = append_module.AdafactorSchedule_append(optimizer, **adafactor_scheduler_arg)
-    print(f"AdafactorSchedule initial lrs: {lr_scheduler.get_lr()}")
-    del list_of_init_lr
-  else:
-    lr_scheduler = get_scheduler_fix(
-        args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps,
-        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-        num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
-  
-  #追加機能の設定をコメントに追記して残す
-  args.training_comment=f"{args.training_comment} optimizer: {args.optimizer} / optimizer_arg: {args.optimizer_arg}"
-  if AdafactorScheduler_Flag:
-    args.training_comment=f"{args.training_comment} split_lora_networks: {args.split_lora_networks}"
-  if args.min_resolution:
-    args.training_comment=f"{args.training_comment} min_resolution: {args.min_resolution} area_step: {args.area_step}"
-
-  # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
-  if args.full_fp16:
-    assert args.mixed_precision == "fp16", "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
-    print("enable full fp16 training.")
-    network.to(weight_dtype)
-
-  # acceleratorがなんかよろしくやってくれるらしい
-  if train_unet and train_text_encoder:
-    unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler)
-  elif train_unet:
-    unet, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet,  network, optimizer, train_dataloader, lr_scheduler)
-  elif train_text_encoder:
-    text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        text_encoder, network, optimizer, train_dataloader, lr_scheduler)
-  else:
-    network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        network, optimizer, train_dataloader, lr_scheduler)
-
-  unet.requires_grad_(False)
-  unet.to(accelerator.device, dtype=weight_dtype)
-  text_encoder.requires_grad_(False)
-  text_encoder.to(accelerator.device)
-  if args.gradient_checkpointing:                       # according to TI example in Diffusers, train is required
-    unet.train()
-    text_encoder.train()
-
-    # set top parameter requires_grad = True for gradient checkpointing works
-    if type(text_encoder) == DDP:
-      text_encoder.module.text_model.embeddings.requires_grad_(True)
-    else:
-      text_encoder.text_model.embeddings.requires_grad_(True)
-  else:
-    unet.eval()
-    text_encoder.eval()
-
-  # support DistributedDataParallel
-  if type(text_encoder) == DDP:
-    text_encoder = text_encoder.module
-    unet = unet.module
-    network = network.module
-
-  network.prepare_grad_etc(text_encoder, unet)
-
-  if not cache_latents:
-    vae.requires_grad_(False)
-    vae.eval()
-    vae.to(accelerator.device, dtype=weight_dtype)
-
-  # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
-  if args.full_fp16:
-    train_util.patch_accelerator_for_fp16_training(accelerator)
-
-  # resumeする
-  if args.resume is not None:
-    print(f"resume training from state: {args.resume}")
-    accelerator.load_state(args.resume)
-
-  # epoch数を計算する
-  num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
-  num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
-  if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
-    args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
-
-  # 学習する
-  total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-  print("running training / 学習開始")
-  print(f"  num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
-  print(f"  num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
-  print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
-  print(f"  num epochs / epoch数: {num_train_epochs}")
-  print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
-  print(f"  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
-  print(f"  gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
-  print(f"  total optimization steps / 学習ステップ数: {args.max_train_steps}")
-
-  metadata = {
-      "ss_session_id": session_id,            # random integer indicating which group of epochs the model came from
-      "ss_training_started_at": training_started_at,          # unix timestamp
-      "ss_output_name": args.output_name,
-      "ss_learning_rate": args.learning_rate,
-      "ss_text_encoder_lr": args.text_encoder_lr,
-      "ss_unet_lr": args.unet_lr,
-      "ss_num_train_images": train_dataset.num_train_images,          # includes repeating
-      "ss_num_reg_images": train_dataset.num_reg_images,
-      "ss_num_batches_per_epoch": len(train_dataloader),
-      "ss_num_epochs": num_train_epochs,
-      "ss_batch_size_per_device": args.train_batch_size,
-      "ss_total_batch_size": total_batch_size,
-      "ss_gradient_checkpointing": args.gradient_checkpointing,
-      "ss_gradient_accumulation_steps": args.gradient_accumulation_steps,
-      "ss_max_train_steps": args.max_train_steps,
-      "ss_lr_warmup_steps": args.lr_warmup_steps,
-      "ss_lr_scheduler": args.lr_scheduler,
-      "ss_network_module": args.network_module,
-      "ss_network_dim": args.network_dim,          # None means default because another network than LoRA may have another default dim
-      "ss_network_alpha": args.network_alpha,      # some networks may not use this value
-      "ss_mixed_precision": args.mixed_precision,
-      "ss_full_fp16": bool(args.full_fp16),
-      "ss_v2": bool(args.v2),
-      "ss_resolution": args.resolution,
-      "ss_clip_skip": args.clip_skip,
-      "ss_max_token_length": args.max_token_length,
-      "ss_color_aug": bool(args.color_aug),
-      "ss_flip_aug": bool(args.flip_aug),
-      "ss_random_crop": bool(args.random_crop),
-      "ss_shuffle_caption": bool(args.shuffle_caption),
-      "ss_cache_latents": bool(args.cache_latents),
-      "ss_enable_bucket": bool(train_dataset.enable_bucket),
-      "ss_min_bucket_reso": train_dataset.min_bucket_reso,
-      "ss_max_bucket_reso": train_dataset.max_bucket_reso,
-      "ss_seed": args.seed,
-      "ss_keep_tokens": args.keep_tokens,
-      "ss_noise_offset": args.noise_offset,
-      "ss_dataset_dirs": json.dumps(train_dataset.dataset_dirs_info),
-      "ss_reg_dataset_dirs": json.dumps(train_dataset.reg_dataset_dirs_info),
-      "ss_tag_frequency": json.dumps(train_dataset.tag_frequency),
-      "ss_bucket_info": json.dumps(train_dataset.bucket_info),
-      "ss_training_comment": args.training_comment,       # will not be updated after training
-      "ss_sd_scripts_commit_hash": train_util.get_git_revision_hash()
-  }
-
-  # uncomment if another network is added
-  # for key, value in net_kwargs.items():
-  #   metadata["ss_arg_" + key] = value
-
-  if args.pretrained_model_name_or_path is not None:
-    sd_model_name = args.pretrained_model_name_or_path
-    if os.path.exists(sd_model_name):
-      metadata["ss_sd_model_hash"] = train_util.model_hash(sd_model_name)
-      metadata["ss_new_sd_model_hash"] = train_util.calculate_sha256(sd_model_name)
-      sd_model_name = os.path.basename(sd_model_name)
-    metadata["ss_sd_model_name"] = sd_model_name
-
-  if args.vae is not None:
-    vae_name = args.vae
-    if os.path.exists(vae_name):
-      metadata["ss_vae_hash"] = train_util.model_hash(vae_name)
-      metadata["ss_new_vae_hash"] = train_util.calculate_sha256(vae_name)
-      vae_name = os.path.basename(vae_name)
-    metadata["ss_vae_name"] = vae_name
-
-  metadata = {k: str(v) for k, v in metadata.items()}
-
-  progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
-  global_step = 0
-
-  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
-                                  num_train_timesteps=1000, clip_sample=False)
-
-  if accelerator.is_main_process:
-    accelerator.init_trackers("network_train")
-
-  loss_list = []
-  loss_total = 0.0
-  for epoch in range(num_train_epochs):
-    print(f"epoch {epoch+1}/{num_train_epochs}")
-    train_dataset.set_current_epoch(epoch + 1)
-
-    metadata["ss_epoch"] = str(epoch+1)
-
-    network.on_epoch_start(text_encoder, unet)
-
-    for step, batch in enumerate(train_dataloader):
-      with accelerator.accumulate(network):
-        with torch.no_grad():
-          if "latents" in batch and batch["latents"] is not None:
-            latents = batch["latents"].to(accelerator.device)
-          else:
-            # latentに変換
-            latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
-          latents = latents * 0.18215
-        b_size = latents.shape[0]
-
-        with torch.set_grad_enabled(train_text_encoder):
-          # Get the text embedding for conditioning
-          input_ids = batch["input_ids"].to(accelerator.device)
-          encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizer, text_encoder, weight_dtype)
-
-        # Sample noise that we'll add to the latents
-        noise = torch.randn_like(latents, device=latents.device)
-        if args.noise_offset:
-          # https://www.crosslabs.org//blog/diffusion-with-offset-noise
-          noise += args.noise_offset * torch.randn((latents.shape[0], latents.shape[1], 1, 1), device=latents.device)
-
-        # Sample a random timestep for each image
-        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
-        timesteps = timesteps.long()
-
-        # Add noise to the latents according to the noise magnitude at each timestep
-        # (this is the forward diffusion process)
-        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
-
-        # Predict the noise residual
-        with autocast():
-          noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
-
-        if args.v_parameterization:
-          # v-parameterization training
-          target = noise_scheduler.get_velocity(latents, noise, timesteps)
-        else:
-          target = noise
-
-        loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
-        loss = loss.mean([1, 2, 3])
-
-        loss_weights = batch["loss_weights"]                      # 各sampleごとのweight
-        loss = loss * loss_weights
-
-        loss = loss.mean()                # 平均なのでbatch_sizeで割る必要なし
-
-        accelerator.backward(loss)
-        if accelerator.sync_gradients:
-          params_to_clip = network.get_trainable_params()
-          accelerator.clip_grad_norm_(params_to_clip, 1.0)  # args.max_grad_norm)
-
-        optimizer.step()
-        lr_scheduler.step()
-        optimizer.zero_grad(set_to_none=True)
-
-      # Checks if the accelerator has performed an optimization step behind the scenes
-      if accelerator.sync_gradients:
-        progress_bar.update(1)
-        global_step += 1
-
-      current_loss = loss.detach().item()
-      if epoch == 0:
-        loss_list.append(current_loss)
-      else:
-        loss_total -= loss_list[step]
-        loss_list[step] = current_loss
-      loss_total += current_loss
-      avr_loss = loss_total / len(loss_list)
-      logs = {"loss": avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
-      progress_bar.set_postfix(**logs)
-
-      if args.logging_dir is not None:
-        logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
-        accelerator.log(logs, step=global_step)
-
-      if global_step >= args.max_train_steps:
-        break
-
-    if args.logging_dir is not None:
-      logs = {"loss/epoch": loss_total / len(loss_list)}
-      accelerator.log(logs, step=epoch+1)
-
-    accelerator.wait_for_everyone()
-
-    if args.save_every_n_epochs is not None:
-      model_name = train_util.DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
-
-      def save_func():
-        ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + '.' + args.save_model_as
-        ckpt_file = os.path.join(args.output_dir, ckpt_name)
-        print(f"saving checkpoint: {ckpt_file}")
-        unwrap_model(network).save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
-
-      def remove_old_func(old_epoch_no):
-        old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + '.' + args.save_model_as
-        old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
-        if os.path.exists(old_ckpt_file):
-          print(f"removing old checkpoint: {old_ckpt_file}")
-          os.remove(old_ckpt_file)
-
-      saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
-      if saving and args.save_state:
-        train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
-
-    # end of epoch
-
-  metadata["ss_epoch"] = str(num_train_epochs)
-
-  is_main_process = accelerator.is_main_process
-  if is_main_process:
-    network = unwrap_model(network)
-
-  accelerator.end_training()
-
-  if args.save_state:
-    train_util.save_state_on_train_end(args, accelerator)
-
-  del accelerator                         # この後メモリを使うのでこれは消す
-
-  if is_main_process:
-    os.makedirs(args.output_dir, exist_ok=True)
-
-    model_name = train_util.DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
-    ckpt_name = model_name + '.' + args.save_model_as
-    ckpt_file = os.path.join(args.output_dir, ckpt_name)
-
-    print(f"save trained model to {ckpt_file}")
-    network.save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
-    print("model saved.")
-
-
-if __name__ == '__main__':
-  parser = argparse.ArgumentParser()
-
-  train_util.add_sd_models_arguments(parser)
-  train_util.add_dataset_arguments(parser, True, True, True)
-  train_util.add_training_arguments(parser, True)
-
-  parser.add_argument("--no_metadata", action='store_true', help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
-  parser.add_argument("--save_model_as", type=str, default="safetensors", choices=[None, "ckpt", "pt", "safetensors"],
-                      help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）")
-
-  parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
-  parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
-  parser.add_argument("--lr_scheduler_num_cycles", type=int, default=1,
-                      help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
-  parser.add_argument("--lr_scheduler_power", type=float, default=1,
-                      help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")
-
-  parser.add_argument("--network_weights", type=str, default=None,
-                      help="pretrained weights for network / 学習するネットワークの初期重み")
-  parser.add_argument("--network_module", type=str, default=None, help='network module to train / 学習対象のネットワークのモジュール')
-  parser.add_argument("--network_dim", type=int, default=None,
-                      help='network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）')
-  parser.add_argument("--network_alpha", type=float, default=1,
-                      help='alpha for LoRA weight scaling, default 1 (same as network_dim for same behavior as old version) / LoRaの重み調整のalpha値、デフォルト1（旧バージョンと同じ動作をするにはnetwork_dimと同じ値を指定）')
-  parser.add_argument("--network_args", type=str, default=None, nargs='*',
-                      help='additional argmuments for network (key=value) / ネットワークへの追加の引数')
-  parser.add_argument("--network_train_unet_only", action="store_true", help="only training U-Net part / U-Net関連部分のみ学習する")
-  parser.add_argument("--network_train_text_encoder_only", action="store_true",
-                      help="only training Text Encoder part / Text Encoder関連部分のみ学習する")
-  parser.add_argument("--training_comment", type=str, default=None,
-                      help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列")
-
-  #Optimizer変更関連のオプション追加
-  append_module.add_append_arguments(parser)
-  args = append_module.get_config(parser)
-
-  if args.resolution==args.min_resolution:
-    args.min_resolution=None
-
-  train(args)
-
-  #学習が終わったら現在のargsを保存する
-#  import yaml
-#  import datetime
-#  _t = datetime.datetime.today().strftime('%Y%m%d_%H%M')
-#  if args.output_name==None:
-#    config_name = f"train_network_config_{_t}.yaml"
-#  else:
-#    config_name = f"train_network_config_{os.path.basename(args.output_name)}_{_t}.yaml"
-#  print(f"{config_name} に設定を書き出し中...")
-#  with open(config_name, mode="w") as f:
-#      yaml.dump(args.__dict__, f, indent=4)
-#  print("done!")
-
-'''
-optimizer設定メモ
-(optimizer_argから設定できるように変更するためのメモ)
-
-AdamWのweight_decay初期値は1e-2
-
-RAdam
-weight_decay=1e-2 (? 初期値は0だけど上手くいかないのでweight_decayを設定したほうがいいのかもしれない)
-
-AdaBelief
-eps=1e-16 betas=0.9,0.999 weight_decouple=True rectify=False fiexed_decay=False
-論文中でTransformerの成績が良かった設定
-weight_decay=1e-4 weight_decouple=True rectify=True fixed_decay=False
-
-DiffGrad
-eps=1e-16
-MADGRAD
-eps=1e-6
-
-NovoGrad
-論文ではAdamWと比較して汎化性能が物凄く良くなってるらしい
-再現実装の比較によるとAdamより性能は高いけど収束に必要なstepが多い
-実験paramater: weight_decay=1e-2 amsgrad=True
-
-QHAdam
-収束早い
-
-Adafactor
-transformerベースのT5学習において最強とかいう噂のoptimizer
-huggingfaceのサンプルパラ
-eps=1e-30,1e-3 clip_threshold=1.0 decay_rate=-0.8 relative_step=False scale_parameter=False warmup_init=False
-
-AggMo
-
-Apollo(学習内容によってやたらとVRAM消費が増え続けるので非推奨)
-
-未整頓のパラメータメモ
-Ranger(収束できる設定が見つかる気がしない)
-RangerVA(収束できる設定が見つかる気がしない)
-'''