File size: 1,356 Bytes
e4e03fd |
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 |
import os
import subprocess
def convert_onnx_to_fp16_corrected(input_dir):
# List all files in the input directory
files = os.listdir(input_dir)
# Filter out files with the .onnx extension
onnx_files = [file for file in files if file.endswith('.onnx')]
# Iterate over each ONNX file to convert it to FP16
for onnx_file in onnx_files:
# Split the file name to insert '.fp16' before the last underscore
parts = onnx_file.rsplit('_', 1)
if len(parts) == 2:
# If there is at least one underscore, insert '.fp16' before the last part
output_model_name = f"{parts[0]}.fp16_{parts[1]}"
else:
# If there's no underscore, just replace '.onnx' with '.fp16.onnx'
output_model_name = onnx_file.replace('.onnx', '.fp16.onnx')
# Construct the command to run the conversion script
command = [
'python3', 'convert_to_fp16.py',
'--input_model', os.path.join(input_dir, onnx_file),
'--output_model', os.path.join(input_dir, output_model_name)
]
# Execute the command
subprocess.run(command, check=True)
if __name__ == '__main__':
# Assuming the current directory is the input directory
input_dir = os.getcwd()
convert_onnx_to_fp16_corrected(input_dir)
|