如何使用 gpu 并行训练 tensorflow.keras 模型？ TensorFlow 版本 2.5.0

Posted 2023-03-11

【中文标题】如何使用 gpu 并行训练 tensorflow.keras 模型？ TensorFlow 版本 2.5.0

【英文标题】：How to train tensorflow.keras models in parallel using gpu? Tensorflow version 2.5.0

【发布时间】：2021-09-24 09:04:17

【问题描述】：

我有以下代码运行我在不同模块中拥有的自定义模型，并将几个参数（学习率、卷积核大小等）作为输入

custom_model是在tensorflow中编译一个tensorflow.keras.models.Model并返回模型的函数。

LOW 是训练数据集

HIGH 是目标数据集

我通过 hdf5 文件加载了这两个文件，但数据集非常大，大约 10 GB。

通常我在 jupyter-lab 中运行它没有问题，并且模型不会消耗 GPU 上的资源。最后，我保存不同参数的权重。

现在我的问题是：

如何将其作为脚本并针对k1 和k2 的不同值并行运行。 我想像 bash 循环之类的东西会做，但我想避免重新读取数据集。 我正在使用 Windows 10 作为操作系统。

import tensorflow as tf
physical_devices = tf.config.list_physical_devices('GPU') 
for gpu_instance in physical_devices: 
    tf.config.experimental.set_memory_growth(gpu_instance, True)
import h5py

from model_custom import custom_model
winx = 100
winz = 10
k1 = 9
k2 = 5

with h5py.File('MYFILE', 'r') as hf:
    LOW = hf['LOW'][:]
    HIGH = hf['HIGH'][:]

with tf.device("/gpu:1"):
    mymodel = custom_model(winx,winz,lrate=0.001,usebias=True,kz1=k1, kz2=k2)
    myhistory = mymodel.fit(LOW, HIGH, batch_size=1, epochs=1)
    mymodel.save_weights('zkernel__kz1__kz2_.hdf5'.format(winz, k1,k2))

【参考方案1】：

我发现这个解决方案对我来说很好用。这使得使用 MPI 和 mpi4py 在 gpus 中运行并行模型训练成为可能。当我尝试加载大文件并同时运行多个进程以使进程数乘以我加载的数据超过我的内存容量时，只有一个问题。

from mpi4py import MPI 
import tensorflow as tf
physical_devices = tf.config.list_physical_devices('GPU') 
for gpu_instance in physical_devices: 
    tf.config.experimental.set_memory_growth(gpu_instance, True)
import h5py
from model_custom import custom_model

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

winx = 100
winy = 100
winz = 10

if rank == 10:
    with h5py.File('mifile.hdf5', 'r') as hf:
        LOW = hf['LOW'][:]
        HIGH = hf['HIGH'][:]
else:
    HIGH = None
    LOW= None
HIGH = comm.bcast(HIGH, root=10)
LOW = comm.bcast(LOW, root=10)
    
if rank < 5:
    with tf.device("/gpu:1"):
        k = 9
        q = rank +1
        mymodel1 = custom_model(winx,winz,lrate=0.001,usebias=True,kz1=k, kz2=q)
        mymodel1._name = '__'.format(winz,k,q)
        myhistory1 = mymodel1.fit(LOW, HIGH, batch_size=1, epochs=1)
        mymodel1.save_weights(mymodel1.name +'winz__k__q_.hdf5'.format(winz, k,q))

elif 5 <= rank < 10: 
    with tf.device("/gpu:2"):
        k = 8
        q = rank +1 -5
        mymodel2 = custom_model(winx,winz,lrate=0.001,usebias=True,kz1=k, kz2=q)
        mymodel2._name = '__'.format(winz,k,q)
        myhistory2 = mymodel2.fit(LOW, HIGH, batch_size=1, epochs=1)
        mymodel2.save_weights(mymodel2.name +'winz__k__q_.hdf5'.format(winz, k,q))

然后我保存到名为 mycode.py 的 python 模块，然后在控制台中运行

mpiexec -n 11 python ./mycode.py