YOLOv5训练损失、精度、mAP绘图功能 | 支持多模型对比,多结果绘在一个图片(消融实验、科研必备
一、本文介绍
本文给大家带来的是YOLOv5系列的绘图功能,我将向大家介绍YOLO系列的绘图功能。我们在进行实验时,经常需要比较多个结果,针对这一问题,我写了点代码来解决这个问题,它可以根据训练结果绘制损失(loss)和mAP(平均精度均值)的对比图。这个工具不仅支持多个文件的对比分析,还允许大家在现有代码的基础上进行修,从而达到数据可视化的功能,大家也可以将对比图来放在论文中进行对比也是非常不错的选择。
效果图
二、核心代码
import os
import pandas as pd
import matplotlib.pyplot as plt
def plot_metrics_and_loss(experiment_names, metrics_info, loss_info, metrics_subplot_layout, loss_subplot_layout,
metrics_figure_size=(15, 10), loss_figure_size=(15, 10), base_directory='runs/train'):
# Plot metrics
plt.figure(figsize=metrics_figure_size)
for i, (metric_name, title) in enumerate(metrics_info):
plt.subplot(*metrics_subplot_layout, i + 1)
for name in experiment_names:
file_path = os.path.join(base_directory, name, 'results.csv')
data = pd.read_csv(file_path)
column_name = [col for col in data.columns if col.strip() == metric_name][0]
plt.plot(data[column_name], label=name)
plt.xlabel('Epoch')
plt.title(title)
plt.legend()
plt.tight_layout()
metrics_filename = 'metrics_curves.png'
plt.savefig(metrics_filename)
plt.show()
# Plot loss
plt.figure(figsize=loss_figure_size)
for i, (loss_name, title) in enumerate(loss_info):
plt.subplot(*loss_subplot_layout, i + 1)
for name in experiment_names:
file_path = os.path.join(base_directory, name, 'results.csv')
data = pd.read_csv(file_path)
column_name = [col for col in data.columns if col.strip() == loss_name][0]
plt.plot(data[column_name], label=name)
plt.xlabel('Epoch')
plt.title(title)
plt.legend()
plt.tight_layout()
loss_filename = 'loss_curves.png'
plt.savefig(loss_filename)
plt.show()
return metrics_filename, loss_filename
# Metrics to plot
metrics_info = [
('metrics/precision', 'Precision'),
('metrics/recall', 'Recall'),
('metrics/mAP_0.5', 'mAP at IoU=0.5'),
('metrics/mAP_0.5:0.95', 'mAP for IoU Range 0.5-0.95')
]
# Loss to plot
loss_info = [
('train/box_loss', 'Training Box Loss'),
('train/cls_loss', 'Training Classification Loss'),
('train/obj_loss', 'Training OBJ Loss'),
('val/box_loss', 'Validation Box Loss'),
('val/cls_loss', 'Validation Classification Loss'),
('val/obj_loss', 'Validation obj Loss')
]
# Plot the metrics and loss from multiple experiments
metrics_filename, loss_filename = plot_metrics_and_loss(
experiment_names=['exp40', 'exp38'],
metrics_info=metrics_info,
loss_info=loss_info,
metrics_subplot_layout=(2, 2),
loss_subplot_layout=(2, 3)
)
三、使用方法
使用方式非常简单,我们首先在yolo文件夹根目录下创建一个python文件,将核心代码粘贴进去,代码其中experiment_names这个参数就代表我们的每个训练结果的名字, 我们只需要修改这个即可,我这里就是五个结果进行对比,修改完成之后大家运行该文件即可。