From ce0d2bdbc5b6f87923f95aacbf587d2ba48ddec1 Mon Sep 17 00:00:00 2001 From: Chang Xu Date: Tue, 12 Apr 2022 16:39:51 +0800 Subject: [PATCH] update_slim_in_trainer (#5622) --- ppdet/engine/trainer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ppdet/engine/trainer.py b/ppdet/engine/trainer.py index 35e1b9710..622b9c3b2 100644 --- a/ppdet/engine/trainer.py +++ b/ppdet/engine/trainer.py @@ -103,6 +103,13 @@ class Trainer(object): if 'slim' in cfg and cfg['slim_type'] == 'OFA': self.model.model.load_meanstd(cfg['TestReader'][ 'sample_transforms']) + elif 'slim' in cfg and cfg['slim_type'] == 'Distill': + self.model.student_model.load_meanstd(cfg['TestReader'][ + 'sample_transforms']) + elif 'slim' in cfg and cfg[ + 'slim_type'] == 'DistillPrune' and self.mode == 'train': + self.model.student_model.load_meanstd(cfg['TestReader'][ + 'sample_transforms']) else: self.model.load_meanstd(cfg['TestReader']['sample_transforms']) -- GitLab