【tensorflow2.0】使用多GPU训练模型

如果使用多GPU训练模型，推荐使用内置fit方法，较为方便，仅需添加2行代码。

在Colab笔记本中：修改->笔记本设置->硬件加速器 中选择 GPU

注：以下代码只能在Colab 上才能正确执行。

可通过以下colab链接测试效果《tf_多GPU》：

https://colab.research.google.com/drive/1j2kp_t0S_cofExSN7IyJ4QtMscbVlXU-

MirroredStrategy过程简介：

• 训练开始前，该策略在所有 N 个计算设备上均各复制一份完整的模型；
• 每次训练传入一个批次的数据时，将数据分成 N 份，分别传入 N 个计算设备（即数据并行）；
• N 个计算设备使用本地变量（镜像变量）分别计算自己所获得的部分数据的梯度；
• 使用分布式计算的 All-reduce 操作，在计算设备间高效交换梯度数据并进行求和，使得最终每个设备都有了所有设备的梯度之和；
• 使用梯度求和的结果更新本地变量（镜像变量）；
• 当所有设备均更新本地变量后，进行下一轮训练（即该并行策略是同步的）。

tensorflow_version 2.x
import tensorflow as tf
print(tf.__version__)
from tensorflow.keras import * 
# 此处在colab上使用1个GPU模拟出两个逻辑GPU进行多GPU训练
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    # 设置两个逻辑GPU模拟多GPU训练
    try:
        tf.config.experimental.set_virtual_device_configuration(gpus[0],
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024),
             tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPU,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        print(e)

2.2.0-rc2

1 Physical GPU, 2 Logical GPUs

一，准备数据

MAX_LEN = 300
BATCH_SIZE = 32
(x_train,y_train),(x_test,y_test) = datasets.reuters.load_data()
x_train = preprocessing.sequence.pad_sequences(x_train,maxlen=MAX_LEN)
x_test = preprocessing.sequence.pad_sequences(x_test,maxlen=MAX_LEN)
 
MAX_WORDS = x_train.max()+1
CAT_NUM = y_train.max()+1
 
ds_train = tf.data.Dataset.from_tensor_slices((x_train,y_train)) 
          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) 
          .prefetch(tf.data.experimental.AUTOTUNE).cache()
 
ds_test = tf.data.Dataset.from_tensor_slices((x_test,y_test)) 
          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) 
          .prefetch(tf.data.experimental.AUTOTUNE).cache()

二，定义模型

tf.keras.backend.clear_session()
def create_model():
 
    model = models.Sequential()
 
    model.add(layers.Embedding(MAX_WORDS,7,input_length=MAX_LEN))
    model.add(layers.Conv1D(filters = 64,kernel_size = 5,activation = "relu"))
    model.add(layers.MaxPool1D(2))
    model.add(layers.Conv1D(filters = 32,kernel_size = 3,activation = "relu"))
    model.add(layers.MaxPool1D(2))
    model.add(layers.Flatten())
    model.add(layers.Dense(CAT_NUM,activation = "softmax"))
    return(model)
 
def compile_model(model):
    model.compile(optimizer=optimizers.Nadam(),
                loss=losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=[metrics.SparseCategoricalAccuracy(),metrics.SparseTopKCategoricalAccuracy(5)]) 
    return(model)

三，训练模型

# 增加以下两行代码
strategy = tf.distribute.MirroredStrategy()  
with strategy.scope(): 
    model = create_model()
    model.summary()
    model = compile_model(model)
 
history = model.fit(ds_train,validation_data = ds_test,epochs = 10)  

WARNING:tensorflow:NCCL is not supported when using virtual GPUs, fallingback to reduction to one device
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1')
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding (Embedding)        (None, 300, 7)            216874    
_________________________________________________________________
conv1d (Conv1D)              (None, 296, 64)           2304      
_________________________________________________________________
max_pooling1d (MaxPooling1D) (None, 148, 64)           0         
_________________________________________________________________
conv1d_1 (Conv1D)            (None, 146, 32)           6176      
_________________________________________________________________
max_pooling1d_1 (MaxPooling1 (None, 73, 32)            0         
_________________________________________________________________
flatten (Flatten)            (None, 2336)              0         
_________________________________________________________________
dense (Dense)                (None, 46)                107502    
=================================================================
Total params: 332,856
Trainable params: 332,856
Non-trainable params: 0
_________________________________________________________________
Epoch 1/10
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:GPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1').
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
281/281 [==============================] - 4s 15ms/step - sparse_categorical_accuracy: 0.3546 - loss: 3.5168 - sparse_top_k_categorical_accuracy: 0.7163 - val_sparse_categorical_accuracy: 0.5000 - val_loss: 3.3722 - val_sparse_top_k_categorical_accuracy: 0.7066
Epoch 2/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.5279 - loss: 3.3386 - sparse_top_k_categorical_accuracy: 0.7267 - val_sparse_categorical_accuracy: 0.5387 - val_loss: 3.3299 - val_sparse_top_k_categorical_accuracy: 0.7173
Epoch 3/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.5583 - loss: 3.3094 - sparse_top_k_categorical_accuracy: 0.7238 - val_sparse_categorical_accuracy: 0.5490 - val_loss: 3.3169 - val_sparse_top_k_categorical_accuracy: 0.7217
Epoch 4/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.5856 - loss: 3.2818 - sparse_top_k_categorical_accuracy: 0.7244 - val_sparse_categorical_accuracy: 0.5574 - val_loss: 3.3077 - val_sparse_top_k_categorical_accuracy: 0.7217
Epoch 5/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.5967 - loss: 3.2693 - sparse_top_k_categorical_accuracy: 0.7242 - val_sparse_categorical_accuracy: 0.5659 - val_loss: 3.2993 - val_sparse_top_k_categorical_accuracy: 0.7248
Epoch 6/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.6030 - loss: 3.2626 - sparse_top_k_categorical_accuracy: 0.7262 - val_sparse_categorical_accuracy: 0.5690 - val_loss: 3.2974 - val_sparse_top_k_categorical_accuracy: 0.7244
Epoch 7/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.6054 - loss: 3.2600 - sparse_top_k_categorical_accuracy: 0.7266 - val_sparse_categorical_accuracy: 0.5677 - val_loss: 3.2980 - val_sparse_top_k_categorical_accuracy: 0.7262
Epoch 8/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.6065 - loss: 3.2581 - sparse_top_k_categorical_accuracy: 0.7273 - val_sparse_categorical_accuracy: 0.5708 - val_loss: 3.2990 - val_sparse_top_k_categorical_accuracy: 0.7262
Epoch 9/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.6091 - loss: 3.2558 - sparse_top_k_categorical_accuracy: 0.7283 - val_sparse_categorical_accuracy: 0.5726 - val_loss: 3.2952 - val_sparse_top_k_categorical_accuracy: 0.7253
Epoch 10/10
281/281 [==============================] - 5s 18ms/step - sparse_categorical_accuracy: 0.6093 - loss: 3.2551 - sparse_top_k_categorical_accuracy: 0.7288 - val_sparse_categorical_accuracy: 0.5726 - val_loss: 3.2908 - val_sparse_top_k_categorical_accuracy: 0.7244

参考：

开源电子书地址：https://lyhue1991.github.io/eat_tensorflow2_in_30_days/

GitHub 项目地址：https://github.com/lyhue1991/eat_tensorflow2_in_30_days