tensorflow2.x 下 eager mode的训练流程
程序员文章站
2024-01-19 08:09:16
...
此前一直用 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)