diff --git a/library/train_util.py b/library/train_util.py index 7d311827..e68444a0 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -1927,7 +1927,7 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: parser.add_argument( "--prior_loss_weight", type=float, default=1.0, help="loss weight for regularization images / 正則化画像のlossの重み" ) - + parser.add_argument("--min_snr_gamma", type=float, default=0, help="gamma for reducing the weight of high loss timesteps. Lower numbers have stronger effect. 5 is recommended by paper.") def verify_training_args(args: argparse.Namespace): if args.v_parameterization and not args.v2: diff --git a/train_network.py b/train_network.py index 7f910df4..5cb08f15 100644 --- a/train_network.py +++ b/train_network.py @@ -548,6 +548,16 @@ def train(args): loss_weights = batch["loss_weights"] # 各sampleごとのweight loss = loss * loss_weights + gamma = args.min_snr_gamma + if gamma: + sigma = torch.sub(noisy_latents, latents) #find noise as applied + zeros = torch.zeros_like(sigma) + alpha_mean_sq = torch.nn.functional.mse_loss(latents.float(), zeros.float(), reduction="none").mean([1, 2, 3]) #trick to get Mean Square + sigma_mean_sq = torch.nn.functional.mse_loss(sigma.float(), zeros.float(), reduction="none").mean([1, 2, 3]) #trick to get Mean Square + snr = torch.div(alpha_mean_sq,sigma_mean_sq) #Signal to Noise Ratio = ratio of Mean Squares + gamma_over_snr = torch.div(torch.ones_like(snr)*gamma,snr) + snr_weight = torch.minimum(gamma_over_snr,torch.ones_like(gamma_over_snr)).float() #from paper + loss = loss * snr_weight loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし