updated grph

Author: stephenoffer

.idea/workspace.xml

@@ -7,7 +7,6 @@
 
 import os
 import time
-import keras
 
 import ray
 
@@ -139,12 +138,15 @@ def make_model():
                 else:
                     time.sleep(1e-3)
 
-    def predict(self, data, distributed=True, stop_ray=False):
+    def predict(self, data, distributed=True, close=False):
         """
+        Run inference on a data batch, returns predictions
 
-        :param data:
-        :param distributed:
-        :return:
+        Arguments:
+            data: numpy array of images
+            distributed: True for distributed inference, false for serial
+            close: boolean value for whether to stop workers
+        return: Predictions
         """
         if distributed:
             n_data = len(data)
@@ -161,14 +163,27 @@ def predict(self, data, distributed=True, stop_ray=False):
             print("Completed!")
         else:
             return self.model.predict(data)
-        if stop_ray:
-            ray.shutdown()
+        if close:
+            self.close()
 
     def close(self, stop_ray=False):
+        """
+        Close the Ray workers for the model
+
+        Arguments:
+            stop_ray: Boolean value for whether to close Ray cluster
+
+        return: None
+        """
         self.data_server.close.remote()
         if stop_ray:
             ray.shutdown()
         time.sleep(5e-2)
 
     def is_ready(self):
+        """
+        Wait for workers to initialize
+
+        :return: True
+        """
         return ray.get(self.data_server.all_ready.remote())

assets/dkeras_design.png

@@ -11,13 +11,15 @@
 
 @ray.remote
 class DataServer(object):
+    """
+    Data server for transferring data to workers
 
-    def __init__(self, n_workers, worker_ids):
-        """
+    Arguments:
+        n_workers: Number of workers
+        worker_ids: List of worker IDs
+    """
 
-        :param n_workers:
-        :param worker_ids:
-        """
+    def __init__(self, n_workers, worker_ids):
         self.n_workers = n_workers
         self.id_indexes = {}
         self.id = 0
@@ -32,8 +34,9 @@ def __init__(self, n_workers, worker_ids):
 
     def pull_results(self):
         """
+        Pull the prediction results
 
-        :return:
+        :return: Numpy array with the output of the workers
         """
         output = []
         for n in self.results:
@@ -45,8 +48,10 @@ def pull_results(self):
 
     def close(self):
         """
+        Set the close flag to True, which will close the workers
+        on the next worker data pull request
 
-        :return:
+        :return: None
         """
         self.closed = True
 

assets/inference_comparison.png

@@ -14,6 +14,7 @@
 @ray.remote(num_cpus=cpus_per_worker, num_gpus=gpus_per_worker)
 def worker_task(weights, ds, make_model):
     """
+    Worker function loop
 
     :param weights:
     :param ds: