keras搬砖系列-多GPU上使用keras

keras搬砖系列-多GPU上使用keras

1，数据并行

数据并行是将目标模型在多个设备上各复制一份，并使用每个设备上的复制品处理整个数据集的不同部分数据。

from keras.utils import multi_gpu_model

parallel_model = multi_gpu_model(model,gpus=8)
parallel_model.compile(loss='categorical_crossentropy',optimizer='rmsprop')
parallel_model.fit(x,y,epochs=20,batch_size=256)

2，设备并行

设备并行是在不同的设备上运行同一个模型的不同部分，当模型含有多个并行的结构，这个方法比较适合。

input_a = keras.Input(shape=(120,256))
input_b = keras.Input(shape=(140,256))

shared_lstm = keras.layers.LSTM(64)

with tf.device_scope('/gpu:0'):
	encoded_a = shared_lstm(tweet_a)
with tf.device_scope('/gpu:1'):
	encoded_b = shared_lstm(tweet_b)

with tf.device_scope('/gpu:0'):
	merged_vector = keras.layers.concatenate([encoded_a,encoded_b],axis=-1)