未验证 提交 03b7daa5 编写于 作者: littletomatodonkey's avatar littletomatodonkey 提交者: GitHub

fix distillation model export and pred save (#3869)

上级 073a591c
...@@ -93,6 +93,9 @@ def main(): ...@@ -93,6 +93,9 @@ def main():
for key in config["Architecture"]["Models"]: for key in config["Architecture"]["Models"]:
config["Architecture"]["Models"][key]["Head"][ config["Architecture"]["Models"][key]["Head"][
"out_channels"] = char_num "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 else: # base rec model
config["Architecture"]["Head"]["out_channels"] = char_num config["Architecture"]["Head"]["out_channels"] = char_num
model = build_model(config["Architecture"]) model = build_model(config["Architecture"])
......
...@@ -121,7 +121,7 @@ def main(): ...@@ -121,7 +121,7 @@ def main():
if len(post_result[key][0]) >= 2: if len(post_result[key][0]) >= 2:
rec_info[key] = { rec_info[key] = {
"label": post_result[key][0][0], "label": post_result[key][0][0],
"score": post_result[key][0][1], "score": float(post_result[key][0][1]),
} }
info = json.dumps(rec_info) info = json.dumps(rec_info)
else: else:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册