Merge 16664ea31b into 709d5e0d9d
swenkel
1 year ago
committed by
GitHub
commit
58d9f09c86
@ -0,0 +1,183 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from timm.models import create_model
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
|
||||
def argument_parser() -> dict:
|
||||
"""Argument Parser
|
||||
|
||||
Returns
|
||||
-------
|
||||
config : dict
|
||||
Python dict containing all settings
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--single_model',
|
||||
action='store_true',
|
||||
help='Flag to converts a single model \
|
||||
(use full path + filename as \
|
||||
--folder_path; input dimensions \
|
||||
(--input_dim) and number of \
|
||||
classes (--num_classes) are required)')
|
||||
parser.add_argument('--model_name',
|
||||
default=None,
|
||||
type=str,
|
||||
help='model name (e.g. mobilenetv3_large_100)\
|
||||
required if single_model')
|
||||
parser.add_argument('--folder_path',
|
||||
type=str,
|
||||
help='Path of folder containing training \
|
||||
output (incl. yaml files)')
|
||||
parser.add_argument('--best_model_only',
|
||||
action='store_true',
|
||||
help='Flag to export model_best.pth.tar in \
|
||||
training folder only')
|
||||
parser.add_argument('--last_model_only',
|
||||
action='store_true',
|
||||
help='Flag to export last.pth.tar in \
|
||||
training folder only')
|
||||
parser.add_argument('--input_dim',
|
||||
default=None,
|
||||
type=int,
|
||||
help='Input shape of model assuming square\
|
||||
input (example: --input_dim 224); \
|
||||
if not provided, input dim is derived\
|
||||
from folder name')
|
||||
parser.add_argument('--num_classes',
|
||||
type=int,
|
||||
default=None,
|
||||
help='Number of output classes')
|
||||
args = parser.parse_args()
|
||||
config = {}
|
||||
config['single_model'] = args.single_model
|
||||
config['best_model_only'] = args.best_model_only
|
||||
config['last_model_only'] = args.last_model_only
|
||||
if config['best_model_only'] and config['last_model_only']:
|
||||
raise Exception('invalid options: best_model_only and\
|
||||
last_model_only - choose one')
|
||||
config['folder_path'] = args.folder_path
|
||||
|
||||
if not config['single_model']:
|
||||
if config['folder_path'][-1] != '/':
|
||||
config['folder_path'] += '/'
|
||||
if not os.path.exists(config['folder_path']):
|
||||
raise Exception("Folder containing training output\
|
||||
does not exist")
|
||||
if not os.path.exists(config['folder_path']+'args.yaml'):
|
||||
raise Exception('args.yaml does not exist in folder')
|
||||
else:
|
||||
# derive data from folder name
|
||||
# training folder name format:
|
||||
# YYYYMMDD-hh:mm:ss-modelName_epochs-inputSize
|
||||
folder_name = config['folder_path'].split('/')[-2]
|
||||
timm_train_data = folder_name.split('-')
|
||||
config['input_dim'] = int(timm_train_data[3])
|
||||
config['timm_cfg'] = yaml.safe_load(open(
|
||||
config['folder_path']+'args.yaml','r'))
|
||||
config['timm_cfg_keys'] = tuple(config['timm_cfg'].keys())
|
||||
if 'model' in config['timm_cfg_keys']:
|
||||
config['model_name'] = config['timm_cfg']['model']
|
||||
else:
|
||||
config['model_name'] = timm_train_data[2]
|
||||
files_in_folder = os.listdir(config['folder_path'])
|
||||
config['models_in_folder'] = \
|
||||
[file for file in files_in_folder if '.pth.tar' in file]
|
||||
if len(config['models_in_folder']) == 0:
|
||||
raise Exception('No models found in folder')
|
||||
if config['last_model_only']:
|
||||
if not 'last.pth.tar' in config['models_in_folder']:
|
||||
raise Exception('model last.pth.tar not found')
|
||||
else:
|
||||
config['models_in_folder'] = ['last.pth.tar']
|
||||
if config['best_model_only']:
|
||||
if not 'model_best.pth.tar' in config['models_in_folder']:
|
||||
raise Exception('model model_best.pth.tar not found')
|
||||
else:
|
||||
config['models_in_folder'] = ['model_best.pth.tar']
|
||||
else:
|
||||
if not os.path.exists(config['folder_path']):
|
||||
raise Exception("Input model not found")
|
||||
config['num_classes'] = args.num_classes
|
||||
if not config['single_model']:
|
||||
if config['num_classes'] is None:
|
||||
if 'num_classes' in config['timm_cfg_keys']:
|
||||
if config['timm_cfg']['num_classes'] is None:
|
||||
raise Exception('Number of classes required')
|
||||
else:
|
||||
config['num_classes'] = config['timm_cfg']['num_classes']
|
||||
else:
|
||||
if config['num_classes'] is None:
|
||||
raise Exception('Number of classes required')
|
||||
if args.input_dim is None:
|
||||
if config['single_model']:
|
||||
raise Exception('Input dimension missing')
|
||||
else:
|
||||
config['input_dim'] = args.input_dim
|
||||
if args.model_name is not None:
|
||||
config['model_name'] = args.model_name
|
||||
else:
|
||||
if config['single_model']:
|
||||
raise Exception('Model name missing')
|
||||
return config
|
||||
|
||||
|
||||
def convert_model(model_in : str,
|
||||
model_out : str,
|
||||
config : dict):
|
||||
"""Loads timm model file and converts it to onnx
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model_in : str
|
||||
filepath of input model
|
||||
model_out : str
|
||||
filepath of converted model
|
||||
config : dict
|
||||
Python dict containing all settings
|
||||
|
||||
Returns
|
||||
-------
|
||||
None
|
||||
"""
|
||||
random_input = torch.randn(1,
|
||||
3,
|
||||
config['input_dim'],
|
||||
config['input_dim'],
|
||||
requires_grad=True)
|
||||
model = create_model(config['model_name'],
|
||||
checkpoint_path=model_in,
|
||||
exportable=True,
|
||||
num_classes=config['num_classes'])
|
||||
model.eval()
|
||||
torch.onnx.export(model,
|
||||
random_input,
|
||||
model_out,
|
||||
export_params=True,
|
||||
opset_version=13,
|
||||
do_constant_folding=True,
|
||||
input_names=['input'],
|
||||
output_names=['output'])
|
||||
|
||||
|
||||
def main():
|
||||
config = argument_parser()
|
||||
print('=' * 72)
|
||||
print('Exporting models to ONNX')
|
||||
if not config['single_model']:
|
||||
for model_file in config['models_in_folder']:
|
||||
print('Converting '+model_file+' to '+
|
||||
str(model_file.split('.')[0])+'.onnx')
|
||||
model_in = config['folder_path']+model_file
|
||||
model_out = model_in.replace('pth.tar', 'onnx')
|
||||
convert_model(model_in, model_out, config)
|
||||
else:
|
||||
model_in = config['folder_path']
|
||||
model_out = model_in.replace('pth.tar', 'onnx')
|
||||
convert_model(model_in, model_out, config)
|
||||
print('=' * 72)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Loading…
Reference in new issue