iasjkk/bbox_detection

1import tensorflow as tf2import keras.backend as K3import keras.layers as KL4import keras.models as KM5 6 7class ParallelModel(KM.Model):8 9    def __init__(self, keras_model, gpu_count):10        """Class constructor.11        keras_model: The Keras model to parallelize12        gpu_count: Number of GPUs. Must be > 113        """14        self.inner_model = keras_model15        self.gpu_count = gpu_count16        merged_outputs = self.make_parallel()17        super(ParallelModel, self).__init__(inputs=self.inner_model.inputs,18                                            outputs=merged_outputs)19 20    def __getattribute__(self, attrname):21        """Redirect loading and saving methods to the inner model. That's where22        the weights are stored."""23        if 'load' in attrname or 'save' in attrname:24            return getattr(self.inner_model, attrname)25        return super(ParallelModel, self).__getattribute__(attrname)26 27    def summary(self, *args, **kwargs):28        """Override summary() to display summaries of both, the wrapper29        and inner models."""30        super(ParallelModel, self).summary(*args, **kwargs)31        self.inner_model.summary(*args, **kwargs)32 33    def make_parallel(self):34        """Creates a new wrapper model that consists of multiple replicas of35        the original model placed on different GPUs.36        """37        # Slice inputs. Slice inputs on the CPU to avoid sending a copy38        # of the full inputs to all GPUs. Saves on bandwidth and memory.39        input_slices = {name: tf.split(x, self.gpu_count)40                        for name, x in zip(self.inner_model.input_names,41                                           self.inner_model.inputs)}42 43        output_names = self.inner_model.output_names44        outputs_all = []45        for i in range(len(self.inner_model.outputs)):46            outputs_all.append([])47 48        # Run the model call() on each GPU to place the ops there49        for i in range(self.gpu_count):50            with tf.device('/gpu:%d' % i):51                with tf.name_scope('tower_%d' % i):52                    # Run a slice of inputs through this replica53                    zipped_inputs = zip(self.inner_model.input_names,54                                        self.inner_model.inputs)55                    inputs = [56                        KL.Lambda(lambda s: input_slices[name][i],57                                  output_shape=lambda s: (None,) + s[1:])(tensor)58                        for name, tensor in zipped_inputs]59                    # Create the model replica and get the outputs60                    outputs = self.inner_model(inputs)61                    if not isinstance(outputs, list):62                        outputs = [outputs]63                    # Save the outputs for merging back together later64                    for l, o in enumerate(outputs):65                        outputs_all[l].append(o)66 67        # Merge outputs on CPU68        with tf.device('/cpu:0'):69            merged = []70            for outputs, name in zip(outputs_all, output_names):71                # Concatenate or average outputs?72                # Outputs usually have a batch dimension and we concatenate73                # across it. If they don't, then the output is likely a loss74                # or a metric value that gets averaged across the batch.75                # Keras expects losses and metrics to be scalars.76                if K.int_shape(outputs[0]) == ():77                    # Average78                    m = KL.Lambda(lambda o: tf.add_n(o) / len(outputs), name=name)(outputs)79                else:80                    # Concatenate81                    m = KL.Concatenate(axis=0, name=name)(outputs)82                merged.append(m)83        return merged84 85 86if __name__ == "__main__":87    # Testing code below. It creates a simple model to train on MNIST and88    # tries to run it on 2 GPUs. It saves the graph so it can be viewed89    # in TensorBoard. Run it as:90    #91    # python3 parallel_model.py92 93    import os94    import numpy as np95    import keras.optimizers96    from keras.datasets import mnist97    from keras.preprocessing.image import ImageDataGenerator98 99    GPU_COUNT = 2100 101    # Root directory of the project102    ROOT_DIR = os.path.abspath("../")103 104    # Directory to save logs and trained model105    MODEL_DIR = os.path.join(ROOT_DIR, "logs")106 107    def build_model(x_train, num_classes):108        # Reset default graph. Keras leaves old ops in the graph,109        # which are ignored for execution but clutter graph110        # visualization in TensorBoard.111        tf.reset_default_graph()112 113        inputs = KL.Input(shape=x_train.shape[1:], name="input_image")114        x = KL.Conv2D(32, (3, 3), activation='relu', padding="same",115                      name="conv1")(inputs)116        x = KL.Conv2D(64, (3, 3), activation='relu', padding="same",117                      name="conv2")(x)118        x = KL.MaxPooling2D(pool_size=(2, 2), name="pool1")(x)119        x = KL.Flatten(name="flat1")(x)120        x = KL.Dense(128, activation='relu', name="dense1")(x)121        x = KL.Dense(num_classes, activation='softmax', name="dense2")(x)122 123        return KM.Model(inputs, x, "digit_classifier_model")124 125    # Load MNIST Data126    (x_train, y_train), (x_test, y_test) = mnist.load_data()127    x_train = np.expand_dims(x_train, -1).astype('float32') / 255128    x_test = np.expand_dims(x_test, -1).astype('float32') / 255129 130    print('x_train shape:', x_train.shape)131    print('x_test shape:', x_test.shape)132 133    # Build data generator and model134    datagen = ImageDataGenerator()135    model = build_model(x_train, 10)136 137    # Add multi-GPU support.138    model = ParallelModel(model, GPU_COUNT)139 140    optimizer = keras.optimizers.SGD(lr=0.01, momentum=0.9, clipnorm=5.0)141 142    model.compile(loss='sparse_categorical_crossentropy',143                  optimizer=optimizer, metrics=['accuracy'])144 145    model.summary()146 147    # Train148    model.fit_generator(149        datagen.flow(x_train, y_train, batch_size=64),150        steps_per_epoch=50, epochs=10, verbose=1,151        validation_data=(x_test, y_test),152        callbacks=[keras.callbacks.TensorBoard(log_dir=MODEL_DIR,153                                               write_graph=True)]154    )