From dfef7db4212163b1dd09d2b82dcf54f6bdee738b Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Mon, 31 Oct 2022 18:26:34 +0000 Subject: [PATCH] moar debugging --- aimodel/src/subcommands/pretrain.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/aimodel/src/subcommands/pretrain.py b/aimodel/src/subcommands/pretrain.py index 2506b8f..5a2f68a 100644 --- a/aimodel/src/subcommands/pretrain.py +++ b/aimodel/src/subcommands/pretrain.py @@ -21,6 +21,13 @@ def parse_args(): return parser + +def count_batches(dataset): + count = 0 + for _ in dataset: + count += 1 + return count + def run(args): if (not hasattr(args, "water_size")) or args.water_size == None: args.water_size = 100 @@ -40,7 +47,11 @@ def run(args): dirpath_input=args.input, batch_size=args.batch_size, ) - dataset_metadata = read_metadata(args.input) + + print("BATCHES_TRAIN", count_batches(dataset_train)) + print("BATCHES_VALIDATE", count_batches(dataset_validate)) + + # for (items, label) in dataset_train: # print("ITEMS", len(items), [ item.shape for item in items ]) @@ -59,4 +70,3 @@ def run(args): ) ai.train(dataset_train, dataset_validate) - \ No newline at end of file