mirror of
https://github.com/vale981/ray
synced 2025-03-10 13:26:39 -04:00
26 lines
677 B
Python
26 lines
677 B
Python
|
import json
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
|
||
|
|
import ray
|
||
|
|
from ray.util.sgd.torch.examples.cifar_pytorch_example import train_cifar
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
ray.init(address=os.environ.get("RAY_ADDRESS", "auto"))
|
||
|
|
start_time = time.time()
|
||
|
|
success = True
|
||
|
|
try:
|
||
|
|
train_cifar(
|
||
|
|
num_workers=2,
|
||
|
|
use_gpu=True,
|
||
|
|
num_epochs=5,
|
||
|
|
fp16=True,
|
||
|
|
test_mode=False)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"The test failed with {e}")
|
||
|
|
success = False
|
||
|
|
|
||
|
|
delta = time.time() - start_time
|
||
|
|
with open(os.environ["TEST_OUTPUT_JSON"], "w") as f:
|
||
|
|
f.write(json.dumps({"train_time": delta, "success": success}))
|