update_slim_in_trainer (#5622)

ce0d2bdb · Chang Xu · GitHub · 0c4166af · ce0d2bdb
隐藏空白更改
内联并排

Showing with 7 addition and 0 deletion

ppdet/engine/trainer.py ppdet/engine/trainer.py +7 -0

未找到文件。
--- 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'])