Tensorflow-keras 2.X多GPU预测

b1payxdu 于 2023-03-23 发布在其他

关注(0)|答案(1)|浏览(123)

我有4GPU（RTX3090）在一台PC。
我只使用了1GPU进行训练和预测，但现在我将使用4GPU。
在训练期间，4GPU激活成功，但只有1GPU处于活动状态以进行预测。

mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
    model = MyModel()

model_checkpoint = ModelCheckpoint(model_path, monitor='loss', verbose=1, save_best_only=True)
history = model.fit(trainData, steps_per_epoch=num_images//batch_size, verbose=1)

使用4 GPU

model.predict(testData, verbose=1)
#for x in testData:
#   model.predict_on_batch(x)

仅使用1个GPU
如何使用我的全GPU？

keras

来源：https://stackoverflow.com/questions/72571288/tensorflow-keras-2-x-multi-gpu-prediction

1条答案

按热度按时间

mitkmikd1#

我解决了！改变了我的测试数据集！
我正在使用测试数据集与发电机。

def testGenerator(image_path):
    for file in image_path:
        img = file
        img = img / 255
        img = np.reshape(img,img.shape+(1,)) if (not flag_multi_class) else img 
    yield img

我想将数据集加载到GPU内存中。
所以我使用了“tf.data.Dataset.from_generator”，我的GPU工作了！
加载保存的模型

model = tf.keras.models.load_model(model_path)

创建TensorFlow数据集

testGene = testGenerator(img_data)
test_dataset = tf.data.Dataset.from_generator(
    lambda: testGene,
    output_types = tf.float64,
    #output_shapes = tf.TensorShape([None]),
    output_shapes = tf.TensorShape([512, 512, 1])
    ).batch(global_batch_size)

result = model.predict(
    test_dataset,
    verbose=1,
    steps=math.ceil(max / global_batch_size),
    callbacks=[CustomCallback(root_file_list, int(max / len(root_file_list) / global_batch_size))])

它使用我的所有GPU工作！

赞(0）回复(0）举报 2023-03-23

我来回答

Tensorflow-keras 2.X多GPU预测

1条答案

相关问题

热门标签

最新问答