test TrackioLogger

Author: Xmaster6y

test/test_loggers.py

@@ -21,6 +21,7 @@
 from torchrl.record.loggers.csv import CSVLogger
 from torchrl.record.loggers.mlflow import _has_mlflow, _has_tv, MLFlowLogger
 from torchrl.record.loggers.tensorboard import _has_tb, TensorboardLogger
+from torchrl.record.loggers.trackio import _has_trackio, TrackioLogger
 from torchrl.record.loggers.wandb import _has_wandb, WandbLogger
 from torchrl.record.recorder import PixelRenderTransform, VideoRecorder
 
@@ -455,6 +456,65 @@ def make_env():
                 env.close()
 
 
+@pytest.fixture(scope="function")
+def trackio_logger():
+    exp_name = "ramala"
+    logger = TrackioLogger(project="test", exp_name=exp_name)
+    yield logger
+    logger.experiment.finish()
+    del logger
+
+
+@pytest.mark.skipif(not _has_trackio, reason="trackio not installed")
+class TestTrackioLogger:
+    @pytest.mark.parametrize("steps", [None, [1, 10, 11]])
+    def test_log_scalar(self, steps, trackio_logger):
+        torch.manual_seed(0)
+
+        values = torch.rand(3)
+        for i in range(3):
+            scalar_name = "foo"
+            scalar_value = values[i].item()
+            trackio_logger.log_scalar(
+                value=scalar_value,
+                name=scalar_name,
+                step=steps[i] if steps else None,
+            )
+
+    def test_log_video(self, trackio_logger):
+        torch.manual_seed(0)
+
+        # creating a sample video (T, C, H, W), where T - number of frames,
+        # C - number of image channels (e.g. 3 for RGB), H, W - image dimensions.
+        # the first 64 frames are black and the next 64 are white
+        video = torch.cat(
+            (torch.zeros(128, 1, 32, 32), torch.full((128, 1, 32, 32), 255))
+        )
+        video = video[None, :]
+        trackio_logger.log_video(
+            name="foo",
+            video=video,
+            fps=4,
+            format="mp4",
+        )
+        trackio_logger.log_video(
+            name="foo_16fps",
+            video=video,
+            fps=16,
+            format="mp4",
+        )
+
+    def test_log_hparams(self, trackio_logger, config):
+        trackio_logger.log_hparams(config)
+        for key, value in config.items():
+            assert trackio_logger.experiment.config[key] == value
+
+    def test_log_histogram(self, trackio_logger):
+        with pytest.raises(NotImplementedError):
+            data = torch.randn(10)
+            trackio_logger.log_histogram("hist", data, step=0, bins=2)
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)