parallel_model code update for keras 2.2.0 compatibility

mrcnn/parallel_model.py

@@ -32,11 +32,11 @@ def __init__(self, keras_model, gpu_count):
         keras_model: The Keras model to parallelize
         gpu_count: Number of GPUs. Must be > 1
         """
+        merged_outputs = self.make_parallel(keras_model, gpu_count)
+        super(ParallelModel, self).__init__(inputs=keras_model.inputs,
+                                            outputs=merged_outputs)
         self.inner_model = keras_model
         self.gpu_count = gpu_count
-        merged_outputs = self.make_parallel()
-        super(ParallelModel, self).__init__(inputs=self.inner_model.inputs,
-                                            outputs=merged_outputs)
 
     def __getattribute__(self, attrname):
         """Redirect loading and saving methods to the inner model. That's where
@@ -51,34 +51,34 @@ def summary(self, *args, **kwargs):
         super(ParallelModel, self).summary(*args, **kwargs)
         self.inner_model.summary(*args, **kwargs)
 
-    def make_parallel(self):
+    def make_parallel(self, inner_model,gpu_count):
         """Creates a new wrapper model that consists of multiple replicas of
         the original model placed on different GPUs.
         """
         # Slice inputs. Slice inputs on the CPU to avoid sending a copy
         # of the full inputs to all GPUs. Saves on bandwidth and memory.
-        input_slices = {name: tf.split(x, self.gpu_count)
-                        for name, x in zip(self.inner_model.input_names,
-                                           self.inner_model.inputs)}
+        input_slices = {name: tf.split(x, gpu_count)
+                        for name, x in zip(inner_model.input_names,
+                                           inner_model.inputs)}
 
-        output_names = self.inner_model.output_names
+        output_names = inner_model.output_names
         outputs_all = []
-        for i in range(len(self.inner_model.outputs)):
+        for i in range(len(inner_model.outputs)):
             outputs_all.append([])
 
         # Run the model call() on each GPU to place the ops there
-        for i in range(self.gpu_count):
+        for i in range(gpu_count):
             with tf.device('/gpu:%d' % i):
                 with tf.name_scope('tower_%d' % i):
                     # Run a slice of inputs through this replica
-                    zipped_inputs = zip(self.inner_model.input_names,
-                                        self.inner_model.inputs)
+                    zipped_inputs = zip(inner_model.input_names,
+                                        inner_model.inputs)
                     inputs = [
                         KL.Lambda(lambda s: input_slices[name][i],
                                   output_shape=lambda s: (None,) + s[1:])(tensor)
                         for name, tensor in zipped_inputs]
                     # Create the model replica and get the outputs
-                    outputs = self.inner_model(inputs)
+                    outputs = inner_model(inputs)
                     if not isinstance(outputs, list):
                         outputs = [outputs]
                     # Save the outputs for merging back together later