tensorflow2.0gpu版本 自定义模型提高GPU利用率

使用tensorflow2.0进行cifar100数据集的图像分类时发现GPU的使用率始终不高，且波动较大，在网上查询资料后尝试从数据处理和使用@tf.function图执行模式入手，其中模型结构使用ResNet18，机器配置显卡为2070super，CPU为AMD2600。
参考资料：
ResNet模型参考：https://blog.csdn.net/abc13526222160/article/details/90057121
数据处理参考：https://tf.wiki/zh_hans/basic/tools.html#tf-data
图执行模式参考：https://mp.weixin.qq.com/s?__biz=MzU1OTMyNDcxMQ==&mid=2247487599&idx=1&sn=13a53532ad1d2528f0ece4f33e3ae143&chksm=fc185b27cb6fd2313992f8f2644b0a10e8dd7724353ff5e93a97d121cd1c7f3a4d4fcbcb82e8&scene=21#wechat_redirect

数据处理：
模型常规训练时会发生cpu在处理数据而gpu在等待cpu传输数据过来，cpu和gpu的串行处理流程会发生空载现象，造成计算资源的浪费，GPU利用率波动较大，就是GPU处于空载的表现，这种情况下可以使用tf.data一些方法提高数据处理效率。
tf.data 的数据集对象为我们提供了 Dataset.prefetch() 方法，使得我们可以让数据集对象 Dataset 在训练时预取出若干个元素，使得在 GPU 训练的同时 CPU 可以准备数据，提升训练流程的效率，这里prefetch()的参数 buffer_size 一般设置为 tf.data.experimental.AUTOTUNE 从而由 TensorFlow 自动选择合适的数值，
常规训练过程：

使用Dataset.prefetch() 方法进行数据预加载后的训练流程：

Dataset.map(f) 是转换函数f映射到数据集每一个元素，Dataset.map() 可以利用多 CPU 资源，充分利用多核心的优势对数据进行并行化变换，num_parallel_calls 设置为 tf.data.experimental.AUTOTUNE 以让 TensorFlow 自动选择合适的数值，数据转换过程多进行执行，设置num_parallel_calls 参数能发挥cpu多核心的优势，

tf.data.Dataset 类其他最常用的数据集预处理方法如：
Dataset.shuffle(buffer_size) ：将数据集打乱（设定一个固定大小的缓冲区（Buffer），取出前 buffer_size 个元素放入，并从缓冲区中随机采样，采样后的数据用后续数据替换）；
Dataset.batch(batch_size) ：将数据集分成批次，即对每 batch_size 个元素，使用 tf.stack() 在第 0 维合并，成为一个元素
通过使用tf.data.Dataset的数据集预处理方法，有效的减低了GPU利用率的波动问题，稳定在60%~78%的GPU使用率，处理代码如下：

ds_train=tf.data.Dataset.from_tensor_slices((x,y)).shuffle(5000).batch(batchs)\
.map(preprocess,num_parallel_calls=tf.data.experimental.AUTOTUNE).prefetch(tf.data.experimental.AUTOTUNE)

图执行模式：
如何进一步提高利用率就自然考虑到tf2提供的图执行模式了，TensorFlow 2 提供了 tf.function 模块，结合 AutoGraph 机制，仅需加入一个简单的 @tf.function 修饰符，就能轻松将模型以图执行模式高效运行。
将我们希望以图执行模式运行的代码封装在一个函数内，并在函数前加上 @tf.function 即可，@tf.function使用要注意，1，函数内尽量使用tf原生函数，如print改为tf.print，range改为tf.range,2，避免在@tf.function修饰的函数内部定义tf.Variable.，3，被@tf.function修饰的函数不可修改该函数外部的Python列表或字典等数据结构变量。
由于本人只是初学者，对于图执行的原理了解不深，网上没找到完整的示范样例，在参考各篇文章及官网资料后尝试如下，tf2通过继承layers.Layer和 tf.keras.Model 这个 Python 类来定义自己的层和模型，在继承类中，我们需要重写 __init__() （构造函数，初始化）和 call(input) （模型调用）两个方法，在call方法前面增加@tf.function修饰，同时将求导的计算也封装成函数，以@tf.function修饰。
在使用图执行模式后，GPU利用率能达到99%，几乎是满载的运行，每一个epoch训练时间从44秒缩短到了26秒，训练效率提高了约40%，

模型完整代码如下：

import tensorflow as tf
import tensorflow.keras as keras 
from tensorflow.keras import layers,Sequential,regularizers

#定义一个3x3卷积
def regularized_padded_conv(*args,**kwargs):
    return layers.Conv2D(*args,**kwargs,padding='same',kernel_regularizer=regularizers.l2(5e-5),
                         use_bias=False,kernel_initializer='glorot_normal')
            

#定义 Basic Block 模块。对应Resnet18和Resnet34
class BasecBlock(layers.Layer):
    expansion=1
    
    def __init__(self,in_channels,out_channels,stride=1):
        super(BasecBlock,self).__init__()
        #1
        self.conv1=regularized_padded_conv(out_channels,kernel_size=3,strides=stride)
        self.bn1=layers.BatchNormalization()
        
        #2
        self.conv2=regularized_padded_conv(out_channels,kernel_size=3,strides=1)
        self.bn2=layers.BatchNormalization()
        
        #3
        if stride!=1 or in_channels!=self.expansion * out_channels:
            self.shortcut= Sequential([regularized_padded_conv(self.expansion * out_channels,kernel_size=3,strides=stride),
                                      layers.BatchNormalization()])
        else :
            self.shortcut= lambda x ,_ : x 
        
    @tf.function
    def call(self,inputs,training=False):
        
        x=self.conv1(inputs)
        x=self.bn1(x,training=training)
        x=tf.nn.relu(x)
        
        x=self.conv2(x)
        x=self.bn2(x,training=training)
        
        x_short=self.shortcut(inputs,training)
        
        x=x+x_short
        out=tf.nn.relu(x)
        
        return out