wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-18 12:20:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[raysgd] Improve raysgd examples (#7818)

Browse Source 
* better_example

* test

* improve some usability things

* submit

* fix

* flake

* Update python/ray/util/sgd/torch/training_operator.py

* trythis

* fix

* fix

* smoke

* fail

* fix

* fix
This commit is contained in:
Richard Liaw
2020-04-01 08:58:39 -07:00
committed by GitHub
parent f4239d27fa
commit 24bf6ad607
11 changed files with 346 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/util/sgd/torch/torch_runner.py
+18 -4
View File
@@ -23,6 +23,14 @@ except ImportError:
pass
def _remind_gpu_usage(use_gpu, is_chief):
if not is_chief:
return
if not use_gpu and torch.cuda.is_available():
logger.info("GPUs detected but not using them. Set `use_gpu` to "
"enable GPU usage. ")
class TorchRunner:
"""Manages a PyTorch model for training.
@@ -36,6 +44,7 @@ class TorchRunner:
torch_trainer.py.
training_operator_cls: see torch_trainer.py
config (dict): see torch_trainer.py.
use_gpu (bool): 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.
@@ -49,6 +58,7 @@ class TorchRunner:
scheduler_creator=None,
training_operator_cls=None,
config=None,
use_gpu=False,
use_fp16=False,
use_tqdm=False,
apex_args=None,
@@ -69,6 +79,7 @@ class TorchRunner:
self.schedulers = None
self.train_loader = None
self.validation_loader = None
self.use_gpu = use_gpu
self.use_fp16 = use_fp16
self.use_tqdm = use_tqdm
self.apex_args = apex_args or {}
@@ -117,8 +128,9 @@ class TorchRunner:
else:
self.criterion = self.loss_creator(self.config)
if torch.cuda.is_available() and hasattr(self.criterion, "cuda"):
self.criterion = self.criterion.cuda()
if self.use_gpu and torch.cuda.is_available():
if hasattr(self.criterion, "cuda"):
self.criterion = self.criterion.cuda()
def _create_schedulers_if_available(self):
# Learning rate schedules are optional.
@@ -138,11 +150,13 @@ class TorchRunner:
def setup(self):
"""Initializes the model."""
_remind_gpu_usage(self.use_gpu, is_chief=True)
self._initialize_dataloaders()
logger.debug("Creating model")
self.models = self.model_creator(self.config)
if not isinstance(self.models, collections.Iterable):
self.models = [self.models]
if torch.cuda.is_available():
if self.use_gpu and torch.cuda.is_available():
self.models = [model.cuda() for model in self.models]
logger.debug("Creating optimizer")
@@ -153,7 +167,6 @@ class TorchRunner:
self._create_schedulers_if_available()
self._try_setup_apex()
self._create_loss()
self._initialize_dataloaders()
self.training_operator = self.training_operator_cls(
self.config,
models=self.models,
@@ -163,6 +176,7 @@ class TorchRunner:
validation_loader=self.validation_loader,
world_rank=0,
schedulers=self.schedulers,
use_gpu=self.use_gpu,
use_fp16=self.use_fp16,
use_tqdm=self.use_tqdm)