__init__.py 1.2 KB

123456789101112131415161718192021222324252627282930
  1. """
  2. :mod:`torch.distributed.optim` exposes DistributedOptimizer, which takes a list
  3. of remote parameters (:class:`~torch.distributed.rpc.RRef`) and runs the
  4. optimizer locally on the workers where the parameters live. The distributed
  5. optimizer can use any of the local optimizer :ref:`optimizer-algorithms` to
  6. apply the gradients on each worker.
  7. """
  8. import torch
  9. from torch import optim
  10. from .apply_optimizer_in_backward import _apply_optimizer_in_backward
  11. from .functional_adadelta import _FunctionalAdadelta
  12. from .functional_adagrad import _FunctionalAdagrad
  13. from .functional_adam import _FunctionalAdam
  14. from .functional_adamax import _FunctionalAdamax
  15. from .functional_adamw import _FunctionalAdamW
  16. from .functional_rmsprop import _FunctionalRMSprop
  17. from .functional_rprop import _FunctionalRprop
  18. from .functional_sgd import _FunctionalSGD
  19. from .named_optimizer import _NamedOptimizer
  20. from .utils import as_functional_optim
  21. # DistributedOptimizer imports torch.distributed.rpc names, so gate availability
  22. # based on RPC being available.
  23. if hasattr(torch._C, "_rpc_init"):
  24. from .optimizer import DistributedOptimizer
  25. from .post_localSGD_optimizer import PostLocalSGDOptimizer
  26. from .zero_redundancy_optimizer import ZeroRedundancyOptimizer