1. 程式人生 > >Keras多GPU訓練以及載入權重無效的問題

Keras多GPU訓練以及載入權重無效的問題

目錄


本文講簡單的探討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)
# 其他程式碼
...                                    

參考連結