fix starnet export (#1850)

* fix starnet export

* fix bug

* add note
This commit is contained in:
zhoujun 2021-01-28 18:16:54 +08:00 committed by GitHub
parent 25bf92295f
commit b0d1dca688
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 2 deletions

View File

@ -52,9 +52,16 @@ def main():
save_path = '{}/inference'.format(config['Global']['save_inference_dir'])
infer_shape = [3, -1, -1]
infer_shape = [3, -1, -1]
if config['Architecture']['model_type'] == "rec":
infer_shape = [3, 32, -1]
infer_shape = [3, 32, -1] # for rec model, H must be 32
if 'Transform' in config['Architecture'] and config['Architecture'][
'Transform'] is not None and config['Architecture'][
'Transform']['name'] == 'TPS':
logger.info(
'When there is tps in the network, variable length input is not supported, and the input size needs to be the same as during training'
)
infer_shape[-1] = 100
model = to_static(
model,