From 03b7daa5be911c0394f6742b5087576c3795143a Mon Sep 17 00:00:00 2001 From: littletomatodonkey Date: Wed, 1 Sep 2021 19:36:16 +0800 Subject: [PATCH] fix distillation model export and pred save (#3869) --- tools/export_model.py | 3 +++ tools/infer_rec.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/export_model.py b/tools/export_model.py index 785aca10..cae87aca 100755 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -93,6 +93,9 @@ def main(): for key in config["Architecture"]["Models"]: config["Architecture"]["Models"][key]["Head"][ "out_channels"] = char_num + # just one final tensor needs to to exported for inference + config["Architecture"]["Models"][key][ + "return_all_feats"] = False else: # base rec model config["Architecture"]["Head"]["out_channels"] = char_num model = build_model(config["Architecture"]) diff --git a/tools/infer_rec.py b/tools/infer_rec.py index 09f5a0c7..cf49348f 100755 --- a/tools/infer_rec.py +++ b/tools/infer_rec.py @@ -121,7 +121,7 @@ def main(): if len(post_result[key][0]) >= 2: rec_info[key] = { "label": post_result[key][0][0], - "score": post_result[key][0][1], + "score": float(post_result[key][0][1]), } info = json.dumps(rec_info) else: -- GitLab