Merge pull request #254 from RangiLyu/fix/ddp_eval

[BUG FIX] Fix evaluation bug when using multi-GPU training with pytorch-lightning.

Author: RangiLyu

nanodet/trainer/task.py

@@ -18,11 +18,12 @@
 import json
 import torch
 import logging
+
 from pytorch_lightning import LightningModule
-from typing import Any, List, Dict, Tuple, Optional
+from typing import Any, List
+from nanodet.util import mkdir, gather_results
 
 from ..model.arch import build_model
-from nanodet.util import mkdir
 
 
 class TrainingTask(LightningModule):
@@ -109,28 +110,32 @@ def validation_epoch_end(self, validation_step_outputs):
         results = {}
         for res in validation_step_outputs:
             results.update(res)
-        eval_results = self.evaluator.evaluate(results, self.cfg.save_dir, rank=self.local_rank)
-        metric = eval_results[self.cfg.evaluator.save_key]
-        # save best model
-        if metric > self.save_flag:
-            self.save_flag = metric
-            best_save_path = os.path.join(self.cfg.save_dir, 'model_best')
-            mkdir(self.local_rank, best_save_path)
-            self.trainer.save_checkpoint(os.path.join(best_save_path, "model_best.ckpt"))
-            txt_path = os.path.join(best_save_path, "eval_results.txt")
-            if self.local_rank < 1:
-                with open(txt_path, "a") as f:
-                    f.write("Epoch:{}\n".format(self.current_epoch+1))
-                    for k, v in eval_results.items():
-                        f.write("{}: {}\n".format(k, v))
+        all_results = gather_results(results)
+        if all_results:
+            eval_results = self.evaluator.evaluate(all_results, self.cfg.save_dir, rank=self.local_rank)
+            metric = eval_results[self.cfg.evaluator.save_key]
+            # save best model
+            if metric > self.save_flag:
+                self.save_flag = metric
+                best_save_path = os.path.join(self.cfg.save_dir, 'model_best')
+                mkdir(self.local_rank, best_save_path)
+                self.trainer.save_checkpoint(os.path.join(best_save_path, "model_best.ckpt"))
+                txt_path = os.path.join(best_save_path, "eval_results.txt")
+                if self.local_rank < 1:
+                    with open(txt_path, "a") as f:
+                        f.write("Epoch:{}\n".format(self.current_epoch+1))
+                        for k, v in eval_results.items():
+                            f.write("{}: {}\n".format(k, v))
+            else:
+                warnings.warn('Warning! Save_key is not in eval results! Only save model last!')
+            if self.log_style == 'Lightning':
+                for k, v in eval_results.items():
+                    self.log('Val_metrics/' + k, v, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            elif self.log_style == 'NanoDet':
+                for k, v in eval_results.items():
+                    self.scalar_summary('Val_metrics/' + k, 'Val', v, self.current_epoch+1)
         else:
-            warnings.warn('Warning! Save_key is not in eval results! Only save model last!')
-        if self.log_style == 'Lightning':
-            for k, v in eval_results.items():
-                self.log('Val_metrics/' + k, v, on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
-        elif self.log_style == 'NanoDet':
-            for k, v in eval_results.items():
-                self.scalar_summary('Val_metrics/' + k, 'Val', v, self.current_epoch+1)
+            self.info('Skip val on rank {}'.format(self.local_rank))
 
     def test_step(self, batch, batch_idx):
         dets = self.predict(batch, batch_idx)
@@ -140,16 +145,20 @@ def test_epoch_end(self, test_step_outputs):
         results = {}
         for res in test_step_outputs:
             results.update(res)
-        res_json = self.evaluator.results2json(results)
-        json_path = os.path.join(self.cfg.save_dir, 'results.json')
-        json.dump(res_json, open(json_path, 'w'))
-
-        if self.cfg.test_mode == 'val':
-            eval_results = self.evaluator.evaluate(results, self.cfg.save_dir, rank=self.local_rank)
-            txt_path = os.path.join(self.cfg.save_dir, "eval_results.txt")
-            with open(txt_path, "a") as f:
-                for k, v in eval_results.items():
-                    f.write("{}: {}\n".format(k, v))
+        all_results = gather_results(results)
+        if all_results:
+            res_json = self.evaluator.results2json(all_results)
+            json_path = os.path.join(self.cfg.save_dir, 'results.json')
+            json.dump(res_json, open(json_path, 'w'))
+
+            if self.cfg.test_mode == 'val':
+                eval_results = self.evaluator.evaluate(all_results, self.cfg.save_dir, rank=self.local_rank)
+                txt_path = os.path.join(self.cfg.save_dir, "eval_results.txt")
+                with open(txt_path, "a") as f:
+                    for k, v in eval_results.items():
+                        f.write("{}: {}\n".format(k, v))
+        else:
+            self.info('Skip test on rank {}'.format(self.local_rank))
 
     def configure_optimizers(self):
         """

nanodet/util/__init__.py

@@ -10,3 +10,4 @@
 from .visualization import Visualizer, overlay_bbox_cv
 from .flops_counter import get_model_complexity_info
 from .misc import multi_apply, images_to_levels, unmap
+from .scatter_gather import gather_results, scatter_kwargs

nanodet/util/scatter_gather.py

@@ -1,5 +1,8 @@
+import pickle
+
 import torch
 from torch.autograd import Variable
+import torch.distributed as dist
 from torch.nn.parallel._functions import Scatter
 
 
@@ -10,6 +13,7 @@ def list_scatter(input, target_gpus, chunk_sizes):
         del input[:size]
     return tuple(ret)
 
+
 def scatter(inputs, target_gpus, dim=0, chunk_sizes=None):
     """
     Slices variables into approximately equal chunks and
@@ -42,4 +46,39 @@ def scatter_kwargs(inputs, kwargs, target_gpus, dim=0, chunk_sizes=None):
         kwargs.extend([{} for _ in range(len(inputs) - len(kwargs))])
     inputs = tuple(inputs)
     kwargs = tuple(kwargs)
-    return inputs, kwargs
+    return inputs, kwargs
+
+
+def gather_results(result_part):
+    rank = -1
+    world_size = 1
+    if dist.is_available() and dist.is_initialized():
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+
+    # dump result part to tensor with pickle
+    part_tensor = torch.tensor(
+        bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
+
+    # gather all result part tensor shape
+    shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
+    shape_list = [shape_tensor.clone() for _ in range(world_size)]
+    dist.all_gather(shape_list, shape_tensor)
+
+    # padding result part tensor to max length
+    shape_max = torch.tensor(shape_list).max()
+    part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
+    part_send[:shape_tensor[0]] = part_tensor
+    part_recv_list = [
+        part_tensor.new_zeros(shape_max) for _ in range(world_size)
+    ]
+
+    # gather all result dict
+    dist.all_gather(part_recv_list, part_send)
+
+    if rank < 1:
+        all_res = {}
+        for recv, shape in zip(part_recv_list, shape_list):
+            all_res.update(
+                pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
+        return all_res