def test_scheduler_validate(ray_start_2_cpus): # noqa: F811
from torch.optim.lr_scheduler import ReduceLROnPlateau
TestOperator = TrainingOperator.from_creators(
model_creator,
optimizer_creator,
data_creator,
scheduler_creator=lambda optimizer, cfg: ReduceLROnPlateau(optimizer),
loss_creator=lambda config: nn.MSELoss())
TestOperator = get_test_operator(TestOperator)
trainer = TorchTrainer(scheduler_step_freq="manual",
training_operator_cls=TestOperator)
trainer.update_scheduler(0.5)
trainer.update_scheduler(0.5)
assert all(
trainer.apply_all_operators(
lambda op: op._schedulers[0].last_epoch == 2))
trainer.shutdown()
def test_resize(ray_start_2_cpus, use_local): # noqa: F811
if not dist.is_available():
return
def single_loader(config):
dataset = LinearDataset(2, 5, size=1000000)
return DataLoader(dataset, batch_size=config.get("batch_size", 32))
start_with_fail = gen_start_with_fail(1)
TestOperator = TrainingOperator.from_creators(
model_creator,
optimizer_creator,
single_loader,
loss_creator=lambda config: nn.MSELoss())
with patch.object(TorchTrainer, "_start_workers", start_with_fail):
trainer1 = TorchTrainer(training_operator_cls=TestOperator,
config={"batch_size": 100000},
use_local=use_local,
num_workers=2)
@ray.remote(num_cpus=1)
class DummyActor:
def get(self):
return 1
dummy_handler = DummyActor.remote()
trainer1.train(max_retries=1)
assert trainer1.worker_group.num_workers == 1
assert trainer1._num_failures == 1
ray.get(dummy_handler.get.remote())
ray.kill(dummy_handler)
time.sleep(1)
# trigger scale up
trainer1.train()
assert trainer1.worker_group.num_workers == 2
trainer1.shutdown(force=True)
def testMultiLoaders(self):
def three_data_loader(config):
return (
LinearDataset(2, 5),
LinearDataset(2, 5, size=400),
LinearDataset(2, 5, size=400),
)
ThreeOperator = TrainingOperator.from_creators(
model_creator,
optimizer_creator,
three_data_loader,
loss_creator=loss_creator,
)
runner = TorchRunner(training_operator_cls=ThreeOperator)
with self.assertRaises(ValueError):
runner.setup_operator()
runner2 = TorchRunner(training_operator_cls=ThreeOperator)
with self.assertRaises(ValueError):
runner2.setup_operator()
def test_fail_twice(ray_start_2_cpus): # noqa: F811
if not dist.is_available():
return
def single_loader(config):
dataset = LinearDataset(2, 5, size=1000000)
return DataLoader(dataset, batch_size=config.get("batch_size", 32))
step_with_fail = gen_step_with_fail(2)
TestOperator = TrainingOperator.from_creators(
model_creator,
optimizer_creator,
single_loader,
loss_creator=lambda config: nn.MSELoss())
with patch.object(TorchTrainer, "_train_epoch", step_with_fail):
trainer1 = TorchTrainer(training_operator_cls=TestOperator,
config={"batch_size": 100000},
num_workers=2)
# MAX RETRIES SHOULD BE ON BY DEFAULT
trainer1.train()
trainer1.shutdown()
def test_fail_with_recover(ray_start_2_cpus): # noqa: F811
if not dist.is_available():
return
def single_loader(config):
dataset = LinearDataset(2, 5, size=1000000)
return DataLoader(dataset, batch_size=config.get("batch_size", 32))
step_with_fail = gen_step_with_fail(3)
TestOperator = TrainingOperator.from_creators(
model_creator,
optimizer_creator,
single_loader,
loss_creator=lambda config: nn.MSELoss())
with patch.object(TorchTrainer, "_train_epoch", step_with_fail):
trainer1 = TorchTrainer(training_operator_cls=TestOperator,
config={"batch_size": 100000},
num_workers=2)
with pytest.raises(RuntimeError):
trainer1.train(max_retries=1)
trainer1.shutdown(force=True)
import torch.nn as nn
import torch.distributed as dist
from torch.utils.data import DataLoader
import ray
from ray.util.sgd.torch import TorchTrainer
from ray.util.sgd.torch.worker_group import RemoteWorkerGroup
from ray.util.sgd.torch.training_operator import TrainingOperator
from ray.util.sgd.torch.examples.train_example import (model_creator,
optimizer_creator,
data_creator,
LinearDataset)
Operator = TrainingOperator.from_creators(model_creator,
optimizer_creator,
data_creator,
loss_creator=nn.MSELoss)
@pytest.fixture
def ray_start_2_cpus():
address_info = ray.init(num_cpus=2)
yield address_info
# The code after the yield will run as teardown code.
ray.shutdown()
# Ensure that tests don't ALL fail
if dist.is_initialized():
dist.destroy_process_group()
@pytest.fixture
def test_num_steps(ray_start_2_cpus, use_local):
"""Tests if num_steps continues training from the subsampled dataset."""
def data_creator(config):
train_dataset = [0] * 5 + [1] * 5
val_dataset = [0] * 5 + [1] * 5
return DataLoader(train_dataset, batch_size=config["batch_size"]), \
DataLoader(val_dataset, batch_size=config["batch_size"])
batch_size = 1
Operator = TrainingOperator.from_creators(model_creator, optimizer_creator,
data_creator)
def train_func(self, iterator, info=None):
total_sum = 0
num_items = 0
for e in iterator:
total_sum += e
num_items += 1
return {"average": total_sum.item() / num_items}
TestOperator = get_test_operator(Operator)
trainer = TorchTrainer(
training_operator_cls=TestOperator,
num_workers=2,
use_local=use_local,
add_dist_sampler=False,
config={
"batch_size": batch_size,
"custom_func": train_func
})
# If num_steps not passed, should do one full epoch.
result = trainer.train()
# Average of 5 0s and 5 1s
assert result["average"] == 0.5
assert result["epoch"] == 1
val_result = trainer.validate()
assert val_result["average"] == 0.5
# Train again with num_steps.
result = trainer.train(num_steps=5)
# 5 zeros
assert result["average"] == 0
assert result["epoch"] == 2
val_result = trainer.validate(num_steps=5)
assert val_result["average"] == 0
# Should continue where last train run left off.
result = trainer.train(num_steps=3)
# 3 ones.
assert result["average"] == 1
assert result["epoch"] == 2
val_result = trainer.validate(num_steps=3)
assert val_result["average"] == 1
# Should continue from last train run, and cycle to beginning.
result = trainer.train(num_steps=5)
# 2 ones and 3 zeros.
assert result["average"] == 0.4
assert result["epoch"] == 3
val_result = trainer.validate(num_steps=5)
assert val_result["average"] == 0.4
# Should continue, and since num_steps not passed in, just finishes epoch.
result = trainer.train()
# 2 zeros and 5 ones.
assert result["average"] == 5 / 7
assert result["epoch"] == 3
val_result = trainer.validate()
assert val_result["average"] == 5 / 7
trainer.shutdown()
def setUp(self):
self.Operator = TrainingOperator.from_creators(
model_creator,
optimizer_creator,
create_dataloaders,
loss_creator=loss_creator)