Clean up pytorch-lightning logger (rusty1s#210)

* initial commit:

* code cov

* move progressbar

bin/kumo-train

@@ -19,6 +19,7 @@ from kumo.config import cfg
 from kumo.model import create_model
 from kumo.train import create_loader
 from kumo.train.encoder import get_emb_size
+from kumo.train.trainer import ProgressBar
 from kumo.utils import tracing
 from kumo.utils.visualization import visualize_scalar_distribution
 
@@ -71,7 +72,9 @@ def main(args):
                                    mode='min')
             callbacks = [early_stopping, ckpt]
         else:
-            callbacks = None
+            callbacks = []
+
+        callbacks.append(ProgressBar())
 
         trainer = pl.Trainer(
             gpus=gpus,

kumo/train/trainer.py

@@ -1,5 +1,6 @@
 from typing import Any, List, Optional
 
+import pytorch_lightning as pl
 import torch
 
 from kumo.task import Task
@@ -20,3 +21,12 @@ def __init__(self, model: torch.nn.Module, task: Task, epochs: int,
 
     def fit(self):
         pass
+
+
+class ProgressBar(pl.callbacks.ProgressBar):
+    def get_metrics(self, trainer, model):
+        # Remove `loss` and `v_num` values from the logger.
+        items = super().get_metrics(trainer, model)
+        items.pop('loss', None)
+        items.pop('v_num', None)
+        return items

test/train/test_trainer.py

@@ -13,6 +13,7 @@
 from kumo.model import create_model
 from kumo.train import create_loader
 from kumo.train.encoder import get_emb_size
+from kumo.train.trainer import ProgressBar
 
 
 @pytest.mark.parametrize('task', [
@@ -49,5 +50,5 @@ def test_trainer(task):
         )
 
         # Training
-        trainer = pl.Trainer(fast_dev_run=True)
+        trainer = pl.Trainer(fast_dev_run=True, callbacks=[ProgressBar()])
         trainer.fit(model, loaders[0], loaders[1])