使用单GPU训练模型 - 文章 - 开发者社区

深度学习的训练过程常常非常耗时，一个模型训练几个小时是家常便饭，训练几天也是常有的事情，有时候甚至要训练几十天。

训练过程的耗时主要来自于两个部分，一部分来自数据准备，另一部分来自参数迭代。

当数据准备过程还是模型训练时间的主要瓶颈时，我们可以使用更多进程来准备数据。

当参数迭代过程成为训练时间的主要瓶颈时，我们通常的方法是应用GPU或者Google的TPU来进行加速。

详见《用GPU加速Keras模型——Colab免费GPU使用攻略》

https://zhuanlan.zhihu.com/p/68509398

本篇我们介绍使用单GPU训练模型的方法，后面两篇分别介绍使用多GPU和使用TPU训练模型的方法。

〇，GPU配置

无论是内置fit方法，还是自定义训练循环，从CPU切换成单GPU训练模型都是非常方便的，无需更改任何代码。 当存在可用的GPU时，如果不特意指定device， tensorflow会自动优先选择使用GPU来创建张量和执行张量计算 。

但如果是在公司或者学校实验室的服务器环境，存在多个GPU和多个使用者时，为了不让单个同学的任务占用全部GPU资源导致其他同学无法使用（tensorflow默认获取全部GPU的全部内存资源权限，但实际上只使用一个GPU的部分资源），我们通常会在开头增加以下几行代码以控制每个任务使用的GPU编号和显存大小，以便其他同学也能够同时训练模型。

在Colab笔记本中：修改->笔记本设置->硬件加速器中选择 GPU

注：以下代码只能在Colab 上才能正确执行。

可通过以下Colab分享链接测试效果《tf_单GPU》：

https://colab.research.google.com/drive/1r5dLoeJq5z01sU72BX2M5UiNSkuxsEFe


          
%tensorflow_version 2.x  
import tensorflow as tf  
print(tf.__version__)


          
from tensorflow.keras import *   
  
#打印时间分割线  
@tf.function  
def printbar():  
    ts = tf.timestamp()  
    today_ts = ts%(24*60*60)  
  
    hour = tf.cast(today_ts//3600+8,tf.int32)%tf.constant(24)  
    minite = tf.cast((today_ts%3600)//60,tf.int32)  
    second = tf.cast(tf.floor(today_ts%60),tf.int32)  
  
    def timeformat(m):  
        if tf.strings.length(tf.strings.format("{}",m))==1:  
            return(tf.strings.format("0{}",m))  
        else:  
            return(tf.strings.format("{}",m))  
  
    timestring = tf.strings.join([timeformat(hour),timeformat(minite),  
                timeformat(second)],separator = ":")  
    tf.print("=========="*8,end = "")  
    tf.print(timestring)

增加以下几行代码控制GPU使用量


          
gpus = tf.config.list_physical_devices("GPU")  
  
if gpus:  
    gpu0 = gpus[0] #如果有多个GPU，仅使用第0个GPU  
    tf.config.experimental.set_memory_growth(gpu0, True) #设置GPU显存用量按需使用  
    # 或者也可以设置GPU显存为固定使用量(例如：4G)  
    #tf.config.experimental.set\_virtual\_device\_configuration(gpu0,  
    #    [tf.config.experimental.VirtualDeviceConfiguration(memory\_limit=4096)])   
    tf.config.set_visible_devices([gpu0],"GPU")

比较GPU和CPU的计算速度


          
printbar()  
with tf.device("/gpu:0"):  
    tf.random.set_seed(0)  
    a = tf.random.uniform((10000,100),minval = 0,maxval = 3.0)  
    b = tf.random.uniform((100,100000),minval = 0,maxval = 3.0)  
    c = a@b  
    tf.print(tf.reduce_sum(tf.reduce_sum(c,axis = 0),axis=0))  
printbar()

picture.image


          
printbar()  
with tf.device("/cpu:0"):  
    tf.random.set_seed(0)  
    a = tf.random.uniform((10000,100),minval = 0,maxval = 3.0)  
    b = tf.random.uniform((100,100000),minval = 0,maxval = 3.0)  
    c = a@b  
    tf.print(tf.reduce_sum(tf.reduce_sum(c,axis = 0),axis=0))  
printbar()

picture.image

一，准备数据


          
MAX_LEN = 300  
BATCH_SIZE = 32  
(x_train,y_train),(x_test,y_test) = datasets.reuters.load_data()  
x_train = preprocessing.sequence.pad_sequences(x_train,maxlen=MAX_LEN)  
x_test = preprocessing.sequence.pad_sequences(x_test,maxlen=MAX_LEN)  
  
MAX_WORDS = x_train.max()+1  
CAT_NUM = y_train.max()+1  
  
ds_train = tf.data.Dataset.from_tensor_slices((x_train,y_train)) \  
          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \  
          .prefetch(tf.data.experimental.AUTOTUNE).cache()  
  
ds_test = tf.data.Dataset.from_tensor_slices((x_test,y_test)) \  
          .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \  
          .prefetch(tf.data.experimental.AUTOTUNE).cache()

二，定义模型


          
tf.keras.backend.clear_session()  
  
def create\_model():  
  
    model = models.Sequential()  
  
    model.add(layers.Embedding(MAX_WORDS,7,input_length=MAX_LEN))  
    model.add(layers.Conv1D(filters = 64,kernel_size = 5,activation = "relu"))  
    model.add(layers.MaxPool1D(2))  
    model.add(layers.Conv1D(filters = 32,kernel_size = 3,activation = "relu"))  
    model.add(layers.MaxPool1D(2))  
    model.add(layers.Flatten())  
    model.add(layers.Dense(CAT_NUM,activation = "softmax"))  
    return(model)  
  
model = create_model()  
model.summary()

picture.image

三，训练模型


          
optimizer = optimizers.Nadam()  
loss_func = losses.SparseCategoricalCrossentropy()  
  
train_loss = metrics.Mean(name='train\_loss')  
train_metric = metrics.SparseCategoricalAccuracy(name='train\_accuracy')  
  
valid_loss = metrics.Mean(name='valid\_loss')  
valid_metric = metrics.SparseCategoricalAccuracy(name='valid\_accuracy')  
  
@tf.function  
def train\_step(model, features, labels):  
    with tf.GradientTape() as tape:  
        predictions = model(features,training = True)  
        loss = loss_func(labels, predictions)  
    gradients = tape.gradient(loss, model.trainable_variables)  
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))  
  
    train_loss.update_state(loss)  
    train_metric.update_state(labels, predictions)  
  
@tf.function  
def valid\_step(model, features, labels):  
    predictions = model(features)  
    batch_loss = loss_func(labels, predictions)  
    valid_loss.update_state(batch_loss)  
    valid_metric.update_state(labels, predictions)  
  
  
def train\_model(model,ds\_train,ds\_valid,epochs):  
    for epoch in tf.range(1,epochs+1):  
  
        for features, labels in ds_train:  
            train_step(model,features,labels)  
  
        for features, labels in ds_valid:  
            valid_step(model,features,labels)  
  
        logs = 'Epoch={},Loss:{},Accuracy:{},Valid Loss:{},Valid Accuracy:{}'  
  
        if epoch%1 ==0:  
            printbar()  
            tf.print(tf.strings.format(logs,  
            (epoch,train_loss.result(),train_metric.result(),valid_loss.result(),valid_metric.result())))  
            tf.print("")  
  
        train_loss.reset_states()  
        valid_loss.reset_states()  
        train_metric.reset_states()  
        valid_metric.reset_states()  
  
train_model(model,ds_train,ds_test,10)

picture.image

猜你喜欢❤️：

《30天吃掉那只 TensorFlow2.0 》开篇辞

公众号后台回复关键字： tensorflow ，获取本书github项目源码和对应数据集！

picture.image