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