0


tensorflow如何使用gpu

目录

1、查看GPU的数量

import tensorflow as tf
# 查看gpu和cpu的数量
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)

2、设置GPU加速

第一种:限制使用的gpu,没有限制消耗内存的大小:

  通过 tf.config.experimental.set_visible_devices 。可以设置当前程序可见的设备范围(当前程序只会使用自己可见的设备,不可见的设备不会被当前程序使用。使用部分gpu加速。如下面使用gpu设备0,1

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_visible_devices(devices=gpus[0:2], device_type='GPU')

或者使用os来进行控制:使用环境变量 CUDA_VISIBLE_DEVICES 也可以控制程序所使用的GPU。

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "2,3"

第二种:动态申请显存,仅在需要时申请显存空间。

  通过 tf.config.experimental.set_memory_growth 将GPU的显存使用策略设置为“仅在需要时申请显存空间”。以下代码将所有GPU设置为仅在需要时申请显存空间:

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

第三种:限制使用的gpu,并且限制使用的内存大小。
  通过 tf.config.experimental.set_virtual_device_configuration 选项并传入 tf.config.experimental.VirtualDeviceConfiguration 实例,设置TensorFlow固定消耗 GPU:0 的1GB显存

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

3、单GPU模拟多GPU环境

  当我们的本地开发环境只有一个GPU,但却需要编写多GPU的程序在工作站上进行训练任务时,TensorFlow为我们提供了一个方便的功能,可以让我们在本地开发环境中建立多个模拟GPU,从而让多GPU的程序调试变得更加方便。以下代码在实体GPU GPU:0 的基础上建立了两个显存均为2GB的虚拟GPU。

gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_virtual_device_configuration(
     gpus[0],
     [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
      tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])

本文转载自: https://blog.csdn.net/To_be_little/article/details/124438800
版权归原作者 浅冲一下 所有, 如有侵权,请联系我们删除。

“tensorflow如何使用gpu”的评论:

还没有评论