如何在Keras中于两块不同GPU上并行运行两个模型

完全可以实现你的需求！你之前的代码之所以是串行运行，核心原因很简单：model.fit()是同步阻塞调用——哪怕你给两个模型指定了不同的GPU，程序也会老老实实等第一个模型的训练完全结束后，才会开始执行第二个模型的训练逻辑。

方案1：用Python多线程实现并行 #

你可以借助Python的threading模块，把两个模型的训练逻辑拆到不同线程里，让它们同时跑。需要注意两个点：一是每个线程里的模型要绑定到对应的GPU，二是要开启显存动态增长，避免某块GPU直接占满所有显存导致另一块无法使用。

给你改好的示例代码：

import threading
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, LSTM, Dense

# 先配置显存动态增长，防止单GPU抢占全部资源
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

# 把训练逻辑封装成函数，方便线程调用
def train_on_gpu(gpu_id, lstm_units):
    with tf.device(f'/gpu:{gpu_id}'):
        model = Sequential()
        model.add(embedding)  # 用你之前定义好的embedding层就行
        model.add(LSTM(lstm_units))
        model.add(Dense(5, activation='softmax'))
        model.compile(loss='categorical_crossentropy', optimizer='adam')
        model.fit(np.array(train_x), np.array(train_y), epochs=15, batch_size=15)

# 创建两个线程，分别对应GPU0和GPU1的模型训练
thread_gpu0 = threading.Thread(target=train_on_gpu, args=(0, 50))
thread_gpu1 = threading.Thread(target=train_on_gpu, args=(1, 100))

# 启动线程，让两个模型同时开始训练
thread_gpu0.start()
thread_gpu1.start()

# 等待两个线程都训练完成
thread_gpu0.join()
thread_gpu1.join()

方案2：用多进程彻底隔离环境（避免线程冲突） #

如果遇到TensorFlow版本对线程支持不友好的情况（比如某些老版本TF的全局状态会干扰线程），可以改用多进程。每个进程有独立的内存空间，能彻底隔离两个模型的运行环境，完全避免资源竞争问题。

示例代码如下：

import multiprocessing
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Embedding, LSTM, Dense

def train_on_gpu(gpu_id, lstm_units, train_x, train_y):
    # 每个进程单独配置GPU可见性和显存
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            # 只让当前进程使用指定的GPU
            tf.config.set_visible_devices(gpus[gpu_id], 'GPU')
            tf.config.experimental.set_memory_growth(gpus[gpu_id], True)
        except RuntimeError as e:
            print(f"GPU配置出错：{e}")
    
    # 构建并训练模型
    model = Sequential()
    model.add(embedding)
    model.add(LSTM(lstm_units))
    model.add(Dense(5, activation='softmax'))
    model.compile(loss='categorical_crossentropy', optimizer='adam')
    model.fit(np.array(train_x), np.array(train_y), epochs=15, batch_size=15)

if __name__ == '__main__':
    # 创建进程，注意多进程中要传入训练数据（如果数据大的话可以考虑用共享内存）
    process_gpu0 = multiprocessing.Process(target=train_on_gpu, args=(0, 50, train_x, train_y))
    process_gpu1 = multiprocessing.Process(target=train_on_gpu, args=(1, 100, train_x, train_y))
    
    # 启动进程
    process_gpu0.start()
    process_gpu1.start()
    
    # 等待进程结束
    process_gpu0.join()
    process_gpu1.join()

几个关键提醒 #

1. 显存必须管！ 不管用线程还是进程，一定要开启显存动态增长，或者手动设置显存限制。默认情况下TensorFlow会占满GPU的全部显存，直接导致另一个模型无法分配资源。
2. 模型完全独立：两个模型的权重、训练状态都是完全分开的，完全符合你做参数网格搜索的需求。
3. TF版本适配：如果用TensorFlow 1.x，需要额外处理每个线程/进程的独立会话（tf.Session）；TF 2.x的即时执行模式会省心很多。

这样调整后，两个模型就能在两块GPU上同时并行训练了，刚好满足你的网格搜索场景。

内容的提问来源于stack exchange，提问作者lefunction