将训练时的每个epoch和三个loss最终生成一个曲线图

在训练过程中，你可以通过将每个epoch的损失（如主损失、辅助损失以及其他你关心的量）记录下来，生成一个折线图或柱状图来可视化训练进展。这种图表通常被称为训练曲线。以下是基本步骤：

1. **数据收集**：每次调用`train_one_epoch`后，记录下当前epoch号以及对应的三种主要损失（主损失、辅助损失和其他指标）。

data_points = {
    "epoch": [],
    "main_loss": [],
    "aux_loss": [],
    # 添加其他你需要跟踪的指标，例如 "mse_loss" 或 "bpp_loss"
}

2. **处理训练日志**：每当训练完成一个epoch，从上述的日志信息中提取这些数值，并添加到`data_points`字典中。

for _ in range(num_epochs):
    data_points["epoch"].append(epoch)
    data_points["main_loss"].append(out_criterion["loss"].item())
    data_points["aux_loss"].append(aux_loss.item())
    # 如果有其他指标，也添加对应值

3. **创建图表**：训练完成后，你可以使用Python的数据分析库如pandas和matplotlib来绘制曲线图。这里是一个简化的例子：

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame(data_points)
plt.figure(figsize=(10, 6))
plt.plot(df['epoch'], df['main_loss'], label='Main Loss')
plt.plot(df['epoch'], df['aux_loss'], label='Auxiliary Loss')
# 如果有更多指标，添加类似这一行
plt.plot(df['epoch'], df['mse_loss'], label='MSE Loss')

plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss Curves')
plt.legend()
plt.show()

这将生成一个显示各损失随时间变化趋势的图表，可以帮助你评估模型收敛情况以及优化策略的有效性。记得定期保存这些图表，以便于后续分析和调试。