keras实现多gpu训练,单gpu保存模型,单gpu预测

2022年12月26日12:58:14

问题:我在训练模型的时候用的是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)

问题解决了!

  • 作者:jin__9981
  • 原文链接:https://blog.csdn.net/jin__9981/article/details/100693949
    更新时间:2022年12月26日12:58:14 ,共 2009 字。