同時在不同的 GPU 上訓練多個 keras/tensorflow 模型

Question

我想在 jupyter notebook 中同時在多個 GPU 上訓練多個模型。 我正在處理一個帶有 4GPU 的節點。 我想將一個 GPU 分配給一個模型並同時訓練 4 個不同的模型。 現在，我通過（例如）為一台筆記本選擇了 GPU：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

def model(...):
    ....

model.fit(...)

在四個不同的筆記本中。 但是，擬合過程的結果和輸出分布在四個不同的筆記本中。 但是，按順序在一個筆記本中運行它們需要很多時間。 您如何將 GPU 分配給各個功能並並行運行它們？

Answer 1

我建議像這樣使用 Tensorflow 范圍：

with tf.device_scope('/gpu:0'):
  model1.fit()
with tf.device_scope('/gpu:1'):
  model2.fit()
with tf.device_scope('/gpu:2'):
  model3.fit()

Answer 2

如果你想在不同的雲 GPU 上訓練模型（例如來自 AWS 的 GPU 實例），試試這個庫：

!pip install aibro==0.0.45 --extra-index-url https://test.pypi.org/simple

from aibro.train import fit
machine_id = 'g4dn.4xlarge' #instance name on AWS
job_id, trained_model, history = fit(
    model=model,
    train_X=train_X,
    train_Y=train_Y,
    validation_data=(validation_X, validation_Y),
    machine_id=machine_id
)

教程： https : //colab.research.google.com/drive/19sXZ4kbic681zqEsrl_CZfB5cegUwuIB#scrollTo=ERqoHEaamR1Y