[RaySGD] Rename PyTorch API endpoints to start with Torch (#7425)

* Start renaming pytorch to torch * Rename PyTorchTrainer to TorchTrainer * Rename PyTorch runners to Torch runners * Finish renaming API * Rename to torch in tests * Finish renaming docs + tests * Run format + fix DeprecationWarning * fix * move tests up * rename Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
2026-07-06 05:16:30 +08:00 · 2020-03-03 16:44:42 -08:00
parent f6883bf725
commit 3a134c7224
22 changed files with 222 additions and 218 deletions
@@ -1,4 +1,9 @@
-from ray.util.sgd.pytorch import PyTorchTrainer
+from ray.util.sgd.torch import TorchTrainer
 from ray.util.sgd.tf import TFTrainer

-__all__ = ["PyTorchTrainer", "TFTrainer"]
+__all__ = ["TorchTrainer", "TFTrainer"]
+
+
+def PyTorchTrainer(**kwargs):
+    raise DeprecationWarning("ray.util.sgd.pytorch.PyTorchTrainer has been "
+                             "renamed to ray.util.sgd.torch.TorchTrainer")
@@ -1,18 +0,0 @@
-import logging
-logger = logging.getLogger(__name__)
-
-PyTorchTrainer = None
-PyTorchTrainable = None
-TrainingOperator = None
-
-try:
-    import torch  # noqa: F401
-
-    from ray.util.sgd.pytorch.pytorch_trainer import (PyTorchTrainer,
-                                                      PyTorchTrainable)
-
-    from ray.util.sgd.pytorch.training_operator import TrainingOperator
-
-    __all__ = ["PyTorchTrainer", "PyTorchTrainable", "TrainingOperator"]
-except ImportError:
-    logger.warning("PyTorch not found. PyTorchTrainer will not be available")
@@ -10,12 +10,12 @@ import torch.distributed as dist

 import ray
 from ray import tune
-from ray.util.sgd.pytorch import PyTorchTrainer, PyTorchTrainable
-from ray.util.sgd.pytorch.training_operator import _TestingOperator
-from ray.util.sgd.pytorch.constants import BATCH_COUNT, SCHEDULER_STEP
+from ray.util.sgd.torch import TorchTrainer, TorchTrainable
+from ray.util.sgd.torch.training_operator import _TestingOperator
+from ray.util.sgd.torch.constants import BATCH_COUNT, SCHEDULER_STEP
 from ray.util.sgd.utils import check_for_failure

-from ray.util.sgd.pytorch.examples.train_example import (
+from ray.util.sgd.torch.examples.train_example import (
    model_creator, optimizer_creator, data_creator, LinearDataset)


@@ -28,7 +28,7 @@ def ray_start_2_cpus():


 def test_single_step(ray_start_2_cpus):  # noqa: F811
-    trainer = PyTorchTrainer(
+    trainer = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -44,7 +44,7 @@ def test_single_step(ray_start_2_cpus):  # noqa: F811
@pytest.mark.parametrize("num_replicas", [1, 2]
                         if dist.is_available() else [1])
 def test_train(ray_start_2_cpus, num_replicas):  # noqa: F811
-    trainer = PyTorchTrainer(
+    trainer = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -107,7 +107,7 @@ def test_multi_model(ray_start_2_cpus, num_replicas):
        ]
        return opts[0], opts[1]

-    trainer1 = PyTorchTrainer(
+    trainer1 = TorchTrainer(
        multi_model_creator,
        data_creator,
        multi_optimizer_creator,
@@ -124,7 +124,7 @@ def test_multi_model(ray_start_2_cpus, num_replicas):

    trainer1.shutdown()

-    trainer2 = PyTorchTrainer(
+    trainer2 = TorchTrainer(
        multi_model_creator,
        data_creator,
        multi_optimizer_creator,
@@ -193,7 +193,7 @@ def test_multi_model_matrix(ray_start_2_cpus, num_replicas):  # noqa: F811
    for model_count in range(1, 3):
        for optimizer_count in range(1, 3):
            for scheduler_count in range(1, 3):
-                trainer = PyTorchTrainer(
+                trainer = TorchTrainer(
                    multi_model_creator,
                    data_creator,
                    multi_optimizer_creator,
@@ -221,7 +221,7 @@ def test_scheduler_freq(ray_start_2_cpus, scheduler_freq):  # noqa: F811
        return torch.optim.lr_scheduler.StepLR(
            optimizer, step_size=30, gamma=0.1)

-    trainer = PyTorchTrainer(
+    trainer = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -239,7 +239,7 @@ def test_scheduler_freq(ray_start_2_cpus, scheduler_freq):  # noqa: F811
 def test_scheduler_validate(ray_start_2_cpus):  # noqa: F811
    from torch.optim.lr_scheduler import ReduceLROnPlateau

-    trainer = PyTorchTrainer(
+    trainer = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -273,7 +273,7 @@ def test_tune_train(ray_start_2_cpus, num_replicas):  # noqa: F811
    }

    analysis = tune.run(
-        PyTorchTrainable,
+        TorchTrainable,
        num_samples=2,
        config=config,
        stop={"training_iteration": 2},
@@ -293,7 +293,7 @@ def test_tune_train(ray_start_2_cpus, num_replicas):  # noqa: F811
@pytest.mark.parametrize("num_replicas", [1, 2]
                         if dist.is_available() else [1])
 def test_save_and_restore(ray_start_2_cpus, num_replicas):  # noqa: F811
-    trainer1 = PyTorchTrainer(
+    trainer1 = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -308,7 +308,7 @@ def test_save_and_restore(ray_start_2_cpus, num_replicas):  # noqa: F811

    trainer1.shutdown()

-    trainer2 = PyTorchTrainer(
+    trainer2 = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -346,8 +346,8 @@ def test_fail_with_recover(ray_start_2_cpus):  # noqa: F811
        success = check_for_failure(worker_stats)
        return success, worker_stats

-    with patch.object(PyTorchTrainer, "_train_epoch", step_with_fail):
-        trainer1 = PyTorchTrainer(
+    with patch.object(TorchTrainer, "_train_epoch", step_with_fail):
+        trainer1 = TorchTrainer(
            model_creator,
            single_loader,
            optimizer_creator,
@@ -376,8 +376,8 @@ def test_resize(ray_start_2_cpus):  # noqa: F811
        success = check_for_failure(worker_stats)
        return success, worker_stats

-    with patch.object(PyTorchTrainer, "_train_epoch", step_with_fail):
-        trainer1 = PyTorchTrainer(
+    with patch.object(TorchTrainer, "_train_epoch", step_with_fail):
+        trainer1 = TorchTrainer(
            model_creator,
            single_loader,
            optimizer_creator,
@@ -412,8 +412,8 @@ def test_fail_twice(ray_start_2_cpus):  # noqa: F811
        success = check_for_failure(worker_stats)
        return success, worker_stats

-    with patch.object(PyTorchTrainer, "_train_epoch", step_with_fail):
-        trainer1 = PyTorchTrainer(
+    with patch.object(TorchTrainer, "_train_epoch", step_with_fail):
+        trainer1 = TorchTrainer(
            model_creator,
            single_loader,
            optimizer_creator,
@@ -4,8 +4,8 @@ import torch.nn as nn
 import unittest
 from unittest.mock import MagicMock

-from ray.util.sgd.pytorch.training_operator import TrainingOperator
-from ray.util.sgd.pytorch.pytorch_runner import PyTorchRunner
+from ray.util.sgd.torch.training_operator import TrainingOperator
+from ray.util.sgd.torch.torch_runner import TorchRunner


 class LinearDataset(torch.utils.data.Dataset):
@@ -45,14 +45,14 @@ def create_dataloaders(config):
    return LinearDataset(2, 5), LinearDataset(2, 5, size=400)


-class TestPyTorchRunner(unittest.TestCase):
+class TestTorchRunner(unittest.TestCase):
    def testValidate(self):
        class MockOperator(TrainingOperator):
            def setup(self, config):
                self.train_epoch = MagicMock(returns=dict(mean_accuracy=10))
                self.validate = MagicMock(returns=dict(mean_accuracy=10))

-        runner = PyTorchRunner(
+        runner = TorchRunner(
            model_creator,
            create_dataloaders,
            optimizer_creator,
@@ -76,7 +76,7 @@ class TestPyTorchRunner(unittest.TestCase):
                self.count += 1
                return {"count": self.count}

-        runner = PyTorchRunner(
+        runner = TorchRunner(
            model_creator,
            create_dataloaders,
            optimizer_creator,
@@ -105,7 +105,7 @@ class TestPyTorchRunner(unittest.TestCase):
            ]
            return opts[0], opts[1], opts[2]

-        runner = PyTorchRunner(
+        runner = TorchRunner(
            three_model_creator,
            single_loader,
            three_optimizer_creator,
@@ -116,8 +116,8 @@ class TestPyTorchRunner(unittest.TestCase):
        self.assertEqual(len(runner.given_models), 3)
        self.assertEqual(len(runner.given_optimizers), 3)

-        runner2 = PyTorchRunner(model_creator, single_loader,
-                                optimizer_creator, loss_creator)
+        runner2 = TorchRunner(model_creator, single_loader, optimizer_creator,
+                              loss_creator)
        runner2.setup()

        self.assertNotEqual(runner2.given_models, runner2.models)
@@ -128,26 +128,26 @@ class TestPyTorchRunner(unittest.TestCase):
            return (LinearDataset(2, 5), LinearDataset(2, 5, size=400),
                    LinearDataset(2, 5, size=400))

-        runner = PyTorchRunner(model_creator, three_data_loader,
-                               optimizer_creator, loss_creator)
+        runner = TorchRunner(model_creator, three_data_loader,
+                             optimizer_creator, loss_creator)
        with self.assertRaises(ValueError):
            runner.setup()

-        runner2 = PyTorchRunner(model_creator, three_data_loader,
-                                optimizer_creator, loss_creator)
+        runner2 = TorchRunner(model_creator, three_data_loader,
+                              optimizer_creator, loss_creator)
        with self.assertRaises(ValueError):
            runner2.setup()

    def testSingleLoader(self):
-        runner = PyTorchRunner(model_creator, single_loader, optimizer_creator,
-                               loss_creator)
+        runner = TorchRunner(model_creator, single_loader, optimizer_creator,
+                             loss_creator)
        runner.setup()
        runner.train_epoch()
        with self.assertRaises(ValueError):
            runner.validate()

    def testNativeLoss(self):
-        runner = PyTorchRunner(
+        runner = TorchRunner(
            model_creator,
            single_loader,
            optimizer_creator,
@@ -165,8 +165,8 @@ class TestPyTorchRunner(unittest.TestCase):
            ]
            return opts[0], opts[1], opts[2]

-        runner = PyTorchRunner(multi_model_creator, single_loader,
-                               multi_optimizer_creator, loss_creator)
+        runner = TorchRunner(multi_model_creator, single_loader,
+                             multi_optimizer_creator, loss_creator)

        with self.assertRaises(ValueError):
            runner.setup()
@@ -0,0 +1,17 @@
+import logging
+logger = logging.getLogger(__name__)
+
+TorchTrainer = None
+TorchTrainable = None
+TrainingOperator = None
+
+try:
+    import torch  # noqa: F401
+
+    from ray.util.sgd.torch.torch_trainer import (TorchTrainer, TorchTrainable)
+
+    from ray.util.sgd.torch.training_operator import TrainingOperator
+
+    __all__ = ["TorchTrainer", "TorchTrainable", "TrainingOperator"]
+except ImportError:
+    logger.warning("PyTorch not found. TorchTrainer will not be available")
@@ -7,24 +7,24 @@ import torch.distributed as dist
 import torch.utils.data
 from torch.nn.parallel import DistributedDataParallel

-from ray.util.sgd.pytorch.pytorch_runner import PyTorchRunner
+from ray.util.sgd.torch.torch_runner import TorchRunner

 logger = logging.getLogger(__name__)


-class DistributedPyTorchRunner(PyTorchRunner):
+class DistributedTorchRunner(TorchRunner):
    """Manages a distributed PyTorch model replica.


    Args:
-        args: Arguments for PyTorchRunner.
+        args: Arguments for TorchRunner.
        backend (string): backend used by distributed PyTorch.
-        kwargs: Keyword arguments for PyTorchRunner.
+        kwargs: Keyword arguments for TorchRunner.

    """

    def __init__(self, *args, backend="gloo", **kwargs):
-        super(DistributedPyTorchRunner, self).__init__(*args, **kwargs)
+        super(DistributedTorchRunner, self).__init__(*args, **kwargs)
        self.backend = backend

    def setup(self, url, world_rank, world_size):
@@ -110,7 +110,7 @@ class DistributedPyTorchRunner(PyTorchRunner):
        """
        if hasattr(self.train_loader.sampler, "set_epoch"):
            self.train_loader.sampler.set_epoch(self.epochs)
-        return super(DistributedPyTorchRunner, self).train_epoch(**kwargs)
+        return super(DistributedTorchRunner, self).train_epoch(**kwargs)

    def _get_model_state_dicts(self):
        """Fetch state from ``model.module`` instead of ``model``.
@@ -132,7 +132,7 @@ class DistributedPyTorchRunner(PyTorchRunner):

    # def shutdown(self):
        """Attempts to shut down the worker."""
-        # super(DistributedPyTorchRunner, self).shutdown()
+        # super(DistributedTorchRunner, self).shutdown()
        # TODO: Temporarily removing since it causes hangs on MacOSX.
        # However, it seems to be harmless to remove permanently
        # since the processes are shutdown anyways. This comment can be
@@ -8,8 +8,8 @@ import torchvision
 import torchvision.transforms as transforms

 import ray
-from ray.util.sgd.pytorch import (PyTorchTrainer, PyTorchTrainable)
-from ray.util.sgd.pytorch.resnet import ResNet18
+from ray.util.sgd.torch import (TorchTrainer, TorchTrainable)
+from ray.util.sgd.torch.resnet import ResNet18


 def initialization_hook():
@@ -62,7 +62,7 @@ def train_example(num_replicas=1,
                  use_gpu=False,
                  use_fp16=False,
                  test_mode=False):
-    trainer1 = PyTorchTrainer(
+    trainer1 = TorchTrainer(
        ResNet18,
        cifar_creator,
        optimizer_creator,
@@ -107,7 +107,7 @@ def tune_example(num_replicas=1, use_gpu=False, test_mode=False):
    }

    analysis = tune.run(
-        PyTorchTrainable,
+        TorchTrainable,
        num_samples=2,
        config=config,
        stop={"training_iteration": 2},
@@ -15,9 +15,9 @@ from torch.nn import functional as F
 from scipy.stats import entropy

 import ray
-from ray.util.sgd import PyTorchTrainer
+from ray.util.sgd import TorchTrainer
 from ray.util.sgd.utils import override
-from ray.util.sgd.pytorch import TrainingOperator
+from ray.util.sgd.torch import TrainingOperator


 def data_creator(config):
@@ -223,9 +223,9 @@ def train_example(num_replicas=1, use_gpu=False, test_mode=False):
        "test_mode": test_mode,
        "classification_model_path": os.path.join(
            os.path.dirname(ray.__file__),
-            "util/sgd/pytorch/examples/mnist_cnn.pt")
+            "util/sgd/torch/examples/mnist_cnn.pt")
    }
-    trainer = PyTorchTrainer(
+    trainer = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -13,7 +13,7 @@ import numpy as np
 import torch
 import torch.nn as nn

-from ray.util.sgd import PyTorchTrainer
+from ray.util.sgd import TorchTrainer


 class LinearDataset(torch.utils.data.Dataset):
@@ -44,7 +44,7 @@ def optimizer_creator(model, config):
 def scheduler_creator(optimizer, config):
    """Returns a learning rate scheduler wrapping the optimizer.

-    You will need to set ``PyTorchTrainer(scheduler_step_freq="epoch")``
+    You will need to set ``TorchTrainer(scheduler_step_freq="epoch")``
    for the scheduler to be incremented correctly.

    If using a scheduler for validation loss, be sure to call
@@ -59,7 +59,7 @@ def data_creator(config):


 def train_example(num_replicas=1, use_gpu=False):
-    trainer1 = PyTorchTrainer(
+    trainer1 = TorchTrainer(
        model_creator,
        data_creator,
        optimizer_creator,
@@ -14,7 +14,7 @@ import torch.nn as nn

 import ray
 from ray import tune
-from ray.util.sgd.pytorch.pytorch_trainer import PyTorchTrainable
+from ray.util.sgd.torch.torch_trainer import TorchTrainable


 class LinearDataset(torch.utils.data.Dataset):
@@ -60,7 +60,7 @@ def tune_example(num_replicas=1, use_gpu=False):
    }

    analysis = tune.run(
-        PyTorchTrainable,
+        TorchTrainable,
        num_samples=12,
        config=config,
        stop={"training_iteration": 2},
@@ -9,8 +9,8 @@ import torch.utils.data
 from torch.utils.data import Dataset

 import ray
-from ray.util.sgd.pytorch.constants import USE_FP16, SCHEDULER_STEP
-from ray.util.sgd.pytorch.training_operator import TrainingOperator
+from ray.util.sgd.torch.constants import USE_FP16, SCHEDULER_STEP
+from ray.util.sgd.torch.training_operator import TrainingOperator
 from ray.util.sgd import utils

 logger = logging.getLogger(__name__)
@@ -23,23 +23,23 @@ except ImportError:
    pass


-class PyTorchRunner:
+class TorchRunner:
    """Manages a PyTorch model for training.

    Args:
-        model_creator (dict -> *): see pytorch_trainer.py
-        data_creator (dict -> Dataset, Dataset): see pytorch_trainer.py.
-        optimizer_creator (models, dict -> optimizers): see pytorch_trainer.py.
-        loss_creator (dict -> loss | Loss class): see pytorch_trainer.py.
+        model_creator (dict -> *): see torch_trainer.py
+        data_creator (dict -> Dataset, Dataset): see torch_trainer.py.
+        optimizer_creator (models, dict -> optimizers): see torch_trainer.py.
+        loss_creator (dict -> loss | Loss class): see torch_trainer.py.
        scheduler_creator (optimizers, dict -> schedulers): see
-            pytorch_trainer.py.
-        training_operator_cls: see pytorch_trainer.py
-        config (dict): see pytorch_trainer.py.
-        dataloader_config (dict): See pytorch_trainer.py.
-        batch_size (int): see pytorch_trainer.py.
-        use_fp16 (bool): see pytorch_trainer.py.
-        apex_args (dict|None): see pytorch_trainer.py.
-        scheduler_step_freq (str): see pytorch_trainer.py.
+            torch_trainer.py.
+        training_operator_cls: see torch_trainer.py
+        config (dict): see torch_trainer.py.
+        dataloader_config (dict): See torch_trainer.py.
+        batch_size (int): see torch_trainer.py.
+        use_fp16 (bool): see torch_trainer.py.
+        apex_args (dict|None): see torch_trainer.py.
+        scheduler_step_freq (str): see torch_trainer.py.
    """

    def __init__(self,
@@ -11,11 +11,11 @@ import ray

 from ray.tune import Trainable
 from ray.tune.trial import Resources
-from ray.util.sgd.pytorch.distributed_pytorch_runner import (
-    DistributedPyTorchRunner)
+from ray.util.sgd.torch.distributed_torch_runner import (
+    DistributedTorchRunner)
 from ray.util.sgd import utils
-from ray.util.sgd.pytorch.pytorch_runner import PyTorchRunner
-from ray.util.sgd.pytorch.constants import VALID_SCHEDULER_STEP
+from ray.util.sgd.torch.torch_runner import TorchRunner
+from ray.util.sgd.torch.constants import VALID_SCHEDULER_STEP

 logger = logging.getLogger(__name__)
 RESIZE_COOLDOWN_S = 10
@@ -29,7 +29,7 @@ def _validate_scheduler_step_freq(scheduler_step_freq):
                    VALID_SCHEDULER_STEP, scheduler_step_freq))


-class PyTorchTrainer:
+class TorchTrainer:
    """Train a PyTorch model using distributed PyTorch.

    Launches a set of actors which connect via distributed PyTorch and
@@ -49,7 +49,7 @@ class PyTorchTrainer:
        def data_creator(config):
            return LinearDataset(2, 5), LinearDataset(2, 5, size=400)

-        trainer = PyTorchTrainer(
+        trainer = TorchTrainer(
            model_creator,
            data_creator,
            optimizer_creator,
@@ -195,7 +195,7 @@ class PyTorchTrainer:
        if num_replicas == 1:
            # Generate actor class
            Runner = ray.remote(
-                num_cpus=1, num_gpus=int(self.use_gpu))(PyTorchRunner)
+                num_cpus=1, num_gpus=int(self.use_gpu))(TorchRunner)
            # Start workers
            self.workers = [
                Runner.remote(
@@ -220,8 +220,7 @@ class PyTorchTrainer:
        else:
            # Generate actor class
            Runner = ray.remote(
-                num_cpus=1,
-                num_gpus=int(self.use_gpu))(DistributedPyTorchRunner)
+                num_cpus=1, num_gpus=int(self.use_gpu))(DistributedTorchRunner)
            # Compute batch size per replica
            batch_size_per_replica = self.batch_size // num_replicas
            if self.batch_size % num_replicas > 0:
@@ -285,7 +284,7 @@ class PyTorchTrainer:
                in case of shared cluster usage.
            checkpoint (str): Path to checkpoint to restore from if retrying.
                If max_retries is set and ``checkpoint == "auto"``,
-                PyTorchTrainer will save a checkpoint before starting to train.
+                TorchTrainer will save a checkpoint before starting to train.
            info (dict): Optional dictionary passed to the training
                operator for ``train_epoch`` and ``train_batch``.

@@ -487,7 +486,7 @@ class PyTorchTrainer:
        return False


-class PyTorchTrainable(Trainable):
+class TorchTrainable(Trainable):
    @classmethod
    def default_resource_request(cls, config):
        return Resources(
@@ -497,7 +496,7 @@ class PyTorchTrainable(Trainable):
            extra_gpu=int(config["use_gpu"]) * config["num_replicas"])

    def _setup(self, config):
-        self._trainer = PyTorchTrainer(**config)
+        self._trainer = TorchTrainer(**config)

    def _train(self):
        train_stats = self._trainer.train()
@@ -2,7 +2,7 @@ import collections
 import torch

 from ray.util.sgd.utils import TimerStat, AverageMeter
-from ray.util.sgd.pytorch.constants import (
+from ray.util.sgd.torch.constants import (
    SCHEDULER_STEP_EPOCH, SCHEDULER_STEP_BATCH, SCHEDULER_STEP, BATCH_COUNT)

 amp = None
@@ -11,7 +11,7 @@ try:
    from apex import amp
 except ImportError:
    # Apex library is not installed, so we cannot enable mixed precision.
-    # We don't log here because logging happens in the pytorch_runner,
+    # We don't log here because logging happens in the torch_runner,
    # where amp is initialized.
    pass

@@ -26,7 +26,7 @@ class TrainingOperator:

    The scheduler will only be called at a batch or epoch frequency, depending
    on the user parameter. Be sure to set ``scheduler_step_freq`` in
-    ``PyTorchTrainer`` to either "batch" or "epoch" to increment the scheduler
+    ``TorchTrainer`` to either "batch" or "epoch" to increment the scheduler
    correctly during training. If using a learning rate scheduler
    that depends on validation loss, you can use ``trainer.update_scheduler``.

@@ -290,7 +290,7 @@ class TrainingOperator:

    @property
    def config(self):
-        """Dictionary as provided into PyTorchTrainer."""
+        """Dictionary as provided into TorchTrainer."""
        return self._config

    @property