Add pin_memory to finetune scripts

kohya-ss · Jan 23, 2025 · 03b35be · 03b35be
1 parent 50d8daa
commit 03b35be
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 0 deletions.
diff --git a/finetune/make_captions.py b/finetune/make_captions.py
@@ -188,6 +188,11 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="enable image reading by DataLoader with this number of workers (faster) / DataLoaderによる画像読み込みを有効にしてこのワーカー数を適用する（読み込みを高速化）",
     )
+    parser.add_argument(
+        "--pin_memory",
+        action="store_true",
+        help="Pin memory for faster GPU loading / GPU の読み込みを高速化するためのピンメモリ",
+    )
     parser.add_argument("--num_beams", type=int, default=1, help="num of beams in beam search /beam search時のビーム数（多いと精度が上がるが時間がかかる）")
     parser.add_argument("--top_p", type=float, default=0.9, help="top_p in Nucleus sampling / Nucleus sampling時のtop_p")
     parser.add_argument("--max_length", type=int, default=75, help="max length of caption / captionの最大長")

diff --git a/finetune/make_captions_by_git.py b/finetune/make_captions_by_git.py
@@ -165,6 +165,11 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="enable image reading by DataLoader with this number of workers (faster) / DataLoaderによる画像読み込みを有効にしてこのワーカー数を適用する（読み込みを高速化）",
     )
+    parser.add_argument(
+        "--pin_memory",
+        action="store_true",
+        help="Pin memory for faster GPU loading / GPU の読み込みを高速化するためのピンメモリ",
+    )
     parser.add_argument("--max_length", type=int, default=50, help="max length of caption / captionの最大長")
     parser.add_argument(
         "--remove_words",

diff --git a/finetune/prepare_buckets_latents.py b/finetune/prepare_buckets_latents.py
@@ -224,6 +224,11 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="enable image reading by DataLoader with this number of workers (faster) / DataLoaderによる画像読み込みを有効にしてこのワーカー数を適用する（読み込みを高速化）",
     )
+    parser.add_argument(
+        "--pin_memory",
+        action="store_true",
+        help="Pin memory for faster GPU loading / GPU の読み込みを高速化するためのピンメモリ",
+    )
     parser.add_argument(
         "--max_resolution",
         type=str,

diff --git a/finetune/tag_images_by_wd14_tagger.py b/finetune/tag_images_by_wd14_tagger.py
@@ -410,6 +410,11 @@ def setup_parser() -> argparse.ArgumentParser:
         default=None,
         help="enable image reading by DataLoader with this number of workers (faster) / DataLoaderによる画像読み込みを有効にしてこのワーカー数を適用する（読み込みを高速化）",
     )
+    parser.add_argument(
+        "--pin_memory",
+        action="store_true",
+        help="Pin memory for faster GPU loading / GPU の読み込みを高速化するためのピンメモリ",
+    )
     parser.add_argument(
         "--caption_extention",
         type=str,