yolov5 保存 loss 变化图片

YOLOv5 是一种流行的物体检测算法，它是 YOLO（You Only Look Once）系列的最新版本。当训练模型时，loss 变化图通常用于可视化训练过程，帮助我们了解模型收敛情况。为了保存每个训练周期的损失变化图像，在训练过程中你可以选择将损失值记录下来，并在每个epoch结束后将其绘制成图表。

以下是基本步骤：

1. **修改训练脚本**：在 `train.py` 或相应的训练函数中，添加代码来记录每次迭代的损失值。这通常是在 `on_epoch_end()` 函数里完成的，可以使用 Python 的内置 `plt` 库来存储损失数据点。

import matplotlib.pyplot as plt

# 在这里添加新的变量来存储每轮的loss
running_loss = []
...

def on_epoch_end(self):
    # 计算平均损失并添加到列表中
    avg_loss = self.loss.item()
    running_loss.append(avg_loss)
    
    # 如果需要，也可以在此处保存损失图像
    if len(running_loss) > 1:  # 避免在第一轮就保存
        save_loss_graph(running_loss)

2. **定义保存图像函数**：
创建一个名为 `save_loss_graph` 的函数，接受列表 `running_loss` 作为输入，然后绘制折线图并保存至指定路径。

def save_loss_graph(loss_history, output_dir='runs'):
    plt.figure()
    plt.plot(loss_history)
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title('Training Loss Curve')
    plt.savefig(os.path.join(output_dir, 'loss_curve.png'))

3. **在适当位置调用保存函数**：
确保在适当的时候调用这个 `save_loss_graph` 函数，比如每个epoch结束后。