mirror of
https://github.com/wassname/ray.git
synced 2026-07-15 11:25:40 +08:00
Richard LiawandGitHub
037aa2b961
* Refactor documentation and directory structurre * update loss * ,ore examples * fix comments * more code * svgs * formatting * more_docs * more writing * comments ready * move * whitespace * examples * fix * bold * pytorch * batch * fix * fix test * Apply suggestions from code review * quarantinegp * tests/ * fix missing
132 lines
3.4 KiB
Python
132 lines
3.4 KiB
Python
import os
|
|
import pytest
|
|
import tempfile
|
|
import numpy as np
|
|
import shutil
|
|
|
|
from ray import tune
|
|
from ray.tests.conftest import ray_start_2_cpus # noqa: F401
|
|
from ray.experimental.sgd.tf import TFTrainer, TFTrainable
|
|
|
|
from ray.experimental.sgd.tf.examples.tensorflow_train_example import (
|
|
simple_model, simple_dataset)
|
|
|
|
SIMPLE_CONFIG = {
|
|
"batch_size": 128,
|
|
"fit_config": {
|
|
"steps_per_epoch": 3,
|
|
},
|
|
"evaluate_config": {
|
|
"steps": 3,
|
|
}
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize( # noqa: F811
|
|
"num_replicas", [1, 2])
|
|
def test_train(ray_start_2_cpus, num_replicas): # noqa: F811
|
|
trainer = TFTrainer(
|
|
model_creator=simple_model,
|
|
data_creator=simple_dataset,
|
|
num_replicas=num_replicas,
|
|
config=SIMPLE_CONFIG)
|
|
|
|
train_stats1 = trainer.train()
|
|
train_stats1.update(trainer.validate())
|
|
|
|
train_stats2 = trainer.train()
|
|
train_stats2.update(trainer.validate())
|
|
|
|
|
|
@pytest.mark.parametrize( # noqa: F811
|
|
"num_replicas", [1, 2])
|
|
def test_tune_train(ray_start_2_cpus, num_replicas): # noqa: F811
|
|
|
|
config = {
|
|
"model_creator": tune.function(simple_model),
|
|
"data_creator": tune.function(simple_dataset),
|
|
"num_replicas": num_replicas,
|
|
"use_gpu": False,
|
|
"trainer_config": SIMPLE_CONFIG
|
|
}
|
|
|
|
tune.run(
|
|
TFTrainable,
|
|
num_samples=2,
|
|
config=config,
|
|
stop={"training_iteration": 2},
|
|
verbose=1)
|
|
|
|
|
|
@pytest.mark.parametrize( # noqa: F811
|
|
"num_replicas", [1, 2])
|
|
def test_save_and_restore(ray_start_2_cpus, num_replicas): # noqa: F811
|
|
trainer1 = TFTrainer(
|
|
model_creator=simple_model,
|
|
data_creator=simple_dataset,
|
|
num_replicas=num_replicas,
|
|
config=SIMPLE_CONFIG)
|
|
trainer1.train()
|
|
|
|
tmpdir = tempfile.mkdtemp()
|
|
filename = os.path.join(tmpdir, "checkpoint")
|
|
trainer1.save(filename)
|
|
|
|
model1 = trainer1.get_model()
|
|
trainer1.shutdown()
|
|
|
|
trainer2 = TFTrainer(
|
|
model_creator=simple_model,
|
|
data_creator=simple_dataset,
|
|
num_replicas=num_replicas,
|
|
config=SIMPLE_CONFIG)
|
|
trainer2.restore(filename)
|
|
|
|
model2 = trainer2.get_model()
|
|
trainer2.shutdown()
|
|
|
|
shutil.rmtree(tmpdir)
|
|
|
|
model1_config = model1.get_config()
|
|
model2_config = model2.get_config()
|
|
assert _compare(model1_config, model2_config, skip_keys=["name"])
|
|
|
|
model1_weights = model1.get_weights()
|
|
model2_weights = model2.get_weights()
|
|
assert _compare(model1_weights, model2_weights)
|
|
|
|
model1_opt_weights = model1.optimizer.get_weights()
|
|
model2_opt_weights = model2.optimizer.get_weights()
|
|
assert _compare(model1_opt_weights, model2_opt_weights)
|
|
|
|
|
|
def _compare(d1, d2, skip_keys=None):
|
|
"""Compare two lists or dictionaries or array"""
|
|
if type(d1) != type(d2):
|
|
return False
|
|
|
|
if isinstance(d1, dict):
|
|
if set(d1) != set(d2):
|
|
return False
|
|
|
|
for key in d1:
|
|
if skip_keys is not None and key in skip_keys:
|
|
continue
|
|
|
|
if not _compare(d1[key], d2[key], skip_keys=skip_keys):
|
|
return False
|
|
|
|
elif isinstance(d1, list):
|
|
for i, _ in enumerate(d1):
|
|
if not _compare(d1[i], d2[i], skip_keys=skip_keys):
|
|
return False
|
|
|
|
elif isinstance(d1, np.ndarray):
|
|
if not np.array_equal(d1, d2):
|
|
return False
|
|
else:
|
|
if d1 != d2:
|
|
return False
|
|
|
|
return True
|