[Enhancement] Support dy2stat for HRNet and DeeplabV3

benchmark/deeplabv3p.yml

@@ -1,5 +1,6 @@
 batch_size: 2
 iters: 80000
+to_static_training: False
 
 train_dataset:
   type: Cityscapes

benchmark/hrnet.yml

@@ -1,5 +1,6 @@
 batch_size: 2
 iters: 80000
+to_static_training: False
 
 train_dataset:
   type: Cityscapes

paddleseg/core/train.py

@@ -68,7 +68,8 @@ def train(model,
           keep_checkpoint_max=5,
           test_config=None,
           fp16=False,
-          profiler_options=None):
+          profiler_options=None,
+          to_static_training=False):
     """
     Launch training.
 
@@ -91,6 +92,7 @@ def train(model,
         test_config(dict, optional): Evaluation config.
         fp16 (bool, optional): Whether to use amp.
         profiler_options (str, optional): The option of train profiler.
+        to_static_training (bool, optional): Whether to use @to_static for training.
     """
     model.train()
     nranks = paddle.distributed.ParallelEnv().nranks
@@ -131,6 +133,10 @@ def train(model,
         from visualdl import LogWriter
         log_writer = LogWriter(save_dir)
 
+    if to_static_training:
+        model = paddle.jit.to_static(model)
+        logger.info("Successfully to apply @to_static")
+
     avg_loss = 0.0
     avg_loss_list = []
     iters_per_epoch = len(batch_sampler)

paddleseg/cvlibs/config.py

@@ -392,6 +392,11 @@ def test_config(self) -> Dict:
     def export_config(self) -> Dict:
         return self.dic.get('export', {})
 
+    @property
+    def to_static_training(self) -> bool:
+        '''Whether to use @to_static for training'''
+        return self.dic.get('to_static_training', False)
+
     def _is_meta_type(self, item: Any) -> bool:
         return isinstance(item, dict) and 'type' in item
 

train.py

@@ -190,7 +190,8 @@ def main(args):
         keep_checkpoint_max=args.keep_checkpoint_max,
         test_config=cfg.test_config,
         fp16=args.fp16,
-        profiler_options=args.profiler_options)
+        profiler_options=args.profiler_options,
+        to_static_training=cfg.to_static_training)
 
 
 if __name__ == '__main__':