自定义训练循环参数优化：如何每个Epoch仅输出一个epoch_loss？

我针对自定义损失函数实现了一个自定义训练循环，该损失函数还整合了无监督参数优化问题的约束条件。当前训练循环会输出多个epoch_loss（输出数量与ab_train中的元素数量一致），请问如何在每个epoch/迭代中仅得到一个epoch_loss？

train_loss_results = []

num_epochs = 200

def loss(model,ab):
    xyz_pred = model(ab)
    return penalty_function(xyz_pred,ab)

def grad(model,ab,xyz):
    with tf.GradientTape() as tape:
        loss_value = loss(model,ab)
    return loss_value, tape.gradient(loss_value,model.trainable_variables)

opt = tf.keras.optimizers.Adam(learning_rate=0.0001)

for epoch in range(num_epochs):        
    for ab in ab_train:        
        ab = ab.reshape(1,2)
        xyz_pred_us = model(ab) #forward pass
        loss_value,grads = grad(model,ab_train,xyz_pred_us)
        
        opt.apply_gradients(zip(grads, model.trainable_variables))
        epoch_loss = loss_value
        train_loss_results.append(epoch_loss)
    print("Epoch {}: Loss: {}".format(epoch,epoch_loss))

问题根源 #

当前代码中，每次遍历ab_train的单个样本时，都会把当前样本的loss赋值给epoch_loss，最终打印的只是该epoch最后一个样本的loss；同时train_loss_results会记录每个样本的loss，导致每个epoch对应多条loss记录。

解决思路 #

要得到每个epoch的单一loss值，需要在每个epoch内累加所有样本的loss，最终计算平均值（或总和）作为该epoch的总loss。

修改后的代码 #

train_loss_results = []
num_epochs = 200

def loss(model, ab):
    xyz_pred = model(ab)
    return penalty_function(xyz_pred, ab)

def grad(model, ab):  # 移除未使用的xyz参数
    with tf.GradientTape() as tape:
        loss_value = loss(model, ab)
    return loss_value, tape.gradient(loss_value, model.trainable_variables)

opt = tf.keras.optimizers.Adam(learning_rate=0.0001)

for epoch in range(num_epochs):
    epoch_total_loss = 0.0  # 初始化当前epoch总损失
    sample_count = len(ab_train)
    
    for ab in ab_train:        
        ab = ab.reshape(1,2)
        # 修正参数：传入当前单个样本ab，而非整个ab_train数据集
        loss_value, grads = grad(model, ab)
        
        opt.apply_gradients(zip(grads, model.trainable_variables))
        # 累加每个样本的loss（转numpy避免张量累积）
        epoch_total_loss += loss_value.numpy()
    
    # 计算当前epoch的平均损失
    epoch_avg_loss = epoch_total_loss / sample_count
    train_loss_results.append(epoch_avg_loss)
    print("Epoch {}: Loss: {}".format(epoch, epoch_avg_loss))

关键修改说明 #

• 初始化累加变量：每个epoch开始时初始化epoch_total_loss，用于累积该epoch所有样本的损失
• 修正参数传递：原代码错误地将整个ab_train传入grad函数，改为传入当前遍历的单个ab样本
• 计算epoch总loss：遍历完所有样本后，用总损失除以样本数得到平均loss，作为该epoch的最终loss
• 精简冗余代码：移除grad函数中未使用的xyz参数，避免无效参数传递

内容的提问来源于stack exchange，提问作者tehomas