News: Keras now supports DataParallel equivalent

FuriouslyCurious · October 12, 2017, 2:31am

Keras can now do multi-GPU models. Seems like they can also run a model on CPU and merge gradients with GPUs to utilize all system resources.

github.com

keras-team/keras/blob/3dd3e8331677e68e7dec6ed4a1cbf16b7ef19f7f/keras/utils/training_utils.py

from ..layers.merge import concatenate
from .. import backend as K
from ..layers.core import Lambda
from ..engine.training import Model


def _get_available_devices():
    from tensorflow.python.client import device_lib
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos]


def multi_gpu_model(model, gpus):
    """Replicates a model on different GPUs.

    Specifically, this function implements single-machine
    multi-GPU data parallelism. It works in the following way:

    - Divide the model's input(s) into multiple sub-batches.
    - Apply a model copy on each sub-batch. Every model copy

This file has been truncated. show original