提交 a22b8f6b 编写于 作者: W weishengyu

remove neck when useless

上级 934de965
...@@ -25,7 +25,8 @@ import paddle.nn as nn ...@@ -25,7 +25,8 @@ import paddle.nn as nn
from ppcls.utils import config from ppcls.utils import config
from ppcls.engine.trainer import Trainer from ppcls.engine.trainer import Trainer
from ppcls.arch import build_model from ppcls.arch import build_model, RecModel
from ppcls.arch.backbone.base.theseus_layer import Identity
from ppcls.utils.save_load import load_dygraph_pretrain from ppcls.utils.save_load import load_dygraph_pretrain
...@@ -38,6 +39,9 @@ class ExportModel(nn.Layer): ...@@ -38,6 +39,9 @@ class ExportModel(nn.Layer):
super().__init__() super().__init__()
self.base_model = build_model(config) self.base_model = build_model(config)
self.infer_output_key = config.get("infer_output_key") self.infer_output_key = config.get("infer_output_key")
if self.infer_output_key == "features" and isinstance(self.base_model,
RecModel):
self.base_model.neck = Identity()
if config.get("infer_add_softmax", True): if config.get("infer_add_softmax", True):
self.softmax = nn.Softmax(axis=-1) self.softmax = nn.Softmax(axis=-1)
else: else:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册