Very easy. Just build a network within the scope of tf.distribute.MirroredStrategy tf.distribute.MirroredStrategy. You only need to change a few lines.
Here, as a simple example, we are building a simple network with only one hidden layer.
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.optimizers import Adam
with tf.distribute.MirroredStrategy().scope():
#Build a network with this block
x = Input(32, dtype=tf.int32, name='x')
hidden = Dense(20, activation='relu', name='hidden')(x)
y = Dense(5, activation='sigmoid', name='y')(hidden)
model = Model(inputs=x, outputs=y)
model.compile(
optimizer=Adam(lr=0.001),
loss='binary_crossentropy',
)
model.fit(
x_train, y_train,
epochs=10,
batch_size=16,
)
The official documentation is as follows.
If you want to use the Keras API as it is, as introduced in this article, but if you implement a custom training loop, there are some additional points to consider. In that case, please refer to the above document.
multi_gpu_model () is scheduled to be abolished after April 2020.