From 9b5e7ce21feb51977d11da4e6a0ed35f502dbfb5 Mon Sep 17 00:00:00 2001 From: Maruyama_Aya Date: Thu, 8 Jun 2023 14:56:56 +0800 Subject: [PATCH] modify shell for check --- examples/images/dreambooth/colossalai.sh | 1 + examples/images/dreambooth/test_ci.sh | 1 + examples/images/dreambooth/train_dreambooth_colossalai.py | 5 +++++ 3 files changed, 7 insertions(+) diff --git a/examples/images/dreambooth/colossalai.sh b/examples/images/dreambooth/colossalai.sh index b2a544928..db4562dbc 100755 --- a/examples/images/dreambooth/colossalai.sh +++ b/examples/images/dreambooth/colossalai.sh @@ -14,4 +14,5 @@ torchrun --nproc_per_node 4 --standalone train_dreambooth_colossalai.py \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ --num_class_images=200 \ + --test_run=True \ --placement="auto" \ diff --git a/examples/images/dreambooth/test_ci.sh b/examples/images/dreambooth/test_ci.sh index 0e3f6efa4..21f45adae 100644 --- a/examples/images/dreambooth/test_ci.sh +++ b/examples/images/dreambooth/test_ci.sh @@ -19,6 +19,7 @@ for plugin in "gemini"; do --learning_rate=5e-6 \ --lr_scheduler="constant" \ --lr_warmup_steps=0 \ + --test_run=True \ --num_class_images=200 \ --placement="auto" # "cuda" done diff --git a/examples/images/dreambooth/train_dreambooth_colossalai.py b/examples/images/dreambooth/train_dreambooth_colossalai.py index 44bde9226..888b28de8 100644 --- a/examples/images/dreambooth/train_dreambooth_colossalai.py +++ b/examples/images/dreambooth/train_dreambooth_colossalai.py @@ -198,6 +198,7 @@ def parse_args(input_args=None): parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.") parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.") parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.") + parser.add_argument("--test_run", default=False, help="Whether to use a smaller dataset for test run.") parser.add_argument( "--hub_model_id", type=str, @@ -267,6 +268,7 @@ class DreamBoothDataset(Dataset): class_prompt=None, size=512, center_crop=False, + test=False, ): self.size = size self.center_crop = center_crop @@ -277,6 +279,8 @@ class DreamBoothDataset(Dataset): raise ValueError("Instance images root doesn't exists.") self.instance_images_path = list(Path(instance_data_root).iterdir()) + if test: + self.instance_images_path = self.instance_images_path[:10] self.num_instance_images = len(self.instance_images_path) self.instance_prompt = instance_prompt self._length = self.num_instance_images @@ -509,6 +513,7 @@ def main(args): tokenizer=tokenizer, size=args.resolution, center_crop=args.center_crop, + test=args.test_run ) def collate_fn(examples):