refactor GANs trainer

Summary: Refactor and simplify the trainer by reusing the _write_metrics defined in SimpleTrainer.

Reviewed By: newstzpz

Differential Revision: D26038293

fbshipit-source-id: 174c386375759ce761ff2db80cb5a076a083bcc8

Author: Tao Xu

detectron2/engine/train_loop.py

@@ -244,7 +244,12 @@ def run_step(self):
         """
         self.optimizer.step()
 
-    def _write_metrics(self, loss_dict: Dict[str, torch.Tensor], data_time: float):
+    def _write_metrics(
+        self,
+        loss_dict: Dict[str, torch.Tensor],
+        data_time: float,
+        prefix: str = "",
+    ):
         """
         Args:
             loss_dict (dict): dict of scalar losses
@@ -281,7 +286,7 @@ def _write_metrics(self, loss_dict: Dict[str, torch.Tensor], data_time: float):
                     f"loss_dict = {metrics_dict}"
                 )
 
-            storage.put_scalar("total_loss", total_losses_reduced)
+            storage.put_scalar("{}total_loss".format(prefix), total_losses_reduced)
             if len(metrics_dict) > 1:
                 storage.put_scalars(**metrics_dict)