Fix distributed training test (#420)

Summary:
Pull Request resolved: #420

The distributed training test was running the local trainer on each process. The `distributed_backend` argument needed to be added since `classy_train.py` now runs the `LocalTrainer` by default.

Reviewed By: vreis

Differential Revision: D20296492

fbshipit-source-id: dbaee937f39d31e7c98db2836b98d7f97307783b

classy_train.py

@@ -43,6 +43,7 @@
 from pathlib import Path
 
 import torch
+from classy_vision.generic.distributed_util import get_rank, get_world_size
 from classy_vision.generic.opts import check_generic_args, parse_train_arguments
 from classy_vision.generic.registry_utils import import_all_packages_from_directory
 from classy_vision.generic.util import load_checkpoint, load_json
@@ -104,6 +105,10 @@ def main(args, config):
 
     trainer = trainer_class(use_gpu=use_gpu, num_dataloader_workers=args.num_workers)
 
+    logging.info(
+        f"Starting training on rank {get_rank()} worker. "
+        f"World size is {get_world_size()}"
+    )
     # That's it! When this call returns, training is done.
     trainer.train(task)
 

test/trainer_distributed_trainer_test.py

@@ -57,7 +57,8 @@ def test_training(self):
             --device={device} \
             --config={self.config_files[config_key]} \
             --num_workers=4 \
-            --log_freq=100
+            --log_freq=100 \
+            --distributed_backend=ddp
             """
             result = subprocess.run(cmd, shell=True)
             success = result.returncode == 0