ray_ddp: support logged_metrics as part of remote worker return value (#156)

Author: chongxiaoc

ray_lightning/ray_ddp.py

@@ -370,14 +370,20 @@ def post_dispatch(self, trainer: "pl.Trainer"):
 
         results = ray.get(self._futures)
         # Get the results, checkpoint path, and model weights from worker 0.
-        results, best_path, state_stream, callback_metrics = results[0]
+        results, best_path, state_stream, callback_metrics, logged_metrics \
+            = results[0]
         self._results = results
 
         # From DDPSpawn.get_queue
         self.lightning_module.trainer.callback_metrics.update(
             apply_to_collection(callback_metrics,
                                 np.ndarray, lambda x: torch.tensor(x)))
 
+        # Same for logged_metrics
+        self.lightning_module.trainer.logged_metrics.update(
+            apply_to_collection(logged_metrics,
+                                np.ndarray, lambda x: torch.tensor(x)))
+
         # DDPSpawnPlugin.__recover_child_process_weights begin
         # Difference here is that instead of writing the model weights to a
         # file and loading it, we use the state dict of the model directly.
@@ -500,8 +506,14 @@ def execute_remote(self,
                 torch.Tensor, lambda x: x.cpu().numpy(
                 ))  # send as numpy to avoid issues with memory sharing
 
+            # Same for logged_metrics
+            logged_metrics: dict = apply_to_collection(
+                self.lightning_module.trainer.logged_metrics,
+                torch.Tensor, lambda x: x.cpu().numpy(
+                ))  # send as numpy to avoid issues with memory sharing
+
             return_val = results, best_model_path, model_state_stream, \
-                callback_metrics
+                callback_metrics, logged_metrics
         else:
             return_val = None
         # __transfer_distrib_spawn_state_on_fit_end end

ray_lightning/tests/test_ddp.py

@@ -1,5 +1,6 @@
 import pytest
 from ray.util.client.ray_client_helpers import ray_start_client_server
+import torch
 from torch.utils.data import DistributedSampler
 
 from pl_bolts.datamodules import MNISTDataModule
@@ -12,7 +13,8 @@
 
 from ray_lightning import RayPlugin
 from ray_lightning.tests.utils import get_trainer, train_test, \
-    load_test, predict_test, BoringModel, LightningMNISTClassifier
+    load_test, predict_test, BoringModel, LightningMNISTClassifier, \
+    XORModel, XORDataModule
 
 
 @pytest.fixture
@@ -319,3 +321,30 @@ def on_train_start(self, trainer, pl_module):
     trainer = get_trainer(
         tmpdir, plugins=[plugin], callbacks=[UnusedParameterCallback()])
     trainer.fit(model)
+
+
+def test_metrics(tmpdir, ray_start_2_cpus):
+    """Tests if metrics are returned correctly"""
+    model = XORModel()
+    plugin = RayPlugin(num_workers=2, find_unused_parameters=False)
+    trainer = get_trainer(
+        tmpdir,
+        plugins=[plugin],
+        max_epochs=1,
+        num_sanity_val_steps=0,
+        reload_dataloaders_every_n_epochs=1)
+    dataset = XORDataModule()
+    trainer.fit(model, dataset)
+    callback_metrics = trainer.callback_metrics
+    logged_metrics = trainer.logged_metrics
+    assert callback_metrics["avg_val_loss"] == logged_metrics["avg_val_loss"]
+    assert logged_metrics["val_foo"] == torch.tensor(1.234)
+    assert callback_metrics["val_foo"] == torch.tensor(1.234)
+    # forked name is used for on_step logged metrics
+    forked_name_loss = "val_loss" + "_step"
+    forked_name_bar = "val_bar" + "_step"
+    assert forked_name_loss in logged_metrics.keys()
+    assert logged_metrics[forked_name_bar] == torch.tensor(5.678)
+    # callback_metrics doesn't record on_step metrics
+    assert forked_name_loss not in callback_metrics.keys()
+    assert forked_name_bar not in callback_metrics.keys()

ray_lightning/tests/utils.py

@@ -148,6 +148,68 @@ def validation_epoch_end(self, outputs):
         self.log("ptl/val_accuracy", avg_acc)
 
 
+class XORModel(LightningModule):
+    def __init__(self, input_dim=2, output_dim=1):
+        super(XORModel, self).__init__()
+        self.save_hyperparameters()
+        self.lin1 = torch.nn.Linear(input_dim, 8)
+        self.lin2 = torch.nn.Linear(8, output_dim)
+
+    def forward(self, features):
+        x = features.float()
+        x = self.lin1(x)
+        x = torch.tanh(x)
+        x = self.lin2(x)
+        x = torch.sigmoid(x)
+        return x
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=0.02)
+
+    def training_step(self, batch, batch_nb):
+        x, y = batch["x"], batch["y"].unsqueeze(1)
+        y_hat = self(x)
+        loss = F.binary_cross_entropy(y_hat, y.float())
+        return loss
+
+    def validation_step(self, batch, batch_nb):
+        x, y = batch["x"], batch["y"].unsqueeze(1)
+        y_hat = self(x)
+        loss = F.binary_cross_entropy(y_hat, y.float())
+        self.log("val_loss", loss, on_step=True)
+        # Log a constant for test purpose
+        self.log("val_bar", torch.tensor(5.678), on_step=True)
+        return loss
+
+    def validation_epoch_end(self, outputs):
+        avg_loss = torch.stack(outputs).mean()
+        self.log("avg_val_loss", avg_loss)
+        # Log a constant for test purpose
+        self.log("val_foo", torch.tensor(1.234))
+
+
+class XORDataModule(LightningDataModule):
+    def train_dataloader(self):
+        input_train = [{
+            "x": torch.tensor([[0.0, 0.0]]),
+            "y": torch.tensor([0])
+        }, {
+            "x": torch.tensor([[1.0, 1.0]]),
+            "y": torch.tensor([0])
+        }]
+        return iter(input_train)
+
+    def val_dataloader(self):
+        input_val = [{
+            "x": torch.tensor([[0.0, 1.0]]),
+            "y": torch.tensor([1])
+        }, {
+            "x": torch.tensor([[1.0, 0.0]]),
+            "y": torch.tensor([1])
+        }]
+        return iter(input_val)
+
+
 def get_trainer(dir,
                 plugins: List[PLUGIN_INPUT],
                 max_epochs: int = 1,

requirements-test.txt

@@ -10,3 +10,4 @@ ray[tune]
 torch==1.8.1
 torchmetrics
 torchvision
+protobuf<=3.20.1