File size: 1,491 Bytes
d4184d4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
#!/usr/bin/env python3
import matplotlib.pyplot as plt
import numpy as np
batch_size_4 = {
"A100": [4.75, 3.26, 3.24, 3.10], # those values are made up
"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)) # the label locations
width = 0.15 # the width of the bars
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
# Add some text for labels, title and custom x-axis tick labels, etc.
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()
|