""" Mask R-CNN Multi-GPU Support for Keras. Copyright (c) 2017 Matterport, Inc. Licensed under the MIT License (see LICENSE for details) Written by Waleed Abdulla Ideas and a small code snippets from these sources: https://github.com/fchollet/keras/issues/2436 https://medium.com/@kuza55/transparent-multi-gpu-training-on-tensorflow-with-keras-8b0016fd9012 https://github.com/avolkov1/keras_experiments/blob/master/keras_exp/multigpu/ https://github.com/fchollet/keras/blob/master/keras/utils/training_utils.py """ import tensorflow as tf import keras.backend as K import keras.layers as KL import keras.models as KM class ParallelModel(KM.Model): """Subclasses the standard Keras Model and adds multi-GPU support. It works by creating a copy of the model on each GPU. Then it slices the inputs and sends a slice to each copy of the model, and then merges the outputs together and applies the loss on the combined outputs. """ def __init__(self, keras_model, gpu_count): """Class constructor. keras_model: The Keras model to parallelize gpu_count: Number of GPUs. Must be > 1 """ 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 the weights are stored.""" if "load" in attrname or "save" in attrname: return getattr(self.inner_model, attrname) return super(ParallelModel, self).__getattribute__(attrname) def summary(self, *args, **kwargs): """Override summary() to display summaries of both, the wrapper and inner models.""" super(ParallelModel, self).summary(*args, **kwargs) self.inner_model.summary(*args, **kwargs) def make_parallel(self): """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) } output_names = self.inner_model.output_names outputs_all = [] for i in range(len(self.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): 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 ) 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) if not isinstance(outputs, list): outputs = [outputs] # Save the outputs for merging back together later for l, o in enumerate(outputs): outputs_all[l].append(o) # Merge outputs on CPU with tf.device("/cpu:0"): merged = [] for outputs, name in zip(outputs_all, output_names): # Concatenate or average outputs? # Outputs usually have a batch dimension and we concatenate # across it. If they don't, then the output is likely a loss # or a metric value that gets averaged across the batch. # Keras expects losses and metrics to be scalars. if K.int_shape(outputs[0]) == (): # Average m = KL.Lambda(lambda o: tf.add_n(o) / len(outputs), name=name)( outputs ) else: # Concatenate m = KL.Concatenate(axis=0, name=name)(outputs) merged.append(m) return merged if __name__ == "__main__": # Testing code below. It creates a simple model to train on MNIST and # tries to run it on 2 GPUs. It saves the graph so it can be viewed # in TensorBoard. Run it as: # # python3 parallel_model.py import os import numpy as np import keras.optimizers from keras.datasets import mnist from keras.preprocessing.image import ImageDataGenerator GPU_COUNT = 2 # Root directory of the project ROOT_DIR = os.path.abspath("../") # Directory to save logs and trained model MODEL_DIR = os.path.join(ROOT_DIR, "logs") def build_model(x_train, num_classes): # Reset default graph. Keras leaves old ops in the graph, # which are ignored for execution but clutter graph # visualization in TensorBoard. tf.reset_default_graph() inputs = KL.Input(shape=x_train.shape[1:], name="input_image") x = KL.Conv2D(32, (3, 3), activation="relu", padding="same", name="conv1")( inputs ) x = KL.Conv2D(64, (3, 3), activation="relu", padding="same", name="conv2")(x) x = KL.MaxPooling2D(pool_size=(2, 2), name="pool1")(x) x = KL.Flatten(name="flat1")(x) x = KL.Dense(128, activation="relu", name="dense1")(x) x = KL.Dense(num_classes, activation="softmax", name="dense2")(x) return KM.Model(inputs, x, "digit_classifier_model") # Load MNIST Data (x_train, y_train), (x_test, y_test) = mnist.load_data() x_train = np.expand_dims(x_train, -1).astype("float32") / 255 x_test = np.expand_dims(x_test, -1).astype("float32") / 255 print("x_train shape:", x_train.shape) print("x_test shape:", x_test.shape) # Build data generator and model datagen = ImageDataGenerator() model = build_model(x_train, 10) # Add multi-GPU support. model = ParallelModel(model, GPU_COUNT) optimizer = keras.optimizers.SGD(lr=0.01, momentum=0.9, clipnorm=5.0) model.compile( loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"], ) model.summary() # Train model.fit_generator( datagen.flow(x_train, y_train, batch_size=64), steps_per_epoch=50, epochs=10, verbose=1, validation_data=(x_test, y_test), callbacks=[keras.callbacks.TensorBoard(log_dir=MODEL_DIR, write_graph=True)], )