mirrors/PaddleClas
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/PaddleClas.git synced 2025-06-03 21:55:06 +08:00
Code Issues Projects Releases Wiki Activity

move identity_head to gears

Browse Source
This commit is contained in:
weishengyu 2021-06-07 14:22:53 +08:00
parent e8959c31d5
commit 07089dcfe7
2 changed files with 7 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
ppcls/arch/gears/identity_head.py Normal file
View File

@ -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}

11
tools/export_model.py
View File

@ -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)