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

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的数量,做以下修改.

d.device_type == 'CPU'