Tensorflow 中自定义训练循环中的 If-Else 语句

Question

我创建了一个模型类，它是 keras.Model 的子类。在训练模型时，我想在一些时期后改变损失函数的权重。为了做到这一点，我为我的模型创建了布尔变量，表明模型应该使用额外的损失函数开始训练。我添加了一个伪代码，主要显示我想要实现的目标。

class MyModel(keras.Model):
    self.start_loss_2 = False


def train_step(self):
    # Check if training with loss_2 started 
    weight_loss_2 = 0.0
    if self.start_loss_2:
        weight_loss_2 = 0.5

    # Pass the data through model
    # Calculate two loss values
    total_loss = loss_1 + weight_loss_2 * loss_2
    # Calculate gradients with tf.Tape
    # Update variables


# This is called via Callback after each epoch
def epoch_finised(epoch_num):
    if epoch_num > START_LOSS_2:
        self.start_loss_2 = True

我的问题是：

• 是否可以使用 if-else 语句，其值在一段时间后发生变化？如果不是，怎么能做到这一点？

Answer 1

是的。您可以创建一个tf.Variable，然后根据一些训练标准为其分配一个新值。

示例：

import numpy as np
import tensorflow as tf


# simple toy network
x_in = tf.keras.Input((10))
x = tf.keras.layers.Dense(25)(x_in)
x_out = tf.keras.layers.Dense(1)(x)

# model
m = tf.keras.Model(x_in, x_out)

# fake data
X = tf.random.normal((100, 10))
y0 = tf.random.normal((100, ))
y1 = tf.random.normal((100, ))

# optimizer
m_opt = tf.keras.optimizers.Adam(1e-2)

# prep data
ds = tf.data.Dataset.from_tensor_slices((X, y0, y1))
ds = ds.repeat().batch(5)
train_iter = iter(ds)

# toy loss function that uses a weight
def loss_fn(y_true0, y_true1, y_pred, weight):
    mse = tf.keras.losses.MSE
    mse_0 = tf.math.reduce_mean(mse(y_true0, y_pred))
    mse_1 = tf.math.reduce_mean(mse(y_true1, y_pred))
    return mse_0 + weight * mse_1
  
NUM_EPOCHS = 4
NUM_BATCHES_PER_EPOCH = 10
START_NEW_LOSS_AT_GLOBAL_STEP = 20

# the weight variable set to 0 initially and then
# will be changed after a certain number of steps
# (or some other training criteria)
w = tf.Variable(0.0, trainable=False)

for epoch in range(NUM_EPOCHS):
    losses = []
    for batch in range(NUM_BATCHES_PER_EPOCH):
        X_train, y0_train, y1_train = next(train_iter)
        with tf.GradientTape() as tape:
            y_hat = m(X_train)
            loss = loss_fn(y0_train, y1_train, y_hat, w)
            losses.append(loss)
    
        m_vars = m.trainable_variables
        m_grads = tape.gradient(loss, m_vars)
        m_opt.apply_gradients(zip(m_grads, m_vars))
    
    print(f"epoch: {epoch}\tloss: {np.mean(losses):.4f}")
    losses = []

    # if the criteria is met assign a huge number to see if the
    # loss spikes up
    if (epoch + 1) * (batch + 1) >= START_NEW_LOSS_AT_GLOBAL_STEP:
        w.assign(10000.0)

# epoch: 0  loss: 1.8226
# epoch: 1  loss: 1.1143
# epoch: 2  loss: 8788.2227    <= looks like assign worked
# epoch: 3  loss: 10999.5449