如何在Seaborn分组柱状图上方正确标注对应speedup值

问题描述 #

我有如下数据集，想要生成以model为X轴、no_dev为分组依据、time为Y轴的分组柱状图，并且在每个柱状条上方标注对应的speedup值。但运行代码后，出现了speedup标注错位的问题。

数据集 #

no_dev  fw          model          time    speedup
0        8  pytorch    efficientnet   2984.223210  12.802327
1        8  pytorch           vgg16   2583.343883   6.794499
2        8  pytorch        resnet50    442.069308  24.661291
3        4  pytorch    efficientnet   5318.658496   7.183202
4        4  pytorch        resnet50    695.629588  15.672134
5        4  pytorch           vgg16   4796.323589   3.659580
6        2  pytorch        resnet50   1041.627414  10.466314
7        2  pytorch           vgg16   9465.335288   1.854401
8        2  pytorch    efficientnet   9365.883145   4.079167
9        1  pytorch        resnet50  10902.000000   1.000000
10       1  pytorch    efficientnet  38205.000000   1.000000
11       1  pytorch           vgg16  17552.527806   1.000000

现有代码 #

df = df.sort_values(by=['model', 'speedup'], ascending=True)

# Set the figure size
fig, ax = plt.subplots(figsize=(10, 6))

# Set the figure size and create the bar plot grouped by 'model' and 'no_npu'
bars = sns.barplot(data=df, x='model', y='time', hue='no_dev', dodge=True)

# Iterate over the bars and the DataFrame rows simultaneously
idx = 0
speedups = df['speedup'].values
print(speedups)
for bar in bars.patches:
    # Use the 'speedup' value from the DataFrame row for the label
    label = speedups[idx]
    print(label)
    
    # Annotate the bar with the 'speedup' value
    bars.annotate(f'{label:.2f}x',  # Formatting the label as a floating point with 'x' to denote speedup
                  xy=(bar.get_x() + bar.get_width() / 2, bar.get_height()),
                  xytext=(0, 3),  # 3 points vertical offset
                  textcoords="offset points",
                  ha='center', va='bottom')
    idx += 1

# Add labels and title
plt.legend()

问题原因 #

标注错位的核心原因是数据排序顺序与seaborn绘制柱状条的顺序不匹配：

• 代码中按['model', 'speedup']升序排序数据，但seaborn绘制分组柱状图时，会先按X轴的model分组，再按hue参数no_dev的默认升序（1→2→4→8）排列每个分组内的柱子，两者顺序不一致，导致标注对应错误。

解决方法 #

调整数据排序规则，让数据顺序与seaborn的绘制逻辑完全对齐：按model升序，再按no_dev升序排序。同时优化标注逻辑，确保每个柱子对应正确的speedup值。

修正后的代码 #

# 调整排序规则：按model升序，再按no_dev升序，和seaborn绘制顺序一致
df = df.sort_values(by=['model', 'no_dev'], ascending=True)

fig, ax = plt.subplots(figsize=(10, 6))
bars = sns.barplot(data=df, x='model', y='time', hue='no_dev', dodge=True)

# 直接按排序后的df提取speedup值进行标注
for idx, bar in enumerate(bars.patches):
    speedup_val = df['speedup'].iloc[idx]
    bars.annotate(f'{speedup_val:.2f}x',
                  xy=(bar.get_x() + bar.get_width()/2, bar.get_height()),
                  xytext=(0, 3),
                  textcoords='offset points',
                  ha='center', va='bottom')

plt.xlabel('模型')
plt.ylabel('耗时')
plt.title('不同设备数下各模型耗时及加速比')
plt.legend(title='设备数量')
plt.show()

说明 #

• 修正排序后，数据的顺序和seaborn绘制每个柱子的顺序完全一致，标注自然不会错位
• 额外添加了中文标签和标题，让图表更易读

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