ray-project · pcmoritz · Sep 20, 2018 · May 15, 2018 · May 15, 2018 · May 16, 2018
diff --git a/python/ray/experimental/sgd/__init__.py b/python/ray/experimental/sgd/__init__.py
diff --git a/python/ray/experimental/sgd/modified_allreduce.py b/python/ray/experimental/sgd/modified_allreduce.py
diff --git a/python/ray/experimental/sgd/sgd.py b/python/ray/experimental/sgd/sgd.py
diff --git a/python/ray/experimental/sgd/test_sgd.py b/python/ray/experimental/sgd/test_sgd.py
@@ -0,0 +1,29 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import ray
+
+import argparse
+import numpy as np
+import tensorflow as tf
+
+from ray.experimental.sgd.tfbench.test_model import TFBenchModel
+from ray.experimental.sgd.sgd import DistributedSGD
+
+if __name__ == "__main__":
+    ray.init()
+
+    model_creator = (
+        lambda worker_idx, device_idx: TFBenchModel(batch=1, use_cpus=True))
+
+    sgd = DistributedSGD(
+        model_creator,
+        num_workers=2,
+        devices_per_worker=2,
+        use_cpus=True,
+        use_plasma_op=False)
+
+    for _ in range(100):
+        loss = sgd.step()
+        print("Current loss", loss)
diff --git a/python/ray/experimental/sgd/tfbench/__init__.py b/python/ray/experimental/sgd/tfbench/__init__.py