From 8211d113299d71bffc0167406fc8102729a9e9fa Mon Sep 17 00:00:00 2001 From: Victor Hall Date: Fri, 16 Dec 2022 23:03:33 -0500 Subject: [PATCH] test dataloader script update --- test/test_dl.py | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/test/test_dl.py b/test/test_dl.py index 9b52b6f..d0de4f0 100644 --- a/test/test_dl.py +++ b/test/test_dl.py @@ -1,18 +1,29 @@ # script to test data loader by itself -# run from training root, edit the data_root manually -# python ldm/data/test_dl.py +# run from trainer root directory +# Ex: +# python test/test_dl.py --data_root "x:/mytestdata/input" --batch_size 2 +import argparse + import ldm.data.data_loader as dl -data_root = "r:/everydream-trainer/test/input" +def main(data_root, batch_size): + data_loader = dl.DataLoaderMultiAspect(data_root=data_root, batch_size=batch_size, debug_level=1) -data_loader = dl.DataLoaderMultiAspect(data_root=data_root, batch_size=2, seed=555, debug_level=2) + image_caption_pairs = data_loader.get_all_images() -image_caption_pairs = data_loader.get_all_images() + print(f"Loaded {len(image_caption_pairs)} image-caption pairs") -print(f"Loaded {len(image_caption_pairs)} image-caption pairs") + print(f"**** Done loading. Loaded {len(image_caption_pairs)} images from data_root: {data_root} ****") -for image_caption_pair in image_caption_pairs: - print(image_caption_pair) +if __name__ == "__main__": + """ + test the data loader by itself, outputs buckets and image counts + data_root: root folder of training data + batch_size: number of images per batch + """ - -print(f"**** Done loading. Loaded {len(image_caption_pairs)} images from data_root: {data_root} ****") \ No newline at end of file + parser = argparse.ArgumentParser() + parser.add_argument("--data_root", type=str, default="input", help="root folder of training data") + parser.add_argument("--batch_size", type=int, default=4, help="number of images per batch") + args = parser.parse_args() + main(data_root=args.data_root, batch_size=args.batch_size) \ No newline at end of file