tensorflow2.x 下 eager mode的训练流程

最新推荐文章于 2021-06-04 21:53:28 发布

chenxin0215

最新推荐文章于 2021-06-04 21:53:28 发布

阅读量699

点赞数

CC 4.0 BY-SA版权

文章标签： tensorflow pytorch 神经网络深度学习

本文链接：https://blog.youkuaiyun.com/chenxin0215/article/details/112978870

本文作者在项目中从PyTorch切换到TensorFlow 2.x，探讨了使用eager mode进行LSTM模型构建，自定义类以获取多层LSTM的sequence输出和state，以及如何在保持类似PyTorch灵活性的同时，适应model.compile和model.fit的训练过程。作者分享了训练函数中的关键步骤，如初始化状态、动态调整学习率等。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

此前一直用 pytorch，习惯了动态图的灵活和方便调试的特性。tf 2019年以来，比较大的改变便是引入eager mode 方式，来直接调用网络层来搭建深度网络，这一点与pytorch 非常类似。

前段日子项目需求，开始使用tensorflow（哈哈入坑有点晚），也是简单keras demo 开始玩，但是越来越发现 model.compile 和 model.fit 的模型建立和训练过程方便有余灵活性不足（也许是我涉世未深）。

于是沿着使用 pytorch 的思路试着用tf 2.x 的eager 模式。

具体任务和场景不多说。把代码和注释 post 出来，做个记录。

数据和数据加载掠过。

模型定义：没有使用 sequential，发现想要同时拿到多层 lstm 的 sequence 输出和 state 输出，sequential 搭建模型总是出错，于是就直接自定义类（继承 tf.keras.Model）和 forward 方式，这里其实还要加上是否为 training 从而区分训练和测试时模型的表现，

类似pytorch 里 .eval() 和 .train() , 看上去整个流程与pytorch 非常类似。

class LstmModel(tf.keras.Model):
    def __init__(self):
        super(LstmModel, self).__init__()
        self.lstm1 = tf.keras.layers.LSTM(121,return_sequences=True,return_state=True)
        self.lstm2 = tf.keras.layers.LSTM(121,return_sequences=True,return_state=True)
        
    def call(self,inputs,h_i,c_i):
        x,h,c = self.lstm1(inputs,initial_state=(h_i,c_i))
        x = self.lstm2(x,initial_state=(h,c))        
        return x

model = LstmModel()

定义训练函数。重要的注意都写在注释里了。

# 训练函数

train_step_signature = [
    tf.TensorSpec(shape=(1,None,3), dtype=tf.float32),
    tf.TensorSpec(shape=(1,None,3), dtype=tf.float32),
]  # 想放开batch 和 time seq 改成 shape = (None,None,3) 就不行，

# 强制标注 输入数据类型 与 shape 
@tf.function(input_signature=train_step_signature)  
def train_step(x,y):
    with tf.GradientTape() as tape: 
        batch_size = x.shape[0]
        # 初始化 lstm h 与 c 状态
        h = tf.zeros([batch_size,121],dtype=tf.float32)
        c = tf.zeros([batch_size,121],dtype=tf.float32)
        
        predictions = model(x,h,c)        # 一次 forward 模型
        loss = loss_calc(y, predictions)  # 自定义 loss 计算  
        tf.print('loss: %f',loss)         # 要用 tf.print()替代 print  
        
        gradients = tape.gradient(loss, model.trainable_variables)  # 反向求导 类似 torch 里的 backward  
        # clip the gradients 为了控制训练的稳定性， 刻意clip 梯度值
        gradients = [tf.clip_by_value(g,clip_value_min=-10.0, clip_value_max=10.0) for g in gradients ]  
        # 更新变量梯度值， 类似 torch里的 step()
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

训练过程，过程没什么，就是注意动态修改学习的操作，如果是采用model.fit 的训练方式实现需要写 callback 函数，

# 训练过程 
initial_lr = 0.0000075
optimizer = keras.optimizers.Adam(learning_rate = initial_lr)
# 提前指定好 optimizer 
for step in range(1000):
    x,y = next(iterator_seq)
    # update the learning rate
    # 动态改变 学习率的方式 
    lr = optimizer.lr.numpy()
    optimizer.lr.assign(lr*0.995)
    train_step(x,y)