wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-23 13:10:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

[raysgd] Cleanup User API (#7384)

Browse Source 
* Init fp16

* fp16 and schedulers

* scheduler linking and fp16

* to fp16

* loss scaling and documentation

* more documentation

* add tests, refactor config

* moredocs

* more docs

* fix logo, add test mode, add fp16 flag

* fix tests

* fix scheduler

* fix apex

* improve safety

* fix tests

* fix tests

* remove pin memory default

* rm

* fix

* Update doc/examples/doc_code/raysgd_torch_signatures.py

* fix

* migrate changes from other PR

* ok thanks

* pass

* signatures

* lint'

* Update python/ray/experimental/sgd/pytorch/utils.py

* Apply suggestions from code review

Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com>

* should address most comments

* comments

* fix this ci

* first_pass

* add overrides

* override

* fixing up operators

* format

* sgd

* constants

* rm

* revert

* save

* failures

* fixes

* trainer

* run test

* operator

* code

* op

* ok done

* operator

* sgd test fixes

* ok

* trainer

* format

* Apply suggestions from code review

Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com>

* Update doc/source/raysgd/raysgd_pytorch.rst

* docstring

* dcgan

* doc

* commits

* nit

* testing

* revert

* 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

* benchmarks

* rename

* remove some args

* better metrics output

* fix up the benchmark

* benchmark-yaml

* horovod-benchmark

* benchmarks

* Remove benchmark code for cleanups

* makedatacreator

* relax

* metrics

* autosetsampler

* profile

* movements

* OK

* smoothen

* fix

* nitdocs

* loss

* comments

* fix

* fix

* runner_tests

* codes

* example

* fix_test

* fix

* tests

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>
Co-authored-by: Maksim Smolin <maximsmol@gmail.com>
This commit is contained in:
Richard Liaw
2020-03-10 08:41:42 -07:00
committed by GitHub
co-authored by Edward Oakes Maksim Smolin
parent 89ec4adb72
commit d192ef0611
16 changed files with 840 additions and 478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/util/sgd/torch/torch_runner.py
+64 -70
View File
@@ -4,9 +4,8 @@ import logging
import inspect
import itertools
import os
import tempfile
import torch
import torch.utils.data
from torch.utils.data import Dataset
import ray
from ray.util.sgd.torch.constants import USE_FP16, SCHEDULER_STEP
@@ -27,16 +26,15 @@ class TorchRunner:
"""Manages a PyTorch model for training.
Args:
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
model_creator (dict -> Model(s)): see torch_trainer.py
data_creator (dict -> Iterable(s)): see torch_trainer.py.
optimizer_creator ((models, dict) -> optimizers): see torch_trainer.py.
loss_creator (torch.nn.*Loss class | dict -> loss):
see torch_trainer.py.
scheduler_creator ((optimizers, dict) -> scheduler): see
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.
@@ -46,36 +44,23 @@ class TorchRunner:
model_creator,
data_creator,
optimizer_creator,
loss_creator,
loss_creator=None,
scheduler_creator=None,
training_operator_cls=None,
config=None,
dataloader_config=None,
batch_size=16,
use_fp16=False,
apex_args=None,
scheduler_step_freq="batch"):
self.model_creator = model_creator
self.data_creator = data_creator
self.optimizer_creator = optimizer_creator
self.loss_creator = loss_creator
self.data_creator = data_creator
self.scheduler_creator = scheduler_creator
self.training_operator_cls = training_operator_cls or TrainingOperator
self.config = {} if config is None else config
self.dataloader_config = {
"num_workers": 2
} if dataloader_config is None else dataloader_config
self.batch_size = batch_size
self.verbose = True
self.timers = utils.TimerCollection()
self.epochs = 0
self._timers = {
k: utils.TimerStat(window_size=1)
for k in [
"setup_proc", "setup_model", "get_state", "set_state",
"validation", "training"
]
}
self.models = None
self.optimizers = None
self.criterion = None
@@ -90,23 +75,40 @@ class TorchRunner:
"https://www.github.com/nvidia/apex to use fp16 training.")
self.scheduler_step_freq = scheduler_step_freq
def _validate_datasets(self, dataset):
assert dataset, "Datasets need to be returned in data_creator."
if issubclass(type(dataset), Dataset):
return dataset, None
elif len(dataset) == 2 and issubclass(type(dataset[0]), Dataset):
return dataset
else:
raise ValueError("Datasets must be <= 2. Got {}".format(dataset))
def _validate_loaders(self, loaders):
assert loaders, "Loaders need to be returned in data_creator."
if isinstance(loaders, (tuple, list)):
if len(loaders) == 1:
return loaders, None
elif len(loaders) == 2:
return loaders
else:
raise ValueError(
"Number of loaders must be <= 2. Got {}".format(loaders))
# No great way of checking type otherwise
return loaders, None
def _initialize_dataloaders(self):
logger.debug("Instantiating dataloaders.")
# When creating loaders, a filelock will be used to ensure no
# race conditions in data downloading among different workers.
with FileLock(os.path.join(tempfile.gettempdir(), ".ray_data.lock")):
loaders = self.data_creator(self.config)
train_loader, val_loader = self._validate_loaders(loaders)
self.train_loader, self.validation_loader = train_loader, val_loader
def _create_loss(self):
if not self.loss_creator:
return
logger.debug("Creating loss.")
if inspect.isclass(self.loss_creator) and issubclass(
self.loss_creator, torch.nn.modules.loss._Loss):
self.criterion = self.loss_creator()
else:
self.criterion = self.loss_creator(self.config)
if torch.cuda.is_available():
if torch.cuda.is_available() and hasattr("cuda", self.criterion):
self.criterion = self.criterion.cuda()
def _create_schedulers_if_available(self):
@@ -142,22 +144,7 @@ class TorchRunner:
self._create_schedulers_if_available()
self._try_setup_apex()
self._create_loss()
logger.debug("Creating dataset")
# When creating datasets, a filelock will be used to ensure no
# race conditions in data downloading among different workers.
with FileLock(os.path.expanduser("~/.ray_data.lock")):
datasets = self.data_creator(self.config)
train_set, val_set = self._validate_datasets(datasets)
self.train_loader = torch.utils.data.DataLoader(
train_set, batch_size=self.batch_size, **self.dataloader_config)
self.validation_loader = None
if val_set:
self.validation_loader = torch.utils.data.DataLoader(
val_set, batch_size=self.batch_size, **self.dataloader_config)
self._initialize_dataloaders()
self.training_operator = self.training_operator_cls(
self.config,
models=self.models,
@@ -174,47 +161,55 @@ class TorchRunner:
"""Finds a free port on the current node."""
return utils.find_free_port()
def train_epoch(self, num_steps=None, info=None):
def train_epoch(self, num_steps=None, profile=False, info=None):
"""Runs a training epoch and updates the model parameters."""
logger.debug("Begin Training Step {}".format(self.epochs + 1))
info = info or {}
self._toggle_profiling(profile=profile)
info.update({
USE_FP16: self.use_fp16,
SCHEDULER_STEP: self.scheduler_step_freq
})
with self._timers["training"]:
with self.timers.record("train_epoch"):
iterator = self.train_loader
if num_steps:
iterator = itertools.islice(iter(self.train_loader), num_steps)
train_stats = self.training_operator.train_epoch(iterator, info)
self.epochs += 1
train_stats.update(self.stats())
return train_stats
# This is so that `epochs` is first in ordering.
stats = dict(epoch=self.epochs, **train_stats)
if profile:
stats.update(profile=self.timers.stats())
return stats
def validate(self, num_steps=None, info=None):
def validate(self, num_steps=None, profile=False, info=None):
"""Evaluates the model on the validation data set."""
if self.validation_loader is None:
raise ValueError("No validation dataloader provided.")
info = info or {}
with self._timers["validation"]:
self._toggle_profiling(profile=profile)
with self.timers.record("validation"):
iterator = self.validation_loader
if num_steps:
iterator = itertools.islice(
iter(self.validation_loader), num_steps)
validation_stats = self.training_operator.validate(iterator, info)
validation_stats.update(self.stats())
validation_stats = self.training_operator.validate(
iterator, info=info)
if profile:
validation_stats.update(profile=self.timers.stats())
return validation_stats
def stats(self):
"""Returns a dictionary of statistics collected."""
stats = {"epoch": self.epochs}
for k, t in self._timers.items():
stats[k + "_time_mean"] = t.mean
stats[k + "_time_total"] = t.sum
t.reset()
return stats
def _toggle_profiling(self, profile=False):
"""Enables/Disables and resets timing profiles."""
if profile:
self.timers.enable()
self.timers.reset()
else:
self.timers.disable()
self.training_operator._set_timers(self.timers)
def _get_model_state_dicts(self):
# This is so that we create a duplicate of weights into CPU rather than
@@ -237,8 +232,7 @@ class TorchRunner:
"epoch": self.epochs,
"operator": self.training_operator.state_dict(),
"models": self._get_model_state_dicts(),
"optimizers": [opt.state_dict() for opt in self.optimizers],
"stats": self.stats()
"optimizers": [opt.state_dict() for opt in self.optimizers]
}
if self.schedulers:
state.update({
@@ -264,7 +258,7 @@ class TorchRunner:
if self.use_fp16 and "amp" in state and amp:
amp.load_state_dict(state["amp"])
self.epochs = state["stats"]["epoch"]
self.epochs = state["epoch"]
self.training_operator.load_state_dict(state_dict)
def apply(self, fn):