From 07089dcfe79f68b0ab718b1ca5da98ce3fdb4905 Mon Sep 17 00:00:00 2001 From: weishengyu Date: Mon, 7 Jun 2021 14:22:53 +0800 Subject: [PATCH] move identity_head to gears --- ppcls/arch/gears/identity_head.py | 6 ++++++ tools/export_model.py | 11 +---------- 2 files changed, 7 insertions(+), 10 deletions(-) create mode 100644 ppcls/arch/gears/identity_head.py diff --git a/ppcls/arch/gears/identity_head.py b/ppcls/arch/gears/identity_head.py new file mode 100644 index 00000000..d4343b9c --- /dev/null +++ b/ppcls/arch/gears/identity_head.py @@ -0,0 +1,6 @@ +class IdentityHead(nn.Layer): + def __init__(self): + super(IdentityHead, self).__init__() + + def forward(self, x, label=None): + return {"features": x, "logits": None} \ No newline at end of file diff --git a/tools/export_model.py b/tools/export_model.py index f705b18b..c1fbc9d9 100644 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -24,10 +24,9 @@ import paddle import paddle.nn as nn from ppcls.utils import config -from ppcls.engine.trainer import Trainer 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.arch.gears.identity_head import IdentityHead class ExportModel(nn.Layer): @@ -62,14 +61,6 @@ class ExportModel(nn.Layer): return x -class IdentityHead(nn.Layer): - def __init__(self): - super(IdentityHead, self).__init__() - - def forward(self, x, label=None): - return {"features": x, "logits": None} - - if __name__ == "__main__": args = config.parse_args() config = config.get_config(args.config, overrides=args.override, show=True) -- GitLab