Initial commit

main.py

@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import argparse
+
+import multiprocessing as mp
+
+import pprint
+import yaml
+
+from src.utils.distributed import init_distributed
+from src.train import main as app_main
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--fname', type=str,
+    help='name of config file to load',
+    default='configs.yaml')
+parser.add_argument(
+    '--devices', type=str, nargs='+', default=['cuda:0'],
+    help='which devices to use on local machine')
+
+
+def process_main(rank, fname, world_size, devices):
+    import os
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(devices[rank].split(':')[-1])
+
+    import logging
+    logging.basicConfig()
+    logger = logging.getLogger()
+    if rank == 0:
+        logger.setLevel(logging.INFO)
+    else:
+        logger.setLevel(logging.ERROR)
+
+    logger.info(f'called-params {fname}')
+
+    # -- load script params
+    params = None
+    with open(fname, 'r') as y_file:
+        params = yaml.load(y_file, Loader=yaml.FullLoader)
+        logger.info('loaded params...')
+        pp = pprint.PrettyPrinter(indent=4)
+        pp.pprint(params)
+
+    world_size, rank = init_distributed(rank_and_world_size=(rank, world_size))
+    logger.info(f'Running... (rank: {rank}/{world_size})')
+    app_main(args=params)
+
+
+if __name__ == '__main__':
+    args = parser.parse_args()
+
+    num_gpus = len(args.devices)
+    mp.set_start_method('spawn')
+
+    for rank in range(num_gpus):
+        mp.Process(
+            target=process_main,
+            args=(rank, args.fname, num_gpus, args.devices)
+        ).start()