使用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