欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

keras 多GPU训练,单GPU预测

程序员文章站 2022-05-19 08:56:04
...

多GPU训练

keras自带模块 multi_gpu_model,此方式为数据并行的方式,将将目标模型在多个设备上各复制一份,并使用每个设备上的复制品处理整个数据集的不同部分数据,最高支持在8片GPU上并行。
使用方式:

from keras.utils import multi_gpu_model
# Replicates `model` on 8 GPUs.
# This assumes that your machine has 8 available GPUs.
parallel_model = multi_gpu_model(model, gpus=4) ##
parallel_model.compile(loss='categorical_crossentropy',
                       optimizer='adam')

# This `fit` call will be distributed on 8 GPUs.
# Since the batch size is 256, each GPU will process 32 samples.
parallel_model.fit(x, y, epochs=40, batch_size=128)

单GPU预测

因为定义的多核训练,所以网络的每一层都是按GPU来命名的,训练时采用多个GPU那么当导入参数的时候必须指定相同数量的GPU才行,如上代码的指定方式。但是,但我们将model切换到单GPU的环境中时,则会出现错误,此时我们必须将参数保存为单GPU的形式。
方法:
在原多GPU环境中导入模型,保存为单GPU版本,代码如下:

from keras.utils import multi_gpu_model
paralleled_model=multi_gpu_model(seg_model,gpus=4)
paralleled_model.load_weights("srf_best_model.h5") # 此时model也自动载入了权重,可用model进行预测
seg_model.save('srf_single_gpu_best_model.h5')

注意:paralleled_model为多GPUmodel,而seg_model为单GPUmodel。
加载单GPU模型:

seg_model.load_weights("srf_single_gpu_best_model.h5") 

参考:

  1. 【keras】一台设备上同时使用多张显卡训练同一个网络模型
  2. Keras 多GPU下模型和参数保存Modelcheckpoint callback报错问题以及在单GPU/CPU下载入