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

fix rec forward bug

Browse Source
This commit is contained in:
dongshuilong 2021-09-02 07:42:22 +00:00
parent 81fcc9cb72
commit d49657ad08
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
ppcls/engine/engine.py
View File

@ -51,6 +51,11 @@ class Engine(object):
self.config = config
self.eval_mode = self.config["Global"].get("eval_mode",
"classification")
if "Head" in self.config["Arch"]:
self.is_rec = True
else:
self.is_rec = False
# init logger
self.output_dir = self.config['Global']['output_dir']
log_file = os.path.join(self.output_dir, self.config["Arch"]["name"],

2
ppcls/engine/train/train.py
View File

@ -79,7 +79,7 @@ def train_epoch(trainer, epoch_id, print_batch_step):
def forward(trainer, batch):
if trainer.eval_mode == "classification":
if not trainer.is_rec:
return trainer.model(batch[0])
else:
return trainer.model(batch[0], batch[1])