tensorflow 计算运算设备的GPU或者CPU的数量

最新推荐文章于 2025-07-09 18:21:25 发布

原创最新推荐文章于 2025-07-09 18:21:25 发布 · 2.5k 阅读

0 ·

CC 4.0 BY-SA版权

tensorflow 专栏收录该内容

91 篇文章

订阅专栏

本文介绍了一段用于检测TensorFlow环境中可用GPU数量的Python代码。该程序能够自动识别多GPU模式，并验证批量大小是否符合GPU数量的要求。对于CPU数量的检测也提供了相应的修改建议。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

tensorflow 计算运算设备的GPU数量的程序

if FLAGS.multi_gpu:
    from tensorflow.python.client import device_lib

    local_device_protos = device_lib.list_local_devices()
    num_gpus = sum([1 for d in local_device_protos if d.device_type == 'GPU'])
    if not num_gpus:
        raise ValueError('Multi-GPU mode was specified, but no GPUs '
                        'were found. To use CPU, run --multi_gpu=False.')

    remainder = batch_size % num_gpus
    if remainder:
        err = ('When running with multiple GPUs, batch size '
                'must be a multiple of the number of available GPUs. '
                'Found {} GPUs with a batch size of {}; try --batch_size={} instead.'
                ).format(num_gpus, batch_size, batch_size - remainder)
        raise ValueError(err)
    return num_gpus

上述程序,返回的为gpu的个数,如1,2,3,4....如果计算CPU的数量,做以下修改.