|
|
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
|
|
batch_size_4 = { |
|
"A100": [4.75, 3.26, 3.24, 3.10], |
|
"A10": [13.94, 9.81, 10.01, 9.35], |
|
"T4": [38.81, 30.09, 29.74, 27.55], |
|
"V100": [9.84, 8.16, 8.09, 7.65], |
|
"3090": [10.04, 7.82, 7.89, 7.47], |
|
"3090TI": [9.07, 7.14, 7.15, 6.81], |
|
} |
|
|
|
batch_size_16 = { |
|
"A100": [18.95, 13.57, 13.67, 12.25], |
|
"A10": [0, 37.55, 38.31, 36.81], |
|
"T4": [0, 111.47, 113.26, 106.93], |
|
"V100": [0, 30.29, 29.84, 28.22], |
|
"3090": [0, 29.06, 29.06, 28.2], |
|
"3090TI": [0, 26.1, 26.28, 25.46], |
|
} |
|
|
|
batch_sizes = batch_size_16 |
|
|
|
methods = { |
|
"Vanilla Attention": [x[0] for x in batch_sizes.values()], |
|
"xFormers": [x[1] for x in batch_sizes.values()], |
|
"PyTorch2.0 SDPA": [x[2] for x in batch_sizes.values()], |
|
"SDPA + torch.compile": [x[3] for x in batch_sizes.values()], |
|
} |
|
|
|
x = np.arange(len(batch_size_4)) |
|
width = 0.15 |
|
multiplier = 0 |
|
|
|
fig, ax = plt.subplots(constrained_layout=True) |
|
|
|
for attribute, measurement in methods.items(): |
|
offset = width * multiplier |
|
rects = ax.bar(x + offset, measurement, width, label=attribute) |
|
ax.bar_label(rects, padding=3) |
|
multiplier += 1 |
|
|
|
|
|
ax.set_ylabel('Time (s)') |
|
ax.set_title('Inference Speed at Batch Size=4') |
|
ax.set_xticks(x + width, batch_size_4) |
|
ax.legend(loc='upper left') |
|
ax.set_ylim(0, 250) |
|
|
|
plt.show() |
|
|