TensorFlow2.0中图运算模式和GPU调用规则-CSDN博客

本文链接：https://blog.csdn.net/hallobike/article/details/112400229

图运算模式（Graph Execution模式）

在TensorFlow2.0中，我们可以使用@tf.function装饰器实现Graph Execution，从而将模型转换为易于部署且高性能的TensorFlow图模型。

@tf.function装饰器使用静态编译将函数内的代码转换为计算图。

@tf.function对函数内可使用的语句有一定的限制（仅支持python语言的一个子集），且需要函数内的操作本身能够被构建为计算图。

建议在函数内只使用TensorFlow的原生操作，不要使用过于复杂的Python语句，函数参数只包括TensorFlow张量或Numpy数组，并最好是能够按照计算图的思想去构建函数。

@tf.function可以带来一定的性能提升。一般当模型由较多小的操作组成的时候，@tf.function带来的提升效果较大。

当被@tf.function修饰的函数第一次被调用的时候，进行以下操作：

（1）在Eager Execution模式关闭的环境下，函数内的代码依次运行。也就是说，每个tf.方法都只是定义了计算节点，而没有进行任何实质的计算。这与TensorFlow1.X的Graph Execution是一致的。

（2）使用AutoGraph将函数中的Python控制流语句转换成TensorFlow计算图中的对应节点。（比如说while和for语句转换为tf.while，if语句转换为tf.cond等等）

（3）基于上面两步，建立函数内代码的计算图表示。（为了保证图的计算顺序，图中还会自动加入一些tf.control_dependencies节点）

（4）运行一次这个计算图

（5）基于函数的名字和输入的函数参数的类型生成一个哈希值，并将建立的计算图缓存到一个哈希表中。

（6）在被@tf.function修饰的函数之后再次被再次被调用的时候，根据函数名和输入的函数参数的类型计算哈希值，检查哈希表中是否已经有了对应计算图的缓存。如果是，则直接使用已缓存的计算图，否则重新按上述步骤建立计算图。

综上：
AutoGraph起到了类似编译器的作用，能够帮助我们通过更加自然的Python控制流轻松地构建带有条件的/循环的计算图，而无需手动使用TensorFlow的API进行构建。

tf.function使用

当定义多个函数实现不同的运算时，仅需要在最后调用的函数上添加@tf.function装饰器即可。这样所有的运算节点都会被编译。

GPU配置与使用策略

当前程序只会使用自己可见的设备，不可见的设备不会被当前程序使用，因此可以设置当前程序可见的设备范围，来规范我们使用的GPU。

设置当前程序可见的设备范围

使用环境变量CUDA_VISIBLE_DEVICES也可以控制程序所使用的GPU

import os

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

即可指定程序只在显卡2,3上运行。

设置显存的使用策略

默认情况下，TensorFlow将使用几乎所有可用的显存，以避免内存碎片化所带来的性能损失。TensorFlow提供两种显存使用策略，让我们能够更灵活地控制程序的显存使用方式：

（1）仅在需要时申请显存空间（程序初始运行时消耗很少的显存，随着程序的运行而动态申请显存）

（2）限制消耗固定大小的显存（程序不会超出限定的显存大小，若超出就报错）

# -*- coding: UTF-8 -*-
"""
Author: LGD
FileName: gpu_usage_strategy
DateTime: 2020/12/29 10:58 
SoftWare: PyCharm
"""
import tensorflow as tf

# gpus = tf.config.experimental.list_logical_devices(device_type='GPU')
# print(gpus)
#
# cpus = tf.config.experimental.list_logical_devices(device_type='CPU')
# print(cpus)

# tf.config.experimental.set_visible_devices(device=gpus[0:2], device_type='gpu')
# 仅在需要时申请显存空间
gpus = tf.config.experimental.list_logical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(device=gpu)

# 设置TensorFlow固定消耗GPU：0的1GB显存：
gpus = tf.config.experimental.list_logical_devices(device_type='GPU')
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]
)