File size: 3,376 Bytes
6a3ad5b |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import onnx
import os
import itertools
import argparse
from onnxconverter_common.float16 import convert_float_to_float16
from onnxruntime.quantization import quantize_dynamic, QuantType
from multiprocessing import Pool
from tabulate import tabulate
def float16(input, output):
model = onnx.load(input)
model_f16 = convert_float_to_float16(model)
onnx.save(model_f16, output)
def qint8(input, output):
quantize_dynamic(input, output, weight_type=QuantType.QInt8)
def quint8(input, output):
quantize_dynamic(input, output, weight_type=QuantType.QUInt8)
def print_table(table):
print(tabulate(table, headers="keys", tablefmt="github"), "\n")
def get_file_mb(path):
try:
stat = os.stat(path)
except FileNotFoundError:
return "N/A"
mb = round(stat.st_size / 1_000_000)
return f"{mb}"
def convert(name, mode, f, markdown):
fname = f.__name__
input = f"models/clip-{name}-{mode}.onnx"
output = f"models/clip-{name}-{mode}-{fname}.onnx"
exists = os.path.exists(output)
if exists:
if not markdown:
print(f"{output} exists")
else:
if not markdown:
print(f"{output} converting")
f(input, output)
if not markdown:
print(f"{output} done")
return [input, output, name, mode, fname, "✅" if exists else "❌"]
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Create variants of converted models')
parser.add_argument(
'--markdown',
action='store_true',
help='Print markdown tables describing the variants'
)
args = parser.parse_args()
names = [
"resnet-50",
"resnet-101",
"resnet-50x4",
"resnet-50x16",
"resnet-50x64",
"resnet-50",
"resnet-50",
"resnet-50",
"vit-base-patch16",
"vit-base-patch32",
"vit-large-patch14",
"vit-large-patch14-336",
]
modes = [
"visual",
"textual"
]
funcs = [
float16,
qint8,
quint8,
]
markdown = args.markdown
if markdown:
print_table({ "Model ID": names })
print_table({ "Mode": modes })
print_table({ "Data Type": [f.__name__ for f in funcs] })
variants = itertools.product(names, modes, funcs, [markdown])
with Pool(8 if not markdown else 1) as p:
variants_table = p.starmap(convert, variants)
if markdown:
# Insert rows for the original models
prev_input = ""
variants_table_with_originals = []
for row in variants_table:
input = row[0]
output = row[1]
if input != prev_input:
prev_input = input
variants_table_with_originals.append(
row[0:1] + row[2:4] + ["float32 (original)", "✅", get_file_mb(input)]
)
file_size = get_file_mb(output)
variants_table_with_originals.append(row[1:] + [file_size])
# Add header
variants_table_with_originals.insert(0, ["Path", "Model ID", "Mode", "Data Type", "Available", "Size (MB)"])
# Print
print(tabulate(variants_table_with_originals, headers="firstrow", tablefmt="github"))
else:
print("done")
|