keras實現多gpu訓練,單gpu保存模型,單gpu預測

問題:我在訓練模型的時候用的是4個gpu,我需要將訓練好的模型加載到樹莓派中,但是樹莓派只有一個cpu or gpu

解決方法:多gpu進行訓練,單gpu保存訓練的模型,用單gpu進行預測

代碼呈現,以mnist數據集爲例:

from keras.datasets import mnist
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dropout, Dense
from keras.utils import to_categorical
from keras import Sequential
from keras import losses
from keras import optimizers
from keras.utils import multi_gpu_model

train_x, train_y = mnist.load_data()[0]
test_x, test_y = mnist.load_data()[1]
train_x = train_x.reshape((-1, 28, 28, 1))
test_x = test_x.reshape((-1, 28, 28, 1))
train_x = train_x.astype('float32')
test_x = test_x.astype('float32')
train_y = to_categorical(train_y, 10)
test_y = to_categorical(test_y, 10)

# build the model
model = Sequential()
model.add(Conv2D(filters=32, kernel_size=(5,5), activation='relu', input_shape=[28, 28, 1]))
model.add(Conv2D(filters=64, kernel_size=(5,5), activation='relu'))
model.add(MaxPooling2D(pool_size=(2,2)))
model.add(Flatten())
model.add(Dropout(0.5))
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(10, activation='softmax'))

model.summary()

# use 4 gpus
model  = multi_gpu_model(model, gpus=4)

model.compile(
    loss = losses.categorical_crossentropy,
    optimizer = optimizers.Adam(),
    metrics = ['accuracy']
)
batch_size = 100
epochs = 10
save_path = 'four_gpus_model.hdf5'

model.fit(train_x, train_y, epochs=epochs, batch_size=batch_size)
model.save_weights(save_path)

上面是多gpu訓練以及保存模型,下面可以看到用單個cpu加載模型進行預測的時候會出現什麼問題:

# 指定cpu
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

model.load_weights('four_gpus_model.hdf5')
model.predict(test_x)

出現的bug是:

ValueError: You are trying to load a weight file containing 1 layers into a model with 4 layers.

解決方法是:多gpu訓練用的model和創建的model命名不同

# notice:the name is 'model_mul_gpu'
model_mul_gpu = multi_gpu_model(model, gpus=4)
batch_size = 100
epochs = 10
save_path = 'single_cpu_model.hdf5'

model_mul_gpu.compile(
    loss = losses.categorical_crossentropy,
    optimizer = optimizers.Adam(),
    metrics = ['accuracy']
)

model_mul_gpu.fit(train_x, train_y, epochs=epochs, batch_size=batch_size)

# the save model is 'model'
model.save_weights(save_path)

再用單cpu加載

model.load_weights('single_cpu_model.hdf5')
model.predict(test_x)

問題解決了!

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章