switch from tqdm to rich

Author: Mathux

README.md

@@ -43,7 +43,7 @@ pip install torchmetrics==0.7
 pip install hydra-core --upgrade
 pip install hydra_colorlog --upgrade
 pip install shortuuid
-pip install tqdm
+pip install rich
 pip install pandas
 pip install transformers
 pip install psutil

configs/evaluate.yaml

@@ -23,8 +23,8 @@ logger_level: INFO
 defaults:
   - /path@path
   - /transforms/rots2joints/smplh@rots2joints
-  - override hydra/job_logging: console
-  - override hydra/hydra_logging: console
+  - override hydra/job_logging: rich # console
+  - override hydra/hydra_logging: rich # console
   - _self_
 
 machine:

configs/hydra/hydra_logging/rich.yaml

@@ -0,0 +1,19 @@
+version: 1
+
+formatters:
+  colorlog:
+    (): colorlog.ColoredFormatter
+    format: '[%(cyan)s%(asctime)s%(reset)s][%(purple)sHYDRA%(reset)s] %(message)s'
+    datefmt: '%d/%m/%y %H:%M:%S'
+
+handlers:
+  console:
+    class: rich.logging.RichHandler # logging.StreamHandler
+    formatter: colorlog
+
+root:
+  level: INFO
+  handlers:
+  - console
+
+disable_existing_loggers: false

configs/hydra/job_logging/rich.yaml

@@ -0,0 +1,64 @@
+version: 1
+
+filters:
+  onlyimportant:
+    (): temos.tools.logging.LevelsFilter
+    levels:
+    - CRITICAL
+    - ERROR
+    - WARNING
+  noimportant:
+    (): temos.tools.logging.LevelsFilter
+    levels:
+    - INFO
+    - DEBUG
+    - NOTSET
+
+formatters:
+  verysimple:
+    format: '%(message)s'
+  simple:
+    format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    datefmt: '%d/%m/%y %H:%M:%S'
+
+  colorlog:
+    (): colorlog.ColoredFormatter
+    format: '[%(cyan)s%(asctime)s%(reset)s][%(blue)s%(name)s%(reset)s][%(log_color)s%(levelname)s%(reset)s]
+      - %(message)s'
+    datefmt: '%d/%m/%y %H:%M:%S'
+
+  log_colors:
+    DEBUG: purple
+    INFO: green
+    WARNING: yellow
+    ERROR: red
+    CRITICAL: red
+
+handlers:
+  console:
+    class: rich.logging.RichHandler # logging.StreamHandler
+    formatter: verysimple # colorlog
+    rich_tracebacks: true
+
+  file_out:
+    class: logging.FileHandler
+    formatter: simple
+    filename: logs.out
+    filters:
+    - noimportant
+
+  file_err:
+    class: logging.FileHandler
+    formatter: simple
+    filename: logs.err
+    filters:
+    - onlyimportant
+
+root:
+    level: ${logger_level}
+    handlers:
+    - console
+    - file_out
+    - file_err
+
+disable_existing_loggers: false

configs/sample.yaml

@@ -33,8 +33,8 @@ defaults:
   - trainer: null
   - sampler: all_conseq
   - /path@path
-  - override hydra/job_logging: custom
-  - override hydra/hydra_logging: custom
+  - override hydra/job_logging: rich # custom
+  - override hydra/hydra_logging: rich # custom
   - _self_
 
 

configs/train.yaml

@@ -22,8 +22,8 @@ defaults:
   - logger: none
   - callback: base
   - /path@path
-  - override hydra/job_logging: custom
-  - override hydra/hydra_logging: custom
+  - override hydra/job_logging: rich # custom
+  - override hydra/hydra_logging: rich # custom
   - _self_
 
 data:

evaluate.py

@@ -24,22 +24,22 @@ def regroup_metrics(metrics):
         APE_pose = dico.pop("APE_pose")
         for name, ape in zip(pose_names, APE_pose):
             dico[f"APE_pose_{name}"] = ape
-    
+
     if "APE_joints" in dico:
         APE_joints = dico.pop("APE_joints")
         for name, ape in zip(mmm_joints, APE_joints):
             dico[f"APE_joints_{name}"] = ape
-    
+
     if "AVE_pose" in dico:
         AVE_pose = dico.pop("AVE_pose")
         for name, ave in zip(pose_names, AVE_pose):
             dico[f"AVE_pose_{name}"] = ave
-    
+
     if "AVE_joints" in dico:
         AVE_joints = dico.pop("AVE_joints")
         for name, ape in zip(mmm_joints, AVE_joints):
             dico[f"AVE_joints_{name}"] = ave
-    
+
     return dico
 
 def sanitize(dico):

sample.py

@@ -115,39 +115,42 @@ def sample(newcfg: DictConfig) -> None:
     dataset = getattr(data_module, f"{cfg.split}_dataset")
 
     from temos.data.sampling import upsample
-    from tqdm import tqdm
+    from rich.progress import Progress
+    from rich.progress import track
 
     # remove printing for changing the seed
     logging.getLogger('pytorch_lightning.utilities.seed').setLevel(logging.WARNING)
 
     import torch
     with torch.no_grad():
-        for keyid in (pbar := tqdm(dataset.keyids)):
-            pbar.set_description(f"Processing {keyid}")
-            for index in range(cfg.number_of_samples):
-                one_data = dataset.load_keyid(keyid)
-                # batch_size = 1 for reproductability
-                batch = collate_datastruct_and_text([one_data])
-                # fix the seed
-                pl.seed_everything(index)
-
-                if cfg.jointstype == "vertices":
-                    vertices = model(batch)[0]
-                    motion = vertices.numpy()
-                    # no upsampling here to keep memory
-                    # vertices = upsample(vertices, cfg.data.framerate, 100)
-                else:
-                    joints = model(batch)[0]
-                    motion = joints.numpy()
-                    # upscaling to compare with other methods
-                    motion = upsample(motion, cfg.data.framerate, 100)
-
-                if cfg.number_of_samples > 1:
-                    npypath = path / f"{keyid}_{index}.npy"
-                else:
-                    npypath = path / f"{keyid}.npy"
-
-                np.save(npypath, motion)
+        with Progress(transient=True) as progress:
+            task = progress.add_task("Sampling", total=len(dataset.keyids))
+            for keyid in dataset.keyids:
+                progress.update(task, description=f"Sampling {keyid}..")
+                for index in range(cfg.number_of_samples):
+                    one_data = dataset.load_keyid(keyid)
+                    # batch_size = 1 for reproductability
+                    batch = collate_datastruct_and_text([one_data])
+                    # fix the seed
+                    pl.seed_everything(index)
+
+                    if cfg.jointstype == "vertices":
+                        vertices = model(batch)[0]
+                        motion = vertices.numpy()
+                        # no upsampling here to keep memory
+                        # vertices = upsample(vertices, cfg.data.framerate, 100)
+                    else:
+                        joints = model(batch)[0]
+                        motion = joints.numpy()
+                        # upscaling to compare with other methods
+                        motion = upsample(motion, cfg.data.framerate, 100)
+
+                    if cfg.number_of_samples > 1:
+                        npypath = path / f"{keyid}_{index}.npy"
+                    else:
+                        npypath = path / f"{keyid}.npy"
+                    np.save(npypath, motion)
+                progress.update(task, advance=1)
 
     logger.info("All the sampling are done")
     logger.info(f"All the sampling are done. You can find them here:\n{path}")

temos/callback/progress.py

@@ -25,10 +25,11 @@ def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule,
         if trainer.sanity_checking:
             logger.info("Sanity checking ok.")
 
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule, **kwargs) -> None:
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule, padding=False, **kwargs) -> None:
         metric_format = f"{{:.{self.precision}e}}"
         line = f"Epoch {trainer.current_epoch}"
-        line = f"{line:>{len('Epoch xxxx')}}"  # Right padding
+        if padding:
+            line = f"{line:>{len('Epoch xxxx')}}"  # Right padding
         metrics_str = []
 
         losses_dict = trainer.callback_metrics

temos/data/kit.py

@@ -9,14 +9,15 @@
 import torch
 from torch import nn
 from torch.utils.data import Dataset
-from tqdm import tqdm
 from pathlib import Path
 
 from temos.tools.easyconvert import matrix_to, axis_angle_to
 from temos.transforms import Transform
 from temos.data.sampling import subsample
 from temos.data.tools.smpl import smpl_data_to_matrix_and_trans
 
+from rich.progress import track
+        
 from .base import BASEDataModule
 from .utils import get_split_keyids
 
@@ -90,7 +91,7 @@ def __init__(self, datapath: str,
                 kitml_correspondances = json.load(correspondance_path_file)
 
         if progress_bar:
-            enumerator = enumerate(tqdm(keyids, f"Loading KIT {split}"))
+            enumerator = enumerate(track(keyids, f"Loading KIT {split}"))
         else:
             enumerator = enumerate(keyids)
 

temos/tools/logging.py

@@ -1,6 +1,4 @@
 import logging
-import tqdm
-
 
 class LevelsFilter(logging.Filter):
     def __init__(self, levels):
@@ -32,6 +30,7 @@ def __init__(self, level=logging.NOTSET):
         super().__init__(level)
 
     def emit(self, record):
+        import tqdm
         try:
             msg = self.format(record)
             tqdm.tqdm.write(msg)

train.py

@@ -50,11 +50,14 @@ def train(cfg: DictConfig) -> None:
         "AVE root": "Metrics/AVE_root",
         "AVE mean pose": "Metrics/AVE_mean_pose"
     }
+
     callbacks = [
+        pl.callbacks.RichProgressBar(),
         instantiate(cfg.callback.progress, metric_monitor=metric_monitor),
         instantiate(cfg.callback.latest_ckpt),
         instantiate(cfg.callback.last_ckpt)
     ]
+
     logger.info("Callbacks initialized")
 
     logger.info("Loading trainer")