Merge dfe1da4d36 into 3265f2edfb

train_network.py

@@ -12,6 +12,8 @@ import json
 from multiprocessing import Value
 import numpy as np
 
+import ast
+
 from tqdm import tqdm
 
 import torch
@@ -1459,8 +1461,9 @@ class NetworkTrainer:
                     optimizer.zero_grad(set_to_none=True)
 
                 if args.scale_weight_norms:
+                    scale_map = args.scale_weight_norms_map if args.scale_weight_norms_map else {}
                     keys_scaled, mean_norm, maximum_norm = accelerator.unwrap_model(network).apply_max_norm_regularization(
-                        args.scale_weight_norms, accelerator.device
+                        args.scale_weight_norms, accelerator.device, scale_map=scale_map
                     )
                     mean_grad_norm = None
                     mean_combined_norm = None
@@ -1728,6 +1731,14 @@ class NetworkTrainer:
 
             logger.info("model saved.")
 
+def parse_dict(input_str):
+    """Convert string input into a dictionary."""
+    try:
+        # Use ast.literal_eval to safely evaluate the string as a Python literal (dict)
+        return ast.literal_eval(input_str)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"Invalid dictionary format: {input_str}")
+
 
 def setup_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
@@ -1831,6 +1842,12 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="Scale the weight of each key pair to help prevent overtraing via exploding gradients. (1 is a good starting point) / 重みの値をスケーリングして勾配爆発を防ぐ（1が初期値としては適当）",
     )
+    parser.add_argument(
+        "--scale_weight_norms_map",
+        type=parse_dict,
+        default="{}",
+        help="Scale the weight of each key pair to help prevent overtraing via exploding gradients. (1 is a good starting point) / 重みの値をスケーリングして勾配爆発を防ぐ（1が初期値としては適当）",
+    )
     parser.add_argument(
         "--base_weights",
         type=str,