allow for multiple datasets for test data loader creation

Summary: Use all training dataset for export instead of just first. This is to support use cases where there is only a small amount of images per jsons but a number of jsons. Since calibration uses the first dataset, it is limited by the number of images in a single dataset.

Reviewed By: ppwwyyxx

Differential Revision: D28902673

fbshipit-source-id: f80146b02d2d1bc04703fbb21ef410f5e26ba64c

tools/exporter.py

@@ -43,7 +43,7 @@ def main(
     # NOTE: train dataset is used to avoid leakage since the data might be used for
     # running calibration for quantization. test_loader is used to make sure it follows
     # the inference behaviour (augmentation will not be applied).
-    datasets = cfg.DATASETS.TRAIN[0]
+    datasets = list(cfg.DATASETS.TRAIN)
     data_loader = runner.build_detection_test_loader(cfg, datasets)
 
     logger.info("Running the pytorch model and print FLOPS ...")