Keras并行训练

Keras如果以Tensorflow作为后端,是自动支持GPU的。也就是说你只管定义模型,compile,fit。如果能检测到可用的GPU,就会自动跑在GPU上。

数据并行

如果要用多GPU训练,只需要用multi_gpu_model函数再包裹 一下。具体如下

## 代码来自keras官方文档  网址 https://keras.io/getting-started/faq/#how-can-i-run-keras-on-gpu
from keras.utils import multi_gpu_model

# Replicates `model` on 8 GPUs.
# This assumes that your machine has 8 available GPUs.
# 返回一个并行模型
parallel_model = multi_gpu_model(model, gpus=8)
parallel_model.compile(loss='categorical_crossentropy',
                       optimizer='rmsprop')

# This `fit` call will be distributed on 8 GPUs.
# Since the batch size is 256, each GPU will process 32 samples.
parallel_model.fit(x, y, epochs=20, batch_size=256)

multi_gpu_model是数据并行方式,即将模型复制到各个GPU上,每个batch被均匀划分到各个GPU上训练,最后合并更新参数。
值得注意的是compile必须是模型定义的最后一步。

结构并行

结构并行将不同的分支放到不同的设备上计算,然后将计算结果合并到某个设备下。需要借助tf的代码。

## 代码来自keras官方文档  网址 https://keras.io/getting-started/faq/#how-can-i-run-keras-on-gpu
# Model where a shared LSTM is used to encode two different sequences in parallel
input_a = keras.Input(shape=(140, 256))
input_b = keras.Input(shape=(140, 256))

shared_lstm = keras.layers.LSTM(64)

# Process the first sequence on one GPU
with tf.device_scope('/gpu:0'):
    encoded_a = shared_lstm(tweet_a)
# Process the next sequence on another GPU
with tf.device_scope('/gpu:1'):
    encoded_b = shared_lstm(tweet_b)

# Concatenate results on CPU
with tf.device_scope('/cpu:0'):
    merged_vector = keras.layers.concatenate([encoded_a, encoded_b],
                                             axis=-1)

你可能感兴趣的:(深度学习)