def test_predict(tmpdir, ray_start_2_cpus, seed, num_workers):
"""Tests if trained model has high accuracy on test set."""
config = {
"layer_1": 32,
"layer_2": 32,
"lr": 1e-2,
"batch_size": 32,
}
model = LightningMNISTClassifier(config, tmpdir)
dm = MNISTDataModule(
data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
plugin = RayPlugin(num_workers=num_workers, use_gpu=False)
trainer = get_trainer(
tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
predict_test(trainer, model, dm)
def test_predict_client(tmpdir, start_ray_client_server_2_cpus, seed,
num_slots):
assert ray.util.client.ray.is_connected()
config = {
"layer_1": 32,
"layer_2": 32,
"lr": 1e-2,
"batch_size": 32,
}
model = LightningMNISTClassifier(config, tmpdir)
dm = MNISTDataModule(data_dir=tmpdir,
num_workers=1,
batch_size=config["batch_size"])
plugin = HorovodRayPlugin(num_slots=num_slots, use_gpu=False)
trainer = get_trainer(tmpdir,
limit_train_batches=20,
max_epochs=1,
plugins=[plugin])
predict_test(trainer, model, dm)
def test_predict_gpu(tmpdir, ray_start_2_gpus, seed, num_slots):
"""Tests if trained model has high accuracy on test set."""
config = {
"layer_1": 32,
"layer_2": 32,
"lr": 1e-2,
"batch_size": 32,
}
model = LightningMNISTClassifier(config, tmpdir)
dm = MNISTDataModule(data_dir=tmpdir,
num_workers=1,
batch_size=config["batch_size"])
accelerator = HorovodRayAccelerator(num_slots=num_slots, use_gpu=True)
trainer = get_trainer(tmpdir,
limit_train_batches=10,
max_epochs=1,
accelerator=accelerator,
use_gpu=True)
predict_test(trainer, model, dm)