Python tensorflow.keras.utils.multi_gpu_model() Examples

def make_multi_gpu_model(model, num_GPUs, logger=None):
    """
    Takes a compiled tf.keras Model object 'model' and applies
    from tensorflow.keras.utils import multi_gpu_model
    ... to mirror it across multiple visible GPUs. Input batches to 'model'
    are split evenly across the GPUs.

    Args:
        model:    (tf.keras Model) A compiled tf.keras Model object.
        num_GPUs: (int)            Number of GPUs to distribute the model over
        logger:   (Logger)         Optional Logger object

    Returns:
        The split, multi-GPU model.
        The original model
        Note: The two will be the same for num_GPUs=1
    """
    org_model = model
    if num_GPUs > 1:
        from tensorflow.keras.utils import multi_gpu_model
        model = multi_gpu_model(org_model, gpus=num_GPUs,
                                cpu_merge=False, cpu_relocation=False)
        logger = logger or ScreenLogger()
        logger("Creating multi-GPU model: N=%i" % num_GPUs)
    return model, org_model 

def load_model(input_shape, weights_path, net, prior_std,
               kernel_size, activation, padding, num_gpus):
    """Loads model from .h5 file.

    If model is saved as multi-gpu, re-saves it as single-gpu.
    """

    # Loads model as multi-gpu, if possible.
    try:
        model = net(input_shape,
                    kernel_size=kernel_size,
                    activation=activation,
                    padding=padding,
                    prior_std=prior_std)
        model = multi_gpu_model(model, gpus=num_gpus)

        # Converts .h5 file to single-gpu.
        model.load_weights(weights_path)
        model = model.layers[-2]
        model.save_weights(weights_path)
    except ValueError as e:
        pass

    # Loads single-gpu model.
    model = net(input_shape,
                kernel_size=kernel_size,
                activation=activation,
                padding=padding,
                prior_std=prior_std)
    model.load_weights(weights_path)

    return model 

def __init__(self, preprocessor, architecture=Architecture(),
                 loss=tversky_loss, metrics=[dice_soft],
                 learninig_rate=0.0001, batch_queue_size=2,
                 workers=1, gpu_number=1):
        # Identify data parameters
        self.three_dim = preprocessor.data_io.interface.three_dim
        self.channels = preprocessor.data_io.interface.channels
        self.classes = preprocessor.data_io.interface.classes
        # Assemble the input shape
        input_shape = (None,)
        # Initialize model for 3D data
        if self.three_dim:
            input_shape = (None, None, None, self.channels)
            self.model = architecture.create_model_3D(input_shape=input_shape,
                                                      n_labels=self.classes)
         # Initialize model for 2D data
        else:
             input_shape = (None, None, self.channels)
             self.model = architecture.create_model_2D(input_shape=input_shape,
                                                       n_labels=self.classes)
        # Transform to Keras multi GPU model
        if gpu_number > 1:
            self.model = multi_gpu_model(self.model, gpu_number)
        # Compile model
        self.model.compile(optimizer=Adam(lr=learninig_rate),
                           loss=loss, metrics=metrics)
        # Cache starting weights
        self.initialization_weights = self.model.get_weights()
        # Cache parameter
        self.preprocessor = preprocessor
        self.loss = loss
        self.metrics = metrics
        self.learninig_rate = learninig_rate
        self.batch_queue_size = batch_queue_size
        self.workers = workers

    #---------------------------------------------#
    #               Class variables               #
    #---------------------------------------------# 

def model(self):
    ''' keras Model before doing `multi_gpu_model` '''
    return self.raw_model.model 

def parallel_model(self):
    ''' `multi_gpu_model` of keras Model '''
    assert self._parallel_model is not None
    return self._parallel_model 

def build(self, multi_gpu=False):
    ''' main entrypoint to build model '''
    assert self.model

    loss = self.get_loss()
    optimizer = self.get_optimizer()

    run_opts, run_metas = self.get_run_opts_metas()

    # compile model
    if self.ngpu > 1 and multi_gpu:
      self._parallel_model = multi_gpu_model(
          self.model, gpus=self.ngpu, cpu_relocation=False, cpu_merge=False)
      self.parallel_model.compile(
          loss=loss,
          optimizer=optimizer,
          metrics=self._metrics_used,
          options=run_opts,
          run_metadata=run_metas)
    else:
      self.model.compile(
          loss=loss,
          optimizer=optimizer,
          metrics=self._metrics_used,
          options=run_opts,
          run_metadata=run_metas)

    # Print model summary
    if self.model.built and self.model._is_graph_network:
      self.model.summary()
    self._built = True 

def _generate_model(self):
        '''to generate the bounding boxes'''
        weights_path = os.path.expanduser(self.weights_path)
        assert weights_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        # Load model, or construct model and load weights.
        num_anchors = len(self.anchors)
        num_classes = len(self.class_names)
        #YOLOv3 model has 9 anchors and 3 feature layers but
        #Tiny YOLOv3 model has 6 anchors and 2 feature layers,
        #so we can calculate feature layers number to get model type
        num_feature_layers = num_anchors//3

        try:
            if num_anchors == 5:
                # YOLOv2 use 5 anchors
                yolo_model, _ = get_yolo2_model(self.model_type, num_anchors, num_classes, input_shape=self.model_image_size + (3,), model_pruning=self.pruning_model)
            else:
                yolo_model, _ = get_yolo3_model(self.model_type, num_feature_layers, num_anchors, num_classes, input_shape=self.model_image_size + (3,), model_pruning=self.pruning_model)
            yolo_model.load_weights(weights_path) # make sure model, anchors and classes match
            if self.pruning_model:
                yolo_model = sparsity.strip_pruning(yolo_model)
            yolo_model.summary()
        except Exception as e:
            print(repr(e))
            assert yolo_model.layers[-1].output_shape[-1] == \
                num_anchors/len(yolo_model.output) * (num_classes + 5), \
                'Mismatch between model and given anchor and class sizes'
        print('{} model, anchors, and classes loaded.'.format(weights_path))
        if self.gpu_num>=2:
            yolo_model = multi_gpu_model(yolo_model, gpus=self.gpu_num)

        return yolo_model 

def main(args):
    include_top = True
    if args.dump_headless:
        include_top = False

    # prepare model
    model, input_shape = get_model(args.model_type, include_top=include_top)
    if args.weights_path:
        model.load_weights(args.weights_path, by_name=True)
    # support multi-gpu training
    if args.gpu_num >= 2:
        model = multi_gpu_model(model, gpus=args.gpu_num)
    model.summary()

    if args.evaluate:
        K.set_learning_phase(0)
        evaluate_model(args, model, input_shape)
    elif args.verify_with_image:
        K.set_learning_phase(0)
        verify_with_image(model, input_shape)
    elif args.dump_headless:
        K.set_learning_phase(0)
        model.save(args.output_model_file)
        print('export headless model to %s' % str(args.output_model_file))
    else:
        train(args, model, input_shape)