[Enhancement] Replace with_xxx with output_op in export.py (#2547)

* [Enhancement] Replace with_xxx with output_op in export.py

Author: juncaipeng

docs/model_export.md

@@ -2,54 +2,53 @@ English | [简体中文](model_export_cn.md)
 
 # Model Export
 
-The trained model needs to be exported as a prediction model before deployment.
+After model training, we can export the inference model and deploy it using inference library, which achieves faster inference speed.
 
 This tutorial will show how to export a trained model。
 
 
-## Acquire the Pre-training Model
+## Acquire trained weight
+
+After model training, the weight with the highest accuracy is saved in ` path/to/save/best_ model/model.pdparams`.
+
+For the convenience of this demo, we run the following commands to download the [trained weight](https://paddleseg.bj.bcebos.com/dygraph/cityscapes/pp_liteseg_stdc1_cityscapes_1024x512_scale0.5_160k/model.pdparams) of PP-LiteSeg.
+
 
-In this example，BiseNetV2 model will be used. Run the following command or click [link](https://paddleseg.bj.bcebos.com/dygraph/cityscapes/bisenet_cityscapes_1024x1024_160k/model.pdparams) to download the pretrained model.
 ```shell
-mkdir bisenet && cd bisenet
-wget https://paddleseg.bj.bcebos.com/dygraph/cityscapes/bisenet_cityscapes_1024x1024_160k/model.pdparams
-cd ..
+wget https://paddleseg.bj.bcebos.com/dygraph/cityscapes/pp_liteseg_stdc1_cityscapes_1024x512_scale0.5_160k/model.pdparams
 ```
 
 ## Export the prediction Model
 
-Make sure you have installed PaddleSeg and are in the PaddleSeg directory.
-
-Run the following command, and the prediction model will be saved in `output` directory.
+Run the following command in the root of PaddleSeg, the inference model is saved in `output/inference_model`.
 
 ```shell
-export CUDA_VISIBLE_DEVICES=0 # Set a usable GPU.
-# If on windows, Run the following command：
-# set CUDA_VISIBLE_DEVICES=0
 python export.py \
-       --config configs/bisenet/bisenet_cityscapes_1024x1024_160k.yml \
-       --model_path bisenet/model.pdparams\
-       --save_dir output
-       --input_shape 1 3 512 1024
+       --config configs/pp_liteseg/pp_liteseg_stdc1_cityscapes_1024x512_scale0.5_160k.yml \
+       --model_path model.pdparams \
+       --save_dir output/inference_model
 ```
 
-### Description of Exported Script Parameters
+**Description of Exported Script Parameters**
 
-|parammeter|purpose|is needed|default|
+|Parammeter|Purpose|Is Needed|Default|
 |-|-|-|-|
-|config|Config file|yes|-|
-|save_dir|Save root path for model and VisualDL log files|no|output|
-|model_path|Path of pre-training model parameters|no|The value in config file|
-|with_softmax|Add softmax operator at the end of the network. Since PaddleSeg networking returns Logits by default, you can set it to True if you want the deployment model to get the probability value|no|False|
-|without_argmax|Whether or not to add argmax operator at the end of the network. Since PaddleSeg networking returns Logits by default, we add argmax operator at the end of the network by default in order to directly obtain the prediction results for the deployment model|no|False|
-|input_shape| Set the input shape of exported model, such as `--input_shape 1 3 1024 1024`。if input_shape is not provided，the Default input shape of exported model is [-1, 3, -1, -1] | no (If the image shape in prediction is consistent, you should set the input_shape) | None |
+|config|The path of config file|yes|-|
+|model_path|The path of trained weight|no|-|
+|save_dir| The save dir for the inference model|no|output/inference_model|
+|input_shape| Set the input shape (`N*C*H*W`) of the inference model, such as `--input_shape 1 3 1024 1024`。if input_shape is not provided，the input shape of the inference model is [-1, 3, -1, -1]. If the image shape in prediction is fixed, you should set the input_shape. | no  | None |
+|output_op | Set the op that is appended to the inference model, should in [`argmax`, `softmax`, `none`]. PaddleSeg models outputs logits (`N*C*H*W`) by default. Adding `argmax` operation, we get the label for every pixel, the dimension of output is `N*H*W`. Adding `softmax` operation, we get the probability of different classes for every pixel. | no | argmax |
+|with_softmax| Deprecated params, please use --output_op. Add softmax operator at the end of the network. Since PaddleSeg networking returns Logits by default, you can set it to True if you want the deployment model to get the probability value|no|False|
+|without_argmax|Deprecated params, please use --output_op. Whether or not to add argmax operator at the end of the network. Since PaddleSeg networking returns Logits by default, we add argmax operator at the end of the network by default in order to directly obtain the prediction results for the deployment model|no|False|
+
 
-**If you encounter shape-relevant issue, please try to set the input_shape.**
+Note that:
+* If you encounter shape-relevant issue, please try to set the input_shape.
 
 ## Prediction Model Files
 
 ```shell
-output
+output/inference_model
   ├── deploy.yaml            # Config file of deployment
   ├── model.pdiparams        # Paramters of static model
   ├── model.pdiparams.info   # Additional information witch is not concerned generally

docs/model_export_cn.md

@@ -2,56 +2,53 @@
 
 # 导出预测模型
 
-PaddleSeg训练好模型后，需要将模型导出为预测模型，才可以进行模型部署。
+使用PaddleSeg训练好模型后，我们将模型导出为预测模型、使用预测库进行部署，可以实现更快的推理速度。
 
-本教程提供一个示例介绍模型导出的过程。
+本教程基于一个示例介绍模型导出的过程。
 
-## 1. 获取预训练莫模型
+## 1. 获取模型权重
 
-本示例中，我们使用BiseNetV2模型，大家执行如下命令或者点击[链接](https://paddleseg.bj.bcebos.com/dygraph/cityscapes/bisenet_cityscapes_1024x1024_160k/model.pdparams)下载模型预训练权重。
+完成模型训练后，精度最高的模型权重保存在`path/to/save/best_model/model.pdparams`。
+
+本示例为了演示方便，大家执行如下命令或者点击[链接](https://paddleseg.bj.bcebos.com/dygraph/cityscapes/pp_liteseg_stdc1_cityscapes_1024x512_scale0.5_160k/model.pdparams)下载PP-LiteSeg模型训练好的权重。
 
 ```shell
-mkdir bisenet && cd bisenet
-wget https://paddleseg.bj.bcebos.com/dygraph/cityscapes/bisenet_cityscapes_1024x1024_160k/model.pdparams
-cd ..
+wget https://paddleseg.bj.bcebos.com/dygraph/cityscapes/pp_liteseg_stdc1_cityscapes_1024x512_scale0.5_160k/model.pdparams
 ```
 
 ## 2. 导出预测模型
 
-确保正确安装PaddleSeg后，在PaddleSeg目录下执行如下命令，则预测模型会保存在output文件夹。
+在PaddleSeg根目录下，执行如下命令，导出预测模型，保存在`output/inference_model`目录。
 
 ```shell
-# 设置1张可用的卡
-export CUDA_VISIBLE_DEVICES=0
-# windows下请执行以下命令
-# set CUDA_VISIBLE_DEVICES=0
 python export.py \
-       --config configs/bisenet/bisenet_cityscapes_1024x1024_160k.yml \
-       --model_path bisenet/model.pdparams \
-       --save_dir output
-       --input_shape 1 3 512 1024
+       --config configs/pp_liteseg/pp_liteseg_stdc1_cityscapes_1024x512_scale0.5_160k.yml \
+       --model_path model.pdparams \
+       --save_dir output/inference_model
 ```
 
-### 导出脚本参数解释
+**导出脚本参数解释：**
 
 |参数名|用途|是否必选项|默认值|
 |-|-|-|-|
-|config|配置文件|是|-|
-|model_path|预训练权重的路径|否|配置文件中指定的预训练权重路径|
-|save_dir|保存预测模型的路径|否|output|
-|input_shape| 设置导出模型的输入shape，比如传入`--input_shape 1 3 1024 1024`。如果不设置input_shape，默认导出模型的输入shape是[-1, 3, -1, -1] | 否（ 如果预测shape固定，建议指定input_shape参数 ） | None |
-|with_softmax|在网络末端添加softmax算子。由于PaddleSeg组网默认返回logits，如果想要部署模型获取概率值，可以置为True|否|False|
-|without_argmax|是否不在网络末端添加argmax算子。由于PaddleSeg组网默认返回logits，为部署模型可以直接获取预测结果，我们默认在网络末端添加argmax算子|否|False|
+|config        | 配置文件的路径       | 是 | - |
+|model_path    | 模型权重的路径      | 否 | - |
+|save_dir      | 预测模型保存的目录  | 否 | `./output/inference_model` |
+|input_shape   | 设置模型的输入shape (`N*C*H*W`)，比如传入`--input_shape 1 3 1024 1024`。如果不设置input_shape，默认导出模型的输入shape是`[-1, 3, -1, -1]`。 预测shape固定时，建议指定input_shape参数。 | 否 | None |
+|output_op     | 设置在模型末端添加的输出算子，支持[`argmax`, `softmax`, `none`]。PaddleSeg模型默认返回logits (`N*C*H*W`)；添加`argmax`算子，可以得到每个像素的分割类别，结果的维度是`N*H*W`、数据类型是`int32`；添加`softmax`算子，可以得到每个像素每类的概率，结果的维度是`N*C*H*W`、数据类型是`float32` | 否 | argmax |
+|with_softmax  | 即将废弃的输入参数，建议使用`--output_op`。在网络末端添加softmax算子。由于PaddleSeg组网默认返回logits，如果想要部署模型获取概率值，可以置为True | 否 | False |
+|without_argmax| 即将废弃的输入参数，建议使用`--output_op`。由于PaddleSeg组网默认返回logits，为部署模型可以直接获取预测结果，我们默认在网络末端添加argmax算子。如果设置`--without_argmax`，则不会在网络末端添加argmax算子。 | 否 | False |
 
-**如果导出出现和shape相关的问题，请尝试指定input_shape。**
+注意：
+* 如果部署模型时，出现和shape相关的问题，请尝试指定input_shape。
 
 ## 3. 预测模型文件
 
 如下是导出的预测模型文件。
 
 ```shell
-output
-  ├── deploy.yaml            # 部署相关的配置文件，主要说明数据预处理的方式
+output/inference_model
+  ├── deploy.yaml            # 部署相关的配置文件，主要说明数据预处理方式等信息
   ├── model.pdmodel          # 预测模型的拓扑结构文件
   ├── model.pdiparams        # 预测模型的权重文件
   └── model.pdiparams.info   # 参数额外信息，一般无需关注

export.py

@@ -24,35 +24,27 @@
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Model export.')
-    # params of training
     parser.add_argument(
-        "--config",
-        dest="cfg",
-        help="The config file.",
-        default=None,
-        type=str,
-        required=True)
+        "--config", help="The config file.", type=str, required=True)
+    parser.add_argument(
+        '--model_path', help='The path of model for export', type=str)
     parser.add_argument(
         '--save_dir',
-        dest='save_dir',
         help='The directory for saving the exported model',
         type=str,
-        default='./output')
+        default='./output/inference_model')
     parser.add_argument(
-        '--model_path',
-        dest='model_path',
-        help='The path of model for export',
-        type=str,
-        default=None)
+        '--output_op',
+        choices=['argmax', 'softmax', 'none'],
+        default="argmax",
+        help="Select which op to be appended to output result, default: argmax")
     parser.add_argument(
         '--without_argmax',
-        dest='without_argmax',
-        help='Do not add the argmax operation at the end of the network',
+        help='Do not add the argmax operation at the end of the network. [Deprecated]',
         action='store_true')
     parser.add_argument(
         '--with_softmax',
-        dest='with_softmax',
-        help='Add the softmax operation at the end of the network',
+        help='Add the softmax operation at the end of the network. [Deprecated]',
         action='store_true')
     parser.add_argument(
         "--input_shape",
@@ -65,41 +57,33 @@ def parse_args():
 
 
 class SavedSegmentationNet(paddle.nn.Layer):
-    def __init__(self, net, without_argmax=False, with_softmax=False):
+    def __init__(self, net, output_op):
         super().__init__()
         self.net = net
-        self.post_processer = PostPorcesser(without_argmax, with_softmax)
+        self.output_op = output_op
+        assert output_op in ['argmax', 'softmax'], \
+            "output_op should in ['argmax', 'softmax']"
 
     def forward(self, x):
         outs = self.net(x)
-        outs = self.post_processer(outs)
-        return outs
-
-
-class PostPorcesser(paddle.nn.Layer):
-    def __init__(self, without_argmax, with_softmax):
-        super().__init__()
-        self.without_argmax = without_argmax
-        self.with_softmax = with_softmax
 
-    def forward(self, outs):
         new_outs = []
         for out in outs:
-            if self.with_softmax:
+            if self.output_op == 'argmax':
+                out = paddle.argmax(out, axis=1, dtype='int32')
+            elif self.output_op == 'softmax':
                 out = paddle.nn.functional.softmax(out, axis=1)
-            if not self.without_argmax:
-                out = paddle.argmax(out, axis=1)
             new_outs.append(out)
         return new_outs
 
 
 def main(args):
     os.environ['PADDLESEG_EXPORT_STAGE'] = 'True'
-    cfg = Config(args.cfg)
+    cfg = Config(args.config)
     cfg.check_sync_info()
     net = cfg.model
 
-    if args.model_path:
+    if args.model_path is not None:
         para_state_dict = paddle.load(args.model_path)
         net.set_dict(para_state_dict)
         logger.info('Loaded trained params of model successfully.')
@@ -109,17 +93,24 @@ def main(args):
     else:
         shape = args.input_shape
 
-    if not args.without_argmax or args.with_softmax:
-        new_net = SavedSegmentationNet(net, args.without_argmax,
-                                       args.with_softmax)
-    else:
-        new_net = net
-
+    output_op = args.output_op
+    if args.without_argmax:
+        logger.warning(
+            '--without_argmax will be deprecated, please use --output_op')
+        output_op = 'none'
+    if args.with_softmax:
+        logger.warning(
+            '--with_softmax will be deprecated, please use --output_op')
+        output_op = 'softmax'
+
+    new_net = net if output_op == 'none' else SavedSegmentationNet(net,
+                                                                   output_op)
     new_net.eval()
     new_net = paddle.jit.to_static(
         new_net,
         input_spec=[paddle.static.InputSpec(
             shape=shape, dtype='float32')])
+
     save_path = os.path.join(args.save_dir, 'model')
     paddle.jit.save(new_net, save_path)
 
@@ -128,19 +119,20 @@ def main(args):
         transforms = cfg.export_config.get('transforms', [{
             'type': 'Normalize'
         }])
+        output_dtype = 'int32' if output_op == 'argmax' else 'float32'
         data = {
             'Deploy': {
-                'transforms': transforms,
                 'model': 'model.pdmodel',
                 'params': 'model.pdiparams',
-                'with_softmax': args.with_softmax,
-                'with_argmax': not args.without_argmax,
-                'input_shape': shape
+                'transforms': transforms,
+                'input_shape': shape,
+                'output_op': output_op,
+                'output_dtype': output_dtype
             }
         }
         yaml.dump(data, file)
 
-    logger.info(f'Model is saved in {args.save_dir}.')
+    logger.info(f'The inference model is saved in {args.save_dir}')
 
 
 if __name__ == '__main__':