【keras】多GPU训练模型及保存

下面是一个小demo，需要注意的地方是保存模型的时候只保存单个GPU的模型，不要保存多GPU训练的模型，否则加载时会报错的。

from sklearn.datasets import make_blobs
from keras.layers import Dense, Input
from keras import Model
from keras.utils import multi_gpu_model

def MLP(hidden=50):
    inputs = Input(shape=(2,))
    x = Dense(units=hidden, activation='relu')(inputs)
    out = Dense(units=3, activation='softmax')(x)
    return Model(inputs, out)


x, y = make_blobs(n_samples=5000, centers=3, n_features=2,
                  cluster_std=2, random_state=2)
trainnum = 2000
train_x, train_y = x[:trainnum], y[:trainnum]
evalu_x, evalu_y = x[trainnum:], y[trainnum:]
print('train_x %s train_y %s   evalu_x %s evalu_y %s' %
      (train_x.shape, train_y.shape, evalu_x.shape, evalu_y.shape))

mlp = MLP()
par_mlp = multi_gpu_model(mlp, gpus=2)
par_mlp.compile(loss='sparse_categorical_crossentropy',
                optimizer='adam', metrics=['accuracy'])
par_mlp.fit(train_x, train_y, validation_data=(
    evalu_x, evalu_y), epochs=50, verbose=1, batch_size=200)
mlp.save_weights('mlp.h5') # 保存模型只保存单个模型的权重

# 加载预测时
model = MLP()
model.compile(loss='sparse_categorical_crossentropy',
              optimizer='adam', metrics=['accuracy'])
model.load_weights('mlp.h5')
_, evalu_acc = model.evaluate(evalu_x, evalu_y, verbose=1)
print('> Evaluate accuracy %.4f' % evalu_acc)