[Fix] Adapt Slimming Tools to New APIs (#2916)

* Adapt to new API

* Optimize distill_train

Author: Bobholamovic

deploy/python/infer_onnx_trt.py

@@ -24,7 +24,7 @@
 import onnx
 import onnxruntime
 
-from paddleseg.cvlibs import Config
+from paddleseg.cvlibs import Config, SegBuilder
 from paddleseg.utils import logger, utils
 """
 Export the Paddle model to ONNX, infer the ONNX model by TRT.
@@ -395,7 +395,8 @@ def export_load_infer(args, model=None):
     # 1. prepare
     if model is None:
         cfg = Config(args.config)
-        model = cfg.model
+        builder = SegBuilder(cfg)
+        model = builder.model
     if args.model_path is not None:
         utils.load_entire_model(model, args.model_path)
         logger.info('Loaded trained params of model successfully')

deploy/slim/distill/distill_train.py

@@ -13,14 +13,11 @@
 # limitations under the License.
 
 import argparse
-import random
 
-import paddle
-import numpy as np
 from paddleslim.dygraph.dist import Distill
 
-from paddleseg.cvlibs import manager, Config
-from paddleseg.utils import get_sys_env, logger, utils
+from paddleseg.cvlibs import Config, SegBuilder
+from paddleseg.utils import logger, utils
 from distill_utils import distill_train
 from distill_config import prepare_distill_adaptor, prepare_distill_config
 
@@ -117,47 +114,29 @@ def prepare_envs(args):
     """
     Set random seed and the device.
     """
-    if args.seed is not None:
-        paddle.seed(args.seed)
-        np.random.seed(args.seed)
-        random.seed(args.seed)
 
-    env_info = get_sys_env()
-    info = ['{}: {}'.format(k, v) for k, v in env_info.items()]
-    info = '\n'.join(['', format('Environment Information', '-^48s')] + info +
-                     ['-' * 48])
-    logger.info(info)
+    utils.set_seed(args.seed)
+    utils.show_env_info()
 
-    place = 'gpu' if env_info['Paddle compiled with cuda'] and env_info[
-        'GPUs used'] else 'cpu'
+    env_info = utils.get_sys_env()
+    place = 'gpu' if env_info['GPUs used'] else 'cpu'
+    utils.set_device(place)
 
-    paddle.set_device(place)
 
+def main(args):
+
+    prepare_envs(args)
 
-def prepare_config(args):
-    """
-    Create and check the config of student and teacher model.
-    Note: we only use the dataset generated by the student config.
-    """
     if args.teather_config is None or args.student_config is None:
         raise RuntimeError('No configuration file specified.')
-
     t_cfg = Config(args.teather_config)
     s_cfg = Config(
         args.student_config,
         learning_rate=args.learning_rate,
         iters=args.iters,
         batch_size=args.batch_size)
-
-    train_dataset = s_cfg.train_dataset
-    val_dataset = s_cfg.val_dataset if args.do_eval else None
-    if train_dataset is None:
-        raise RuntimeError(
-            'The training dataset is not specified in the configuration file.')
-    elif len(train_dataset) == 0:
-        raise ValueError(
-            'The length of train_dataset is 0. Please check if your dataset is valid'
-        )
+    t_builder = SegBuilder(t_cfg)
+    s_builder = SegBuilder(s_cfg)
 
     msg = '\n---------------Teacher Config Information---------------\n'
     msg += str(t_cfg)
@@ -169,21 +148,12 @@ def prepare_config(args):
     msg += '------------------------------------------------'
     logger.info(msg)
 
-    return t_cfg, s_cfg, train_dataset, val_dataset
-
-
-def main(args):
-
-    prepare_envs(args)
-
-    t_cfg, s_cfg, train_dataset, val_dataset = prepare_config(args)
-
     distill_config = prepare_distill_config()
 
     s_adaptor, t_adaptor = prepare_distill_adaptor()
 
-    t_model = t_cfg.model
-    s_model = s_cfg.model
+    t_model = t_builder.model
+    s_model = s_builder.model
     t_model.eval()
     s_model.train()
 
@@ -192,19 +162,19 @@ def main(args):
 
     distill_train(
         distill_model=distill_model,
-        train_dataset=train_dataset,
-        val_dataset=val_dataset,
-        optimizer=s_cfg.optimizer,
+        train_dataset=s_builder.train_dataset,
+        val_dataset=s_builder.val_dataset,
+        optimizer=s_builder.optimizer,
         save_dir=args.save_dir,
-        iters=s_cfg.iters,
-        batch_size=s_cfg.batch_size,
+        iters=s_builder.iters,
+        batch_size=s_builder.batch_size,
         resume_model=args.resume_model,
         save_interval=args.save_interval,
         log_iters=args.log_iters,
         num_workers=args.num_workers,
         use_vdl=args.use_vdl,
-        losses=s_cfg.loss,
-        distill_losses=s_cfg.distill_loss,
+        losses=s_builder.loss,
+        distill_losses=s_builder.distill_loss,
         keep_checkpoint_max=args.keep_checkpoint_max,
         test_config=s_cfg.test_config, )
 

deploy/slim/prune/prune.py

@@ -18,14 +18,14 @@
 from functools import partial
 
 import yaml
-
 import paddle
 from paddleslim.dygraph import L1NormFilterPruner
 from paddleslim.analysis import dygraph_flops
-from paddleseg.cvlibs.config import Config
+
+from paddleseg.cvlibs.config import Config, SegBuilder
 from paddleseg.core.val import evaluate
 from paddleseg.core.train import train
-from paddleseg.utils import get_sys_env, logger
+from paddleseg.utils import logger, utils
 
 
 def parse_args():
@@ -87,9 +87,9 @@ def eval_fn(net, eval_dataset, num_workers):
     return miou
 
 
-def export_model(net, cfg, save_dir):
+def export_model(net, val_dataset, cfg, save_dir):
     net.forward = paddle.jit.to_static(net.forward)
-    input_shape = [1] + list(cfg.val_dataset[0]['img'].shape)
+    input_shape = [1] + list(val_dataset[0]['img'].shape)
     input_var = paddle.ones(input_shape)
     out = net(input_var)
 
@@ -98,7 +98,7 @@ def export_model(net, cfg, save_dir):
 
     yml_file = os.path.join(save_dir, 'deploy.yaml')
     with open(yml_file, 'w') as file:
-        transforms = cfg.dic['val_dataset']['transforms']
+        transforms = cfg.val_dataset_cfg['transforms']
         data = {
             'Deploy': {
                 'transforms': transforms,
@@ -110,11 +110,10 @@ def export_model(net, cfg, save_dir):
 
 
 def main(args):
-    env_info = get_sys_env()
+    env_info = utils.get_sys_env()
 
-    place = 'gpu' if env_info['Paddle compiled with cuda'] and env_info[
-        'GPUs used'] else 'cpu'
-    paddle.set_device(place)
+    place = 'gpu' if env_info['GPUs used'] else 'cpu'
+    utils.set_device(place)
 
     if not (0.0 < args.pruning_ratio < 1.0):
         raise RuntimeError(
@@ -123,24 +122,18 @@ def main(args):
     if not os.path.exists(args.save_dir):
         os.makedirs(args.save_dir)
 
+    os.environ['PADDLESEG_EXPORT_STAGE'] = 'True'
+
     cfg = Config(
         args.cfg,
         iters=args.retraining_iters,
         batch_size=args.batch_size,
         learning_rate=args.learning_rate)
+    builder = SegBuilder(cfg)
 
-    train_dataset = cfg.train_dataset
-    if not train_dataset:
-        raise RuntimeError(
-            'The training dataset is not specified in the configuration file.')
-
-    val_dataset = cfg.val_dataset
-    if not val_dataset:
-        raise RuntimeError(
-            'The validation dataset is not specified in the c;onfiguration file.'
-        )
-    os.environ['PADDLESEG_EXPORT_STAGE'] = 'True'
-    net = cfg.model
+    train_dataset = builder.train_dataset
+    val_dataset = builder.val_dataset
+    net = builder.model
 
     if args.model_path:
         para_state_dict = paddle.load(args.model_path)
@@ -180,17 +173,17 @@ def main(args):
     train(
         net,
         train_dataset,
-        optimizer=cfg.optimizer,
+        optimizer=builder.optimizer,
         save_dir=args.save_dir,
         num_workers=args.num_workers,
         iters=cfg.iters,
         batch_size=cfg.batch_size,
-        losses=cfg.loss)
+        losses=builder.loss)
 
     evaluate(net, val_dataset)
 
     if paddle.distributed.get_rank() == 0:
-        export_model(net, cfg, args.save_dir)
+        export_model(net, val_dataset, cfg, args.save_dir)
 
         ckpt = os.path.join(args.save_dir, f'iter_{args.retraining_iters}')
         if os.path.exists(ckpt):

deploy/slim/quant/ptq.py

@@ -17,10 +17,11 @@
 
 import numpy as np
 import shutil
-from paddleseg.cvlibs import manager, Config
 import paddle
 from paddleslim.quant import quant_post_static
 
+from paddleseg.cvlibs import Config, SegBuilder
+
 paddle.enable_static()
 
 
@@ -63,16 +64,11 @@ def __reader__():
 def main(args):
     fp32_model_dir = args.model_dir
     quant_output_dir = 'quant_model'
+
     cfg = Config(args.cfg)
-    val_dataset = cfg.val_dataset
-    if val_dataset is None:
-        raise RuntimeError(
-            'The verification dataset is not specified in the configuration file.'
-        )
-    elif len(val_dataset) == 0:
-        raise ValueError(
-            'The length of val_dataset is 0. Please check if your dataset is valid'
-        )
+    builder = SegBuilder(cfg)
+
+    val_dataset = builder.val_dataset
 
     use_gpu = True
     place = paddle.CUDAPlace(0) if use_gpu else paddle.CPUPlace()

deploy/slim/quant/qat_export.py

@@ -15,11 +15,11 @@
 import argparse
 import os
 
-import paddle
 import yaml
+import paddle
 from paddleslim import QAT
 
-from paddleseg.cvlibs import Config
+from paddleseg.cvlibs import Config, SegBuilder
 from paddleseg.utils import logger, utils
 from paddleseg.deploy.export import WrappedModel
 from qat_config import quant_config
@@ -59,7 +59,9 @@ def parse_args():
 def main(args):
     os.environ['PADDLESEG_EXPORT_STAGE'] = 'True'
     cfg = Config(args.config)
-    net = cfg.model
+    builder = SegBuilder(cfg)
+
+    net = builder.model
 
     skip_quant(net)
     quantizer = QAT(config=quant_config)
@@ -92,7 +94,7 @@ def main(args):
 
     yml_file = os.path.join(args.save_dir, 'deploy.yaml')
     with open(yml_file, 'w') as file:
-        transforms = cfg.val_dataset_config.get('transforms', [{
+        transforms = cfg.val_dataset_cfg.get('transforms', [{
             'type': 'Normalize'
         }])
         data = {