import pandas as pd
import matplotlib.pyplot as plt

# 读取CSV文件
data = pd.read_csv('training_data.csv')

# 创建图表
plt.figure(figsize=(12, 9))

# 绘制梯度范数变化曲线
plt.subplot(3, 1, 1) # 修改为 3行1列的第1个
plt.plot(data['Step'], data['grad_norm'], label='Gradient Norm')
plt.xlabel('Step')
plt.ylabel('Gradient Norm')
plt.legend()

# 绘制损失值变化曲线
plt.subplot(3, 1, 2) # 修改为 3行1列的第2个
plt.plot(data['Step'], data['loss'], label='Loss', color='orange')
plt.xlabel('Step')
plt.ylabel('Loss')
plt.legend()

# 绘制学习率变化曲线
plt.subplot(3, 1, 3) # 修改为 3行1列的第3个
plt.plot(data['Step'], data['learning_rate'], label='Learning Rate', color='green')
plt.xlabel('Step')
plt.ylabel('Learning Rate')
plt.legend()

# 调整布局并保存图片
plt.tight_layout()
plt.savefig('training_metrics.png')
plt.show()