Repository: pytorch/ignite

Issue:
Improve Neptune logger 
## 🚀 Feature

Idea is to be able use neptune API directly from an instance of `NeptuneLogger`. Same is done for MLflow logger:
https://github.com/pytorch/ignite/blob/29058e2f0e7d753cbd8dbb254652288afa81bda0/ignite/contrib/handlers/mlflow_logger.py#L248-L254


Discussion:

Resolving patch:
diff --git a/ignite/contrib/handlers/neptune_logger.py b/ignite/contrib/handlers/neptune_logger.py
--- a/ignite/contrib/handlers/neptune_logger.py
+++ b/ignite/contrib/handlers/neptune_logger.py
@@ -87,12 +87,12 @@ def global_step_transform(*args, **kwargs):
                              event_name=Events.EPOCH_COMPLETED)
 
     Args:
-        tag (str): common title for all produced plots. For example, 'training'
+        tag (str): common title for all produced plots. For example, "training"
         metric_names (list of str, optional): list of metric names to plot or a string "all" to plot all available
             metrics.
         output_transform (callable, optional): output transform function to prepare `engine.state.output` as a number.
             For example, `output_transform = lambda output: output`
-            This function can also return a dictionary, e.g `{'loss': loss1, 'another_loss': loss2}` to label the plot
+            This function can also return a dictionary, e.g `{"loss": loss1, "another_loss": loss2}` to label the plot
             with corresponding keys.
         another_engine (Engine): Deprecated (see :attr:`global_step_transform`). Another engine to use to provide the
             value of event. Typically, user can provide
@@ -121,7 +121,7 @@ def __init__(self, tag, metric_names=None, output_transform=None, another_engine
     def __call__(self, engine, logger, event_name):
 
         if not isinstance(logger, NeptuneLogger):
-            raise RuntimeError("Handler 'OutputHandler' works only with NeptuneLogger")
+            raise RuntimeError("Handler OutputHandler works only with NeptuneLogger")
 
         metrics = self._setup_output_metrics(engine)
 
@@ -135,10 +135,10 @@ def __call__(self, engine, logger, event_name):
 
         for key, value in metrics.items():
             if isinstance(value, numbers.Number) or isinstance(value, torch.Tensor) and value.ndimension() == 0:
-                logger.experiment.log_metric("{}/{}".format(self.tag, key), x=global_step, y=value)
+                logger.log_metric("{}/{}".format(self.tag, key), x=global_step, y=value)
             elif isinstance(value, torch.Tensor) and value.ndimension() == 1:
                 for i, v in enumerate(value):
-                    logger.experiment.log_metric("{}/{}/{}".format(self.tag, key, i), x=global_step, y=v.item())
+                    logger.log_metric("{}/{}/{}".format(self.tag, key, i), x=global_step, y=v.item())
             else:
                 warnings.warn("NeptuneLogger output_handler can not log " "metrics value type {}".format(type(value)))
 
@@ -170,7 +170,7 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler):
     Args:
         optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
         param_name (str): parameter name
-        tag (str, optional): common title for all produced plots. For example, 'generator'
+        tag (str, optional): common title for all produced plots. For example, generator
     """
 
     def __init__(self, optimizer, param_name="lr", tag=None):
@@ -178,7 +178,7 @@ def __init__(self, optimizer, param_name="lr", tag=None):
 
     def __call__(self, engine, logger, event_name):
         if not isinstance(logger, NeptuneLogger):
-            raise RuntimeError("Handler 'OptimizerParamsHandler' works only with NeptuneLogger")
+            raise RuntimeError("Handler OptimizerParamsHandler works only with NeptuneLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
         tag_prefix = "{}/".format(self.tag) if self.tag else ""
@@ -188,7 +188,7 @@ def __call__(self, engine, logger, event_name):
         }
 
         for k, v in params.items():
-            logger.experiment.log_metric(k, x=global_step, y=v)
+            logger.log_metric(k, x=global_step, y=v)
 
 
 class WeightsScalarHandler(BaseWeightsScalarHandler):
@@ -220,7 +220,7 @@ class WeightsScalarHandler(BaseWeightsScalarHandler):
     Args:
         model (torch.nn.Module): model to log weights
         reduction (callable): function to reduce parameters into scalar
-        tag (str, optional): common title for all produced plots. For example, 'generator'
+        tag (str, optional): common title for all produced plots. For example, generator
 
     """
 
@@ -230,7 +230,7 @@ def __init__(self, model, reduction=torch.norm, tag=None):
     def __call__(self, engine, logger, event_name):
 
         if not isinstance(logger, NeptuneLogger):
-            raise RuntimeError("Handler 'WeightsScalarHandler' works only with NeptuneLogger")
+            raise RuntimeError("Handler WeightsScalarHandler works only with NeptuneLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
         tag_prefix = "{}/".format(self.tag) if self.tag else ""
@@ -239,7 +239,7 @@ def __call__(self, engine, logger, event_name):
                 continue
 
             name = name.replace(".", "/")
-            logger.experiment.log_metric(
+            logger.log_metric(
                 "{}weights_{}/{}".format(tag_prefix, self.reduction.__name__, name),
                 x=global_step,
                 y=self.reduction(p.data),
@@ -275,7 +275,7 @@ class GradsScalarHandler(BaseWeightsScalarHandler):
     Args:
         model (torch.nn.Module): model to log weights
         reduction (callable): function to reduce parameters into scalar
-        tag (str, optional): common title for all produced plots. For example, 'generator'
+        tag (str, optional): common title for all produced plots. For example, generator
 
     """
 
@@ -284,7 +284,7 @@ def __init__(self, model, reduction=torch.norm, tag=None):
 
     def __call__(self, engine, logger, event_name):
         if not isinstance(logger, NeptuneLogger):
-            raise RuntimeError("Handler 'GradsScalarHandler' works only with NeptuneLogger")
+            raise RuntimeError("Handler GradsScalarHandler works only with NeptuneLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
         tag_prefix = "{}/".format(self.tag) if self.tag else ""
@@ -293,7 +293,7 @@ def __call__(self, engine, logger, event_name):
                 continue
 
             name = name.replace(".", "/")
-            logger.experiment.log_metric(
+            logger.log_metric(
                 "{}grads_{}/{}".format(tag_prefix, self.reduction.__name__, name),
                 x=global_step,
                 y=self.reduction(p.grad),
@@ -354,7 +354,7 @@ class NeptuneLogger(BaseLogger):
 
             # Attach the logger to the trainer to log training loss at each iteration
             npt_logger.attach(trainer,
-                             log_handler=OutputHandler(tag="training", output_transform=lambda loss: {'loss': loss}),
+                             log_handler=OutputHandler(tag="training", output_transform=lambda loss: {"loss": loss}),
                              event_name=Events.ITERATION_COMPLETED)
 
             # Attach the logger to the evaluator on the training dataset and log NLL, Accuracy metrics after each epoch
@@ -397,11 +397,11 @@ class NeptuneLogger(BaseLogger):
             from ignite.handlers import Checkpoint
 
             def score_function(engine):
-                return engine.state.metrics['accuracy']
+                return engine.state.metrics["accuracy"]
 
-            to_save = {'model': model}
+            to_save = {"model": model}
             handler = Checkpoint(to_save, NeptuneSaver(npt_logger), n_saved=2,
-                                 filename_prefix='best', score_function=score_function,
+                                 filename_prefix="best", score_function=score_function,
                                  score_name="validation_accuracy",
                                  global_step_transform=global_step_from_engine(trainer))
             validation_evaluator.add_event_handler(Events.COMPLETED, handler)
@@ -425,11 +425,20 @@ def score_function(engine):
                 # Attach the logger to the trainer to log training loss at each iteration
                 npt_logger.attach(trainer,
                                  log_handler=OutputHandler(tag="training",
-                                                           output_transform=lambda loss: {'loss': loss}),
+                                                           output_transform=lambda loss: {"loss": loss}),
                                  event_name=Events.ITERATION_COMPLETED)
 
     """
 
+    def __getattr__(self, attr):
+
+        import neptune
+
+        def wrapper(*args, **kwargs):
+            return getattr(neptune, attr)(*args, **kwargs)
+
+        return wrapper
+
     def __init__(self, *args, **kwargs):
         try:
             import neptune
@@ -444,8 +453,9 @@ def __init__(self, *args, **kwargs):
             neptune.init(project_qualified_name="dry-run/project", backend=neptune.OfflineBackend())
         else:
             self.mode = "online"
-            neptune.init(api_token=kwargs["api_token"], project_qualified_name=kwargs["project_name"])
+            neptune.init(api_token=kwargs.get("api_token"), project_qualified_name=kwargs.get("project_name"))
 
+        kwargs["name"] = kwargs.pop("experiment_name", None)
         self._experiment_kwargs = {
             k: v for k, v in kwargs.items() if k not in ["api_token", "project_name", "offline_mode"]
         }
@@ -453,7 +463,7 @@ def __init__(self, *args, **kwargs):
         self.experiment = neptune.create_experiment(**self._experiment_kwargs)
 
     def close(self):
-        self.experiment.stop()
+        self.stop()
 
 
 class NeptuneSaver:
@@ -486,20 +496,20 @@ class NeptuneSaver:
             from ignite.handlers import Checkpoint
 
             def score_function(engine):
-                return engine.state.metrics['accuracy']
+                return engine.state.metrics["accuracy"]
 
-            to_save = {'model': model}
+            to_save = {"model": model}
 
             # pass neptune logger to NeptuneServer
 
             handler = Checkpoint(to_save, NeptuneSaver(npt_logger), n_saved=2,
-                                 filename_prefix='best', score_function=score_function,
+                                 filename_prefix="best", score_function=score_function,
                                  score_name="validation_accuracy",
                                  global_step_transform=global_step_from_engine(trainer))
 
             evaluator.add_event_handler(Events.COMPLETED, handler)
 
-            # We need to close the logger when we are done
+        # We need to close the logger when we are done
             npt_logger.close()
 
     For example, you can access model checkpoints and download them from here:
@@ -508,13 +518,13 @@ def score_function(engine):
     """
 
     def __init__(self, neptune_logger: NeptuneLogger):
-        self._experiment = neptune_logger.experiment
+        self._logger = neptune_logger
 
     def __call__(self, checkpoint: Mapping, filename: str) -> None:
 
         with tempfile.NamedTemporaryFile() as tmp:
             torch.save(checkpoint, tmp.name)
-            self._experiment.log_artifact(tmp.name, filename)
+            self._logger.log_artifact(tmp.name, filename)
 
     def remove(self, filename: str) -> None:
-        self._experiment.delete_artifacts(filename)
+        self._logger.delete_artifacts(filename)


Test patch:
diff --git a/tests/ignite/contrib/handlers/test_neptune_logger.py b/tests/ignite/contrib/handlers/test_neptune_logger.py
--- a/tests/ignite/contrib/handlers/test_neptune_logger.py
+++ b/tests/ignite/contrib/handlers/test_neptune_logger.py
@@ -18,7 +18,7 @@ def test_optimizer_params_handler_wrong_setup():
 
     mock_logger = MagicMock()
     mock_engine = MagicMock()
-    with pytest.raises(RuntimeError, match="Handler 'OptimizerParamsHandler' works only with NeptuneLogger"):
+    with pytest.raises(RuntimeError, match="Handler OptimizerParamsHandler works only with NeptuneLogger"):
         handler(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
@@ -26,20 +26,20 @@ def test_optimizer_params():
     optimizer = torch.optim.SGD([torch.Tensor(0)], lr=0.01)
     wrapper = OptimizerParamsHandler(optimizer=optimizer, param_name="lr")
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.iteration = 123
 
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
-    mock_logger.experiment.log_metric.assert_called_once_with("lr/group_0", y=0.01, x=123)
+    mock_logger.log_metric.assert_called_once_with("lr/group_0", y=0.01, x=123)
 
     wrapper = OptimizerParamsHandler(optimizer, param_name="lr", tag="generator")
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
+    mock_logger.log_metric = MagicMock()
 
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
-    mock_logger.experiment.log_metric.assert_called_once_with("generator/lr/group_0", y=0.01, x=123)
+    mock_logger.log_metric.assert_called_once_with("generator/lr/group_0", y=0.01, x=123)
 
 
 def test_output_handler_with_wrong_logger_type():
@@ -47,61 +47,56 @@ def test_output_handler_with_wrong_logger_type():
 
     mock_logger = MagicMock()
     mock_engine = MagicMock()
-    with pytest.raises(RuntimeError, match="Handler 'OutputHandler' works only with NeptuneLogger"):
+    with pytest.raises(RuntimeError, match="Handler OutputHandler works only with NeptuneLogger"):
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
 def test_output_handler_output_transform():
     wrapper = OutputHandler("tag", output_transform=lambda x: x)
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.output = 12345
     mock_engine.state.iteration = 123
 
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
-
-    mock_logger.experiment.log_metric.assert_called_once_with("tag/output", y=12345, x=123)
+    mock_logger.log_metric.assert_called_once_with("tag/output", y=12345, x=123)
 
     wrapper = OutputHandler("another_tag", output_transform=lambda x: {"loss": x})
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
+    mock_logger.log_metric = MagicMock()
 
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
-    mock_logger.experiment.log_metric.assert_called_once_with("another_tag/loss", y=12345, x=123)
+    mock_logger.log_metric.assert_called_once_with("another_tag/loss", y=12345, x=123)
 
 
 def test_output_handler_metric_names():
     wrapper = OutputHandler("tag", metric_names=["a", "b"])
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State(metrics={"a": 12.23, "b": 23.45})
     mock_engine.state.iteration = 5
 
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
-    assert mock_logger.experiment.log_metric.call_count == 2
-    mock_logger.experiment.log_metric.assert_has_calls(
-        [call("tag/a", y=12.23, x=5), call("tag/b", y=23.45, x=5),], any_order=True
-    )
+    assert mock_logger.log_metric.call_count == 2
+    mock_logger.log_metric.assert_has_calls([call("tag/a", y=12.23, x=5), call("tag/b", y=23.45, x=5),], any_order=True)
 
     wrapper = OutputHandler("tag", metric_names=["a",])
 
     mock_engine = MagicMock()
+    mock_logger.log_metric = MagicMock()
     mock_engine.state = State(metrics={"a": torch.Tensor([0.0, 1.0, 2.0, 3.0])})
     mock_engine.state.iteration = 5
 
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
-    assert mock_logger.experiment.log_metric.call_count == 4
-    mock_logger.experiment.log_metric.assert_has_calls(
+    assert mock_logger.log_metric.call_count == 4
+    mock_logger.log_metric.assert_has_calls(
         [
             call("tag/a/0", y=0.0, x=5),
             call("tag/a/1", y=1.0, x=5),
@@ -114,40 +109,37 @@ def test_output_handler_metric_names():
     wrapper = OutputHandler("tag", metric_names=["a", "c"])
 
     mock_engine = MagicMock()
+    mock_logger.log_metric = MagicMock()
     mock_engine.state = State(metrics={"a": 55.56, "c": "Some text"})
     mock_engine.state.iteration = 7
 
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
+    mock_logger.log_metric = MagicMock()
 
     with pytest.warns(UserWarning):
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
-    assert mock_logger.experiment.log_metric.call_count == 1
-    mock_logger.experiment.log_metric.assert_has_calls([call("tag/a", y=55.56, x=7),], any_order=True)
+    assert mock_logger.log_metric.call_count == 1
+    mock_logger.log_metric.assert_has_calls([call("tag/a", y=55.56, x=7),], any_order=True)
 
     # all metrics
     wrapper = OutputHandler("tag", metric_names="all")
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State(metrics={"a": 12.23, "b": 23.45})
     mock_engine.state.iteration = 5
 
     wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
-    assert mock_logger.experiment.log_metric.call_count == 2
-    mock_logger.experiment.log_metric.assert_has_calls(
-        [call("tag/a", y=12.23, x=5), call("tag/b", y=23.45, x=5),], any_order=True
-    )
+    assert mock_logger.log_metric.call_count == 2
+    mock_logger.log_metric.assert_has_calls([call("tag/a", y=12.23, x=5), call("tag/b", y=23.45, x=5),], any_order=True)
 
 
 def test_output_handler_both():
     wrapper = OutputHandler("tag", metric_names=["a", "b"], output_transform=lambda x: {"loss": x})
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State(metrics={"a": 12.23, "b": 23.45})
     mock_engine.state.epoch = 5
@@ -155,8 +147,8 @@ def test_output_handler_both():
 
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    assert mock_logger.experiment.log_metric.call_count == 3
-    mock_logger.experiment.log_metric.assert_has_calls(
+    assert mock_logger.log_metric.call_count == 3
+    mock_logger.log_metric.assert_has_calls(
         [call("tag/a", y=12.23, x=5), call("tag/b", y=23.45, x=5), call("tag/loss", y=12345, x=5)], any_order=True
     )
 
@@ -167,8 +159,6 @@ def global_step_transform(*args, **kwargs):
 
     wrapper = OutputHandler("tag", output_transform=lambda x: {"loss": x}, global_step_transform=global_step_transform)
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.epoch = 5
@@ -191,16 +181,15 @@ def test_output_handler_with_global_step_from_engine():
     )
 
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.epoch = 1
     mock_engine.state.output = 0.123
 
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
-    assert mock_logger.experiment.log_metric.call_count == 1
-    mock_logger.experiment.log_metric.assert_has_calls(
+    assert mock_logger.log_metric.call_count == 1
+    mock_logger.log_metric.assert_has_calls(
         [call("tag/loss", y=mock_engine.state.output, x=mock_another_engine.state.epoch)]
     )
 
@@ -208,8 +197,8 @@ def test_output_handler_with_global_step_from_engine():
     mock_engine.state.output = 1.123
 
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
-    assert mock_logger.experiment.log_metric.call_count == 2
-    mock_logger.experiment.log_metric.assert_has_calls(
+    assert mock_logger.log_metric.call_count == 2
+    mock_logger.log_metric.assert_has_calls(
         [call("tag/loss", y=mock_engine.state.output, x=mock_another_engine.state.epoch)]
     )
 
@@ -220,16 +209,15 @@ def global_step_transform(*args, **kwargs):
 
     wrapper = OutputHandler("tag", output_transform=lambda x: {"loss": x}, global_step_transform=global_step_transform)
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.epoch = 5
     mock_engine.state.output = 12345
 
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
-    assert mock_logger.experiment.log_metric.call_count == 1
-    mock_logger.experiment.log_metric.assert_has_calls([call("tag/loss", y=12345, x=10)])
+    assert mock_logger.log_metric.call_count == 1
+    mock_logger.log_metric.assert_has_calls([call("tag/loss", y=12345, x=10)])
 
 
 def test_weights_scalar_handler_wrong_setup():
@@ -246,7 +234,7 @@ def test_weights_scalar_handler_wrong_setup():
     wrapper = WeightsScalarHandler(model)
     mock_logger = MagicMock()
     mock_engine = MagicMock()
-    with pytest.raises(RuntimeError, match="Handler 'WeightsScalarHandler' works only with NeptuneLogger"):
+    with pytest.raises(RuntimeError, match="Handler WeightsScalarHandler works only with NeptuneLogger"):
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
@@ -257,8 +245,7 @@ def test_weights_scalar_handler(dummy_model_factory):
     def _test(tag=None):
         wrapper = WeightsScalarHandler(model, tag=tag)
         mock_logger = MagicMock(spec=NeptuneLogger)
-        mock_logger.experiment = MagicMock()
-
+        mock_logger.log_metric = MagicMock()
         mock_engine = MagicMock()
         mock_engine.state = State()
         mock_engine.state.epoch = 5
@@ -267,8 +254,8 @@ def _test(tag=None):
 
         tag_prefix = "{}/".format(tag) if tag else ""
 
-        assert mock_logger.experiment.log_metric.call_count == 4
-        mock_logger.experiment.log_metric.assert_has_calls(
+        assert mock_logger.log_metric.call_count == 4
+        mock_logger.log_metric.assert_has_calls(
             [
                 call(tag_prefix + "weights_norm/fc1/weight", y=0.0, x=5),
                 call(tag_prefix + "weights_norm/fc1/bias", y=0.0, x=5),
@@ -287,26 +274,25 @@ def test_weights_scalar_handler_frozen_layers(dummy_model_factory):
 
     wrapper = WeightsScalarHandler(model)
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.epoch = 5
 
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    mock_logger.experiment.log_metric.assert_has_calls(
+    mock_logger.log_metric.assert_has_calls(
         [call("weights_norm/fc2/weight", y=12.0, x=5), call("weights_norm/fc2/bias", y=math.sqrt(12.0), x=5),],
         any_order=True,
     )
 
     with pytest.raises(AssertionError):
-        mock_logger.experiment.log_metric.assert_has_calls(
+        mock_logger.log_metric.assert_has_calls(
             [call("weights_norm/fc1/weight", y=12.0, x=5), call("weights_norm/fc1/bias", y=math.sqrt(12.0), x=5),],
             any_order=True,
         )
 
-    assert mock_logger.experiment.log_metric.call_count == 2
+    assert mock_logger.log_metric.call_count == 2
 
 
 def test_grads_scalar_handler_wrong_setup():
@@ -320,7 +306,7 @@ def test_grads_scalar_handler_wrong_setup():
     wrapper = GradsScalarHandler(model)
     mock_logger = MagicMock()
     mock_engine = MagicMock()
-    with pytest.raises(RuntimeError, match="Handler 'GradsScalarHandler' works only with NeptuneLogger"):
+    with pytest.raises(RuntimeError, match="Handler GradsScalarHandler works only with NeptuneLogger"):
         wrapper(mock_engine, mock_logger, Events.ITERATION_STARTED)
 
 
@@ -331,8 +317,7 @@ def test_grads_scalar_handler(dummy_model_factory, norm_mock):
     def _test(tag=None):
         wrapper = GradsScalarHandler(model, reduction=norm_mock, tag=tag)
         mock_logger = MagicMock(spec=NeptuneLogger)
-        mock_logger.experiment = MagicMock()
-
+        mock_logger.log_metric = MagicMock()
         mock_engine = MagicMock()
         mock_engine.state = State()
         mock_engine.state.epoch = 5
@@ -342,7 +327,7 @@ def _test(tag=None):
 
         tag_prefix = "{}/".format(tag) if tag else ""
 
-        mock_logger.experiment.log_metric.assert_has_calls(
+        mock_logger.log_metric.assert_has_calls(
             [
                 call(tag_prefix + "grads_norm/fc1/weight", y=ANY, x=5),
                 call(tag_prefix + "grads_norm/fc1/bias", y=ANY, x=5),
@@ -351,7 +336,7 @@ def _test(tag=None):
             ],
             any_order=True,
         )
-        assert mock_logger.experiment.log_metric.call_count == 4
+        assert mock_logger.log_metric.call_count == 4
         assert norm_mock.call_count == 4
 
     _test()
@@ -363,8 +348,7 @@ def test_grads_scalar_handler_frozen_layers(dummy_model_factory, norm_mock):
 
     wrapper = GradsScalarHandler(model, reduction=norm_mock)
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_metric = MagicMock()
     mock_engine = MagicMock()
     mock_engine.state = State()
     mock_engine.state.epoch = 5
@@ -372,15 +356,15 @@ def test_grads_scalar_handler_frozen_layers(dummy_model_factory, norm_mock):
 
     wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    mock_logger.experiment.log_metric.assert_has_calls(
+    mock_logger.log_metric.assert_has_calls(
         [call("grads_norm/fc2/weight", y=ANY, x=5), call("grads_norm/fc2/bias", y=ANY, x=5),], any_order=True
     )
 
     with pytest.raises(AssertionError):
-        mock_logger.experiment.log_metric.assert_has_calls(
+        mock_logger.log_metric.assert_has_calls(
             [call("grads_norm/fc1/weight", y=ANY, x=5), call("grads_norm/fc1/bias", y=ANY, x=5),], any_order=True
         )
-    assert mock_logger.experiment.log_metric.call_count == 2
+    assert mock_logger.log_metric.call_count == 2
     assert norm_mock.call_count == 2
 
 
@@ -400,7 +384,7 @@ def update_fn(engine, batch):
 
     def dummy_handler(engine, logger, event_name):
         global_step = engine.state.get_event_attrib_value(event_name)
-        logger.experiment.log_metric("test_value", global_step, global_step)
+        logger.log_metric("test_value", global_step, global_step)
 
     npt_logger.attach(trainer, log_handler=dummy_handler, event_name=Events.EPOCH_COMPLETED)
 
@@ -423,7 +407,7 @@ def update_fn(engine, batch):
 
         def dummy_handler(engine, logger, event_name):
             global_step = engine.state.get_event_attrib_value(event_name)
-            logger.experiment.log_metric("test_value", global_step, global_step)
+            logger.log_metric("test_value", global_step, global_step)
 
         npt_logger.attach(trainer, log_handler=dummy_handler, event_name=Events.EPOCH_COMPLETED)
 
@@ -433,8 +417,7 @@ def dummy_handler(engine, logger, event_name):
 def test_neptune_saver_serializable(dummy_model_factory, dirname):
 
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
-
+    mock_logger.log_artifact = MagicMock()
     model = torch.nn.Module()
     to_save_serializable = {"model": model}
 
@@ -442,13 +425,13 @@ def test_neptune_saver_serializable(dummy_model_factory, dirname):
     fname = "test.pt"
     saver(to_save_serializable, fname)
 
-    assert mock_logger.experiment.log_artifact.call_count == 1
+    assert mock_logger.log_artifact.call_count == 1
 
 
 def test_neptune_saver_non_serializable(dirname):
 
     mock_logger = MagicMock(spec=NeptuneLogger)
-    mock_logger.experiment = MagicMock()
+    mock_logger.log_artifact = MagicMock()
 
     to_save_non_serializable = {"model": lambda x: x}
 
@@ -463,7 +446,7 @@ def test_neptune_saver_non_serializable(dirname):
     except Exception:
         pass
 
-    assert mock_logger.experiment.log_artifact.call_count == 0
+    assert mock_logger.log_artifact.call_count == 0
 
 
 @pytest.fixture