Keras多GPU訓練以及載入權重無效的問題
阿新 • • 發佈:2018-11-15
目錄
本文講簡單的探討Keras中使用多GPU訓練的方法以及需要注意的地方。有兩種方法可在多個 GPU 上執行單個模型:資料並行和裝置並行(Keras官方建議使用 TensorFlow 後端)。第一部分講如何使用資料並行的方式使用多GPU加速;第二部分給出一個裝置並行的例子。
1、資料並行
1.1、單GPU或者無GPU訓練的程式碼如下:
# 定義資料集
train_dataset, train_steps = ...
val_dataset, val_steps = ...
# 定義一個模型
model = ...
model.summary()
# 編譯模型。因為是多類分類問題,所以loss是categorical_crossentropy
# lr一定不要太高,不然訓練的精度有時能上去,有時上不去
model.compile(optimizer=keras.optimizers.Adam(lr=0.0001, decay=1e-5),
loss=keras.losses.categorical_crossentropy,
metrics= [keras.metrics.categorical_accuracy])
# 設定回撥
callbacks = [
# 把TensorBoard日誌寫到'logs'裡面
keras.callbacks.TensorBoard(log_dir='./logs'),
# 當categorical_accuracy,也就是分類精度在10個epoh之內都沒提升時,降低learning rate
keras.callbacks.ReduceLROnPlateau(monitor='categorical_accuracy', patience=10, verbose=2),
# 當categorical_accuracy在15個epoch內沒有提升的時候,停止訓練
keras.callbacks.EarlyStopping(monitor='categorical_accuracy', patience=15, verbose=2)
]
# 開始訓練
model.fit(train_dataset, epochs=NUM_TRAIN_EPOCH, steps_per_epoch=train_steps,
validation_data=val_dataset, validation_steps=val_steps, callbacks=callbacks)
# 儲存模型的權重
model.save_weights('model/latest.weights.h5', save_format='h5')
1.2、資料並行的多GPU 訓練
資料並行包括在每個裝置上覆制一次目標模型,並使用每個模型副本處理不同部分的輸入資料。Keras 有一個內建的實用函式 keras.utils.multi_gpu_model,它可以生成任何模型的資料並行版本,在多達 8 個 GPU 上實現準線性加速。資料並行的意思是,把模型複製,然後把一批(Batch)資料分為同等大小的幾個子批在不同的GPU上訓練,然後再在CPU撒謊嗯合併結果。
下面給出1.1中對應的使用多GPU訓練的程式碼:
# 定義資料集
train_dataset, train_steps = ...
val_dataset, val_steps = ...
# 定義一個模型
model = ...
model.summary()
# 獲取資料並行的模型,其中2是GPU的數量。
parallel_model = keras.utils.multi_gpu_model(model, 2)
# 編譯模型。因為是多類分類問題,所以loss是categorical_crossentropy
# lr一定不要太高,不然訓練的精度有時能上去,有時上不去
parallel_model.compile(optimizer=keras.optimizers.Adam(lr=0.0001, decay=1e-5),
loss=keras.losses.categorical_crossentropy,
metrics=[keras.metrics.categorical_accuracy])
# 設定回撥
callbacks = [
# 把TensorBoard日誌寫到'logs'裡面
keras.callbacks.TensorBoard(log_dir='./logs'),
# 當categorical_accuracy,也就是分類精度在10個epoh之內都沒提升時,降低learning rate
keras.callbacks.ReduceLROnPlateau(monitor='categorical_accuracy', patience=10, verbose=2),
# 當categorical_accuracy在15個epoch內沒有提升的時候,停止訓練
keras.callbacks.EarlyStopping(monitor='categorical_accuracy', patience=15, verbose=2)
]
# 開始訓練
parallel_model.fit(train_dataset, epochs=NUM_TRAIN_EPOCH, steps_per_epoch=train_steps,
validation_data=val_dataset, validation_steps=val_steps, callbacks=callbacks)
# 儲存模型的權重
model.save_weights('model/latest.weights.h5', save_format='h5')
注意:儲存權重時,一定要使用model,也就是不是並行的模型,來儲存權重。 如果使用parallel_model.save_weights儲存權重,在使用model.load_weights載入權重時,會出現類似於如下錯誤:
ValueError: You are trying to load a weight file containing 1 layers into a model with 4 layers.
在model.load_weights加上by_name=True時,倒是不會報錯,但是載入的權重是無效的,可以自行測試。
2、裝置並行
裝置並行性包括在不同裝置上運行同一模型的不同部分。對於具有並行體系結構的模型,例如有兩個分支的模型,這種方式很合適。這種並行可以通過使用 tf.device來實現。一個簡單的例子:
# 模型中共享的 LSTM 用於並行編碼兩個不同的序列
input_a = keras.Input(shape=(140, 256))
input_b = keras.Input(shape=(140, 256))
shared_lstm = keras.layers.LSTM(64)
# 在一個 GPU 上處理第一個序列
with tf.device('/device:GPU:0'):
encoded_a = shared_lstm(input_a)
# 在另一個 GPU上 處理下一個序列
with tf.device('/device:GPU:1'):
encoded_b = shared_lstm(input_b)
# 在 CPU 上連線結果
with tf.device('/cpu:0'):
merged_vector = keras.layers.concatenate([encoded_a, encoded_b], axis=-1)
# 整個全連線層
fc1 = keras.layers.Dense(50, activation='softmax')(merged_vector)
# 新建模型
model = keras.Model(inputs=[input_a, input_b], outputs=fc1)
# 其他程式碼
...