瀏覽代碼

Merge pull request #1 from PaddlePaddle/develop

add
SunAhong1993 5 年之前
父節點
當前提交
00fa48225f

+ 6 - 4
deploy/openvino/python/converter.py

@@ -20,6 +20,8 @@ import sys
 import yaml
 import paddlex as pdx
 
+assert pdx.__version__ >= '1.2.6', "paddlex >= 1.2.6 is required."
+
 
 def arg_parser():
     parser = argparse.ArgumentParser()
@@ -50,10 +52,11 @@ def arg_parser():
 
 def export_openvino_model(model, args):
 
+    onnx_save_file = os.path.join(args.save_dir, 'paddle2onnx_model.onnx')
     if model.__class__.__name__ == "YOLOv3":
-        pdx.converter.export_onnx_model(model, args.save_dir)
+        pdx.converter.export_onnx_model(model, onnx_save_file)
     else:
-        pdx.converter.export_onnx_model(model, args.save_dir, 11)
+        pdx.converter.export_onnx_model(model, onnx_save_file, 11)
 
     import mo.main as mo
     from mo.utils.cli_parser import get_onnx_cli_parser
@@ -61,8 +64,7 @@ def export_openvino_model(model, args):
     onnx_parser.add_argument("--model_dir", type=_text_type)
     onnx_parser.add_argument("--save_dir", type=_text_type)
     onnx_parser.add_argument("--fixed_input_shape")
-    onnx_input = os.path.join(args.save_dir, 'paddle2onnx_model.onnx')
-    onnx_parser.set_defaults(input_model=onnx_input)
+    onnx_parser.set_defaults(input_model=onnx_save_file)
     onnx_parser.set_defaults(output_dir=args.save_dir)
     shape_list = args.fixed_input_shape[1:-1].split(',')
     with open(osp.join(args.model_dir, "model.yml")) as f:

+ 4 - 4
deploy/openvino/python/deploy.py

@@ -173,9 +173,9 @@ class Predictor:
         score_name = next(it)
         score_map = np.squeeze(preds[score_name])
         score_map = np.transpose(score_map, (1, 2, 0))
-        
+
         im_info = preprocessed_inputs['im_info']
-        
+
         for info in im_info[0][::-1]:
             if info[0] == 'resize':
                 w, h = info[1][1], info[1][0]
@@ -192,8 +192,8 @@ class Predictor:
         """
         outputs = self.net.outputs
         for name in outputs:
-            if (len(outputs[name].shape) == 3):
-                output = preds[name][0]
+            if (len(outputs[name].shape) == 2):
+                output = preds[name]
         result = []
         for out in output:
             if (out[0] >= 0):

+ 2 - 2
deploy/openvino/src/paddlex.cpp

@@ -164,12 +164,12 @@ bool Model::predict(const cv::Mat& im, DetResult* result) {
   InferenceEngine::OutputsDataMap out_maps = network_.getOutputsInfo();
   std::string outputName;
   for (const auto & output_map : out_maps) {
-    if (output_map.second->getTensorDesc().getDims().size() == 3) {
+    if (output_map.second->getTensorDesc().getDims().size() == 2) {
       outputName = output_map.first;
     }
   }
   if (outputName.empty()) {
-    std::cerr << "get result node failed!" << std::endl:
+    std::cerr << "get result node failed!" << std::endl;
     return false;
   }
   InferenceEngine::Blob::Ptr output = infer_request.GetBlob(outputName);

+ 1 - 1
docs/apis/analysis.md

@@ -29,7 +29,7 @@ Seg分析器的分析接口,完成以下信息的分析统计:
 
 [代码示例](https://github.com/PaddlePaddle/PaddleX/blob/develop/examples/multi-channel_remote_sensing/tools/analysis.py)
 
-[统计信息示例](../../examples/multi-channel_remote_sensing/analysis.html#id2)
+[统计信息示例](../examples/multi-channel_remote_sensing/analysis.html#id2)
 
 ### cal_clipped_mean_std
 ```python

+ 1 - 1
docs/apis/models/semantic_segmentation.md

@@ -13,7 +13,7 @@ paddlex.seg.DeepLabv3p(num_classes=2, backbone='MobileNetV2_x1.0', output_stride
 > > - **num_classes** (int): 类别数。
 > > - **backbone** (str): DeepLabv3+的backbone网络,实现特征图的计算,取值范围为['Xception65', 'Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5', 'MobileNetV2_x1.0', 'MobileNetV2_x1.5', 'MobileNetV2_x2.0', 'MobileNetV3_large_x1_0_ssld'],默认值为'MobileNetV2_x1.0'。
 > > - **output_stride** (int): backbone 输出特征图相对于输入的下采样倍数,一般取值为8或16。默认16。
-> > - **aspp_with_sep_conv** (bool):  decoder模块是否采用separable convolutions。默认True。
+> > - **aspp_with_sep_conv** (bool):  aspp模块是否采用separable convolutions。默认True。
 > > - **decoder_use_sep_conv** (bool): decoder模块是否采用separable convolutions。默认True。
 > > - **encoder_with_aspp** (bool): 是否在encoder阶段采用aspp模块。默认True。
 > > - **enable_decoder** (bool): 是否使用decoder模块。默认True。

+ 1 - 1
docs/apis/transforms/cls_transforms.md

@@ -68,7 +68,7 @@ paddlex.cls.transforms.RandomCrop(crop_size=224, lower_scale=0.08, lower_ratio=3
 * **crop_size** (int): 随机裁剪后重新调整的目标边长。默认为224。
 * **lower_scale** (float): 裁剪面积相对原面积比例的最小限制。默认为0.08。
 * **lower_ratio** (float): 宽变换比例的最小限制。默认为3. / 4。
-* **upper_ratio** (float): 宽变换比例的最限制。默认为4. / 3。
+* **upper_ratio** (float): 宽变换比例的最限制。默认为4. / 3。
 
 ## RandomHorizontalFlip
 ```python

+ 1 - 1
docs/data/annotation/classification.md

@@ -22,7 +22,7 @@ MyDataset/ # 图像分类数据集根目录
 
 在模型进行训练时,我们需要划分训练集,验证集和测试集,因此需要对如上数据进行划分,直接使用paddlex命令即可将数据集随机划分成70%训练集,20%验证集和10%测试集
 ```
-paddlex --split_dataset ImageNet --dataset_dir MyDataset --val_value 0.2 --test_value 0.1
+paddlex --split_dataset --format ImageNet --dataset_dir MyDataset --val_value 0.2 --test_value 0.1
 ```
 
 划分好的数据集会额外生成`labels.txt`, `train_list.txt`, `val_list.txt`, `test_list.txt`四个文件,之后可直接进行训练。

+ 1 - 1
docs/data/annotation/semantic_segmentation.md

@@ -19,7 +19,7 @@
 
 ## 格式转换
 
-LabelMe标注后的数据还需要进行转换为SEG格式,才可以用于实例分割任务的训练,创建保存目录`D:\dataset_seg`,在python环境中安装paddlex后,使用如下命令即可
+LabelMe标注后的数据还需要进行转换为SEG格式,才可以用于语义分割任务的训练,创建保存目录`D:\dataset_seg`,在python环境中安装paddlex后,使用如下命令即可
 ```
 paddlex --data_conversion --source labelme --to SEG \
         --pics D:\MyDataset\JPEGImages \

+ 2 - 2
docs/deploy/openvino/export_openvino_model.md

@@ -17,7 +17,7 @@ paddle模型转openvino之前需要先把paddle模型导出为inference格式模
 paddlex --export_inference --model_dir=/path/to/paddle_model --save_dir=./inference_model --fixed_input_shape=[w,h]
 ```
 
-**注意**:需要转OpenVINO模型时,导出inference模型请务必指定`--fixed_input_shape`参数来固定模型的输入大小,且模型的输入大小需要与训练时一致
+**注意**:需要转OpenVINO模型时,导出inference模型请务必指定`--fixed_input_shape`参数来固定模型的输入大小,且模型的输入大小需要与训练时一致。 PaddleX客户端在发布模型时没有固定输入大小,因此对于可视化客户端,请找到任务所在目录,从里面的`output`文件夹找到`best_model`模型目录,将此目录使用如上命令进行固定shape导出即可。
 
 ## 导出OpenVINO模型
 
@@ -37,7 +37,7 @@ python converter.py --model_dir /path/to/inference_model --save_dir /path/to/ope
 | --model_dir  | Paddle模型路径,请确保__model__, \_\_params__model.yml在同一个目录|
 | --save_dir  | OpenVINO模型保存路径 |
 | --fixed_input_shape  | 模型输入的[W,H] |
-| --data type(option)  | (可选)FP32、FP16,默认为FP32,VPU下的IR需要为FP16 |  
+| --data_type(option)  | (可选)FP32、FP16,默认为FP32,VPU下的IR需要为FP16 |  
 
 **注意**:
 - 由于OpenVINO 从2021.1版本开始支持ONNX的resize-11 OP的原因,请下载OpenVINO 2021.1+的版本

+ 17 - 0
docs/deploy/server/encryption.md

@@ -165,3 +165,20 @@ D:\projects\images\xiaoduxiongn.jpeg
 .\paddlex_inference\detector.exe --model_dir=D:\projects\paddlex_encrypted_model --image_list=D:\projects\images_list.txt --use_gpu=1 --save_dir=output --key=kLAl1qOs5uRbFt0/RrIDTZW2+tOf5bzvUIaHGF8lJ1c=
 ```
 `--key`传入加密工具输出的密钥,例如`kLAl1qOs5uRbFt0/RrIDTZW2+tOf5bzvUIaHGF8lJ1c=`, 图片文件`可视化预测结果`会保存在`save_dir`参数设置的目录下。
+
+## 附录-PaddlePaddle其它模型加密使用
+
+PaddleX提供的加密方案不仅只适配于PaddleX训练的模型,对于PaddlePaddle其它模型套件或用户自行开发的模型同样适用。用户完全可以使用PaddleX的加密工具来加密如PaddleDetection、PaddleSeg、PaddleClas、PaddleOCR等套件训练的模型。
+
+### 加密方法
+1. 下载PaddleX加密工具
+- [Linux版本](https://bj.bcebos.com/paddlex/tools/1.2.0/paddlex-encryption.zip)
+- [Windows版本](https://bj.bcebos.com/paddlex/tools/win/1.2.0/paddlex-encryption.zip)
+
+2. 加密模型
+一般我们使用Paddle训练导出的部署模型都可以保存为`__model__`和`__params__`两个文件,PaddleX在模型保存时,还会额外再保存一个`model.yml`文件,用于存储模型的一些信息。PaddleX在加密模型时,会读取模型目录中这三个文件,进行加密,生成`__model__.encrypted`, `__params__.encrypted`和`model.yml.encrypted`三个文件。  
+因此在使用PaddleX加密工具加密别的套件或用户自行开发训练导出的模型时,需在模型目录中包含`__model__`,`__params__`和`model.yml`三个文件。非PaddleX训练的模型,用户可自行创建一个`model.yml`的同名文件,加密后无需理会此文件即可。
+> 模型加密命令参考本文档前面步骤
+
+3. 加载加密模型
+在使用Paddle预测库加载加密模型时,C++代码开发参考[Paddle部署代码模型加载函数](https://github.com/PaddlePaddle/PaddleX/blob/develop/deploy/cpp/src/paddlex.cpp#L46),同时需要引入新的头文件`paddle_model_decrypt.h`,相应依赖就在下载的加密工具中,CMakelist中参考[PaddleX的依赖](https://github.com/PaddlePaddle/PaddleX/blob/develop/deploy/cpp/CMakeLists.txt#L52)即可。

+ 1 - 1
docs/examples/change_detection.md

@@ -43,7 +43,7 @@ python prepare_data.py
 
 **注意:**
 
-* tiff格式的图片PaddleX统一使用gdal库读取,gdal安装可参考[文档](https://paddlex.readthedocs.io/zh_CN/develop/examples/multi-channel_remote_sensing/README.html#id2)。若数据是tiff格式的三通道RGB图像,如果不想安装gdal,需自行转成jpeg、bmp、png格式图片。
+* tiff格式的图片PaddleX统一使用gdal库读取,gdal安装可参考[gdal文档](https://paddlex.readthedocs.io/zh_CN/develop/examples/multi-channel_remote_sensing/README.html#id2)。若数据是tiff格式的三通道RGB图像,如果不想安装gdal,需自行转成jpeg、bmp、png格式图片。
 
 * label文件需为单通道的png格式图片,且标注从0开始计数,标注255表示该类别不参与计算。例如本案例中,0表示`unchanged`类,1表示`changed`类。
 

+ 27 - 0
docs/examples/human_segmentation.md

@@ -17,6 +17,33 @@
 > * Inference Model和Quant Inference Model为预测部署模型,包含`__model__`计算图结构、`__params__`模型参数和`model.yaml`基础的模型配置信息。
 > * 其中Inference Model适用于服务端的CPU和GPU预测部署,Qunat Inference Model为量化版本,适用于通过Paddle Lite进行移动端等端侧设备部署。
 
+#### 关于预测锯齿问题
+在训练完模型后,可能会遇到预测出来结果存在『锯齿』的问题,这个可能存在的原因是由于模型在预测过程中,经历了原图缩放再放大的过程,如下流程所示,
+
+```
+原图输入 -> 预处理transforms将图像缩放至目标大小 -> Paddle模型预测 -> 预测结果放大至原图大小 
+```
+对于这种原因导致的问题,可以手动修改模型中的`model.yml`文件,将预处理中的目标大小**调整到更高**优化此问题,如在本文档中提供的人像分割server端模型中`model.yml`文件内容,修改`target_size`至1024*1024(这样也会带来模型预测所需的资源更多,预测速度更慢)
+```
+Model: DeepLabv3p
+Transforms:
+- Resize:
+    interp: LINEAR
+    target_size:
+    - 512
+    - 512
+```
+修改为
+```
+Model: DeepLabv3p
+Transforms:
+- Resize:
+    interp: LINEAR
+    target_size:
+    - 1024
+    - 1024
+```
+
 
 预训练模型的存储大小和推理时长如下所示,其中移动端模型的运行环境为cpu:骁龙855,内存:6GB,图片大小:192*192
 

+ 1 - 0
docs/index.rst

@@ -82,3 +82,4 @@ PaddleX是基于飞桨核心框架、开发套件和工具组件的深度学习
    appendix/metrics.md
    appendix/interpret.md
    appendix/how_to_offline_run.md
+   change_log.md

+ 2 - 0
docs/slim/quant.md

@@ -19,6 +19,8 @@ dataset = pdx.datasets.ImageNet(
 
 # 开始量化
 pdx.slim.export_quant_model(model, dataset, 
+			  batch_size=4,
+			  batch_num=5,
 	                  save_dir='./quant_mobilenet', 
 	                  cache_dir='./tmp')
 ```

+ 1 - 5
docs/train/model_export.md

@@ -36,10 +36,6 @@ PaddleX作为开放开源的套件,其中的大部分模型均支持导出为O
 ### 如何导出
 
 - 1. 参考文档[部署模型导出](../deploy/eport_model.md),将训练保存的模型导出为部署模型  
-- 2. 安装paddle2onnx `pip install paddle2onnx`,转换命令如下,通过`--opset_version`指定版本(9/10/11)
-
-```
-paddle2onnx --model_dir paddlex_infere_model --save_file model.onnx --opset_version 10 --enable_onnx_checker True
-```
+- 2. 安装paddle2onnx `pip install paddle2onnx`,转换命令如下,通过`--opset_version`指定版本(9/10/11),转换使用方法参考[Paddle2ONNX说明](https://github.com/PaddlePaddle/paddle2onnx)
 
 - 附: Paddle2ONNX参阅 [https://github.com/PaddlePaddle/paddle2onnx](https://github.com/PaddlePaddle/paddle2onnx)

二進制
examples/human_segmentation/data/background.jpg


二進制
examples/human_segmentation/data/human_image.jpg


+ 11 - 5
paddlex/__init__.py

@@ -14,7 +14,7 @@
 
 from __future__ import absolute_import
 
-__version__ = '1.2.3'
+__version__ = '1.2.6'
 
 import os
 if 'FLAGS_eager_delete_tensor_gb' not in os.environ:
@@ -27,10 +27,16 @@ if "CUDA_VISIBLE_DEVICES" in os.environ:
 
 import paddle
 version = paddle.__version__.strip().split('.')
-if version[0] != '1' and version[1] != '8':
-    raise Exception(
-        "For running paddlex(v{}), Version of paddlepaddle should be greater than 1.8.3 and less than 2.0.0".
-        format(__version__))
+if version[0] == '1':
+    if version[1] != '8':
+        raise Exception(
+            'For running paddlex(v{}), Version of paddlepaddle should be greater than 1.8.3'.
+            format(__version__))
+elif version[0] == '2':
+    print(
+        "[WARNING] You are using paddlepaddle(v{}) which may not compatible with paddlex(v{}), paddlepaddle==1.8.4 is strongly recommended.".
+        format(paddle.__version__, __version__))
+    paddle.enable_static()
 
 from .utils.utils import get_environ_info
 from . import cv

+ 3 - 1
paddlex/command.py

@@ -161,6 +161,7 @@ def main():
         assert args.save_dir is not None, "--save_dir should be defined to create onnx model"
 
         model = pdx.load_model(args.model_dir)
+
         if model.status == "Normal" or model.status == "Prune":
             logging.error(
                 "Only support inference model, try to export model first as below,",
@@ -168,7 +169,8 @@ def main():
             logging.error(
                 "paddlex --export_inference --model_dir model_path --save_dir infer_model"
             )
-        pdx.converter.export_onnx_model(model, args.save_dir, args.onnx_opset)
+        save_file = os.path.join(args.save_dir, 'paddle2onnx_model.onnx') 
+        pdx.converter.export_onnx_model(model, save_file, args.onnx_opset)
 
     if args.data_conversion:
         assert args.source is not None, "--source should be defined while converting dataset"

+ 242 - 405
paddlex/converter.py

@@ -19,17 +19,244 @@ import sys
 import paddlex as pdx
 import paddlex.utils.logging as logging
 
-__all__ = ['export_onnx']
 
-
-def export_onnx(model_dir, save_dir, fixed_input_shape):
-    assert len(fixed_input_shape) == 2, "len of fixed input shape must == 2"
-    model = pdx.load_model(model_dir, fixed_input_shape)
-    model_name = os.path.basename(model_dir.strip('/')).split('/')[-1]
-    export_onnx_model(model, save_dir)
-
-
-def export_onnx_model(model, save_dir, opset_version=10):
+class MultiClassNMS4OpenVINO():
+    """
+    Convert the paddle multiclass_nms to onnx op.
+    This op is get the select boxes from origin boxes.
+    """
+    @classmethod
+    def opset_10(cls, graph, node, **kw):
+        from paddle2onnx.constant import dtypes
+        import numpy as np
+        result_name = node.output('Out', 0)
+        background = node.attr('background_label')
+        normalized = node.attr('normalized')
+        if normalized == False:
+            logging.warning(
+                        "The parameter normalized of multiclass_nms OP of Paddle is False, which has diff with ONNX." \
+                        " Please set normalized=True in multiclass_nms of Paddle, see doc Q1 in" \
+                        " https://github.com/PaddlePaddle/paddle2onnx/blob/develop/FAQ.md")
+    
+        #convert the paddle attribute to onnx tensor
+        node_score_threshold = graph.make_node(
+            'Constant',
+            inputs=[],
+            dtype=dtypes.ONNX.FLOAT,
+            value=[float(node.attr('score_threshold'))])
+    
+        node_iou_threshold = graph.make_node(
+            'Constant',
+            inputs=[],
+            dtype=dtypes.ONNX.FLOAT,
+            value=[float(node.attr('nms_threshold'))])
+    
+        node_keep_top_k = graph.make_node(
+            'Constant',
+            inputs=[],
+            dtype=dtypes.ONNX.INT64,
+            value=[np.int64(node.attr('keep_top_k'))])
+    
+        node_keep_top_k_2D = graph.make_node(
+            'Constant',
+            inputs=[],
+            dtype=dtypes.ONNX.INT64,
+            dims=[1, 1],
+            value=[node.attr('keep_top_k')])
+    
+        # the paddle data format is x1,y1,x2,y2
+        kwargs = {'center_point_box': 0}
+    
+        node_select_nms= graph.make_node(
+            'NonMaxSuppression',
+            inputs=[node.input('BBoxes', 0), node.input('Scores', 0), node_keep_top_k,\
+                node_iou_threshold, node_score_threshold])
+    
+        # step 1 nodes select the nms class
+        # create some const value to use
+        node_const_value = [result_name+"@const_0",
+            result_name+"@const_1",\
+            result_name+"@const_2",\
+            result_name+"@const_-1"]
+        value_const_value = [0, 1, 2, -1]
+        for name, value in zip(node_const_value, value_const_value):
+            graph.make_node(
+                'Constant',
+                layer_name=name,
+                inputs=[],
+                outputs=[name],
+                dtype=dtypes.ONNX.INT64,
+                value=[value])
+    
+        # In this code block, we will deocde the raw score data, reshape N * C * M to 1 * N*C*M
+        # and the same time, decode the select indices to 1 * D, gather the select_indices
+        node_gather_1 = graph.make_node(
+            'Gather',
+            inputs=[node_select_nms, result_name + "@const_1"],
+            axis=1)
+    
+        node_gather_1 = graph.make_node(
+            'Unsqueeze', inputs=[node_gather_1], axes=[0])
+    
+        node_gather_2 = graph.make_node(
+            'Gather',
+            inputs=[node_select_nms, result_name + "@const_2"],
+            axis=1)
+    
+        node_gather_2 = graph.make_node(
+            'Unsqueeze', inputs=[node_gather_2], axes=[0])
+    
+        # reshape scores N * C * M to (N*C*M) * 1
+        node_reshape_scores_rank1 = graph.make_node(
+            "Reshape",
+            inputs=[node.input('Scores', 0), result_name + "@const_-1"])
+    
+        # get the shape of scores
+        node_shape_scores = graph.make_node(
+            'Shape', inputs=node.input('Scores'))
+    
+        # gather the index: 2 shape of scores
+        node_gather_scores_dim1 = graph.make_node(
+            'Gather',
+            inputs=[node_shape_scores, result_name + "@const_2"],
+            axis=0)
+    
+        # mul class * M
+        node_mul_classnum_boxnum = graph.make_node(
+            'Mul', inputs=[node_gather_1, node_gather_scores_dim1])
+    
+        # add class * M * index
+        node_add_class_M_index = graph.make_node(
+            'Add', inputs=[node_mul_classnum_boxnum, node_gather_2])
+    
+        # Squeeze the indices to 1 dim
+        node_squeeze_select_index = graph.make_node(
+            'Squeeze', inputs=[node_add_class_M_index], axes=[0, 2])
+    
+        # gather the data from flatten scores
+        node_gather_select_scores = graph.make_node(
+            'Gather',
+            inputs=[node_reshape_scores_rank1, node_squeeze_select_index],
+            axis=0)
+    
+        # get nums to input TopK
+        node_shape_select_num = graph.make_node(
+            'Shape', inputs=[node_gather_select_scores])
+    
+        node_gather_select_num = graph.make_node(
+            'Gather',
+            inputs=[node_shape_select_num, result_name + "@const_0"],
+            axis=0)
+    
+        node_unsqueeze_select_num = graph.make_node(
+            'Unsqueeze', inputs=[node_gather_select_num], axes=[0])
+    
+        node_concat_topK_select_num = graph.make_node(
+            'Concat',
+            inputs=[node_unsqueeze_select_num, node_keep_top_k_2D],
+            axis=0)
+    
+        node_cast_concat_topK_select_num = graph.make_node(
+            'Cast', inputs=[node_concat_topK_select_num], to=6)
+        # get min(topK, num_select)
+        node_compare_topk_num_select = graph.make_node(
+            'ReduceMin', inputs=[node_cast_concat_topK_select_num], keepdims=0)
+    
+        # unsqueeze the indices to 1D tensor
+        node_unsqueeze_topk_select_indices = graph.make_node(
+            'Unsqueeze', inputs=[node_compare_topk_num_select], axes=[0])
+    
+        # cast the indices to INT64
+        node_cast_topk_indices = graph.make_node(
+            'Cast', inputs=[node_unsqueeze_topk_select_indices], to=7)
+    
+        # select topk scores  indices
+        outputs_topk_select_topk_indices = [result_name + "@topk_select_topk_values",\
+            result_name + "@topk_select_topk_indices"]
+        node_topk_select_topk_indices = graph.make_node(
+            'TopK',
+            inputs=[node_gather_select_scores, node_cast_topk_indices],
+            outputs=outputs_topk_select_topk_indices)
+    
+        # gather topk label, scores, boxes
+        node_gather_topk_scores = graph.make_node(
+            'Gather',
+            inputs=[
+                node_gather_select_scores, outputs_topk_select_topk_indices[1]
+            ],
+            axis=0)
+    
+        node_gather_topk_class = graph.make_node(
+            'Gather',
+            inputs=[
+                node_gather_1, outputs_topk_select_topk_indices[1]
+            ],
+            axis=1)
+    
+        # gather the boxes need to gather the boxes id, then get boxes
+        node_gather_topk_boxes_id = graph.make_node(
+            'Gather',
+            inputs=[
+                node_gather_2, outputs_topk_select_topk_indices[1]
+            ],
+            axis=1)
+    
+        # squeeze the gather_topk_boxes_id to 1 dim
+        node_squeeze_topk_boxes_id = graph.make_node(
+            'Squeeze', inputs=[node_gather_topk_boxes_id], axes=[0, 2])
+    
+        node_gather_select_boxes = graph.make_node(
+            'Gather',
+            inputs=[node.input('BBoxes', 0), node_squeeze_topk_boxes_id],
+            axis=1)
+    
+        # concat the final result
+        # before concat need to cast the class to float
+        node_cast_topk_class = graph.make_node(
+            'Cast', inputs=[node_gather_topk_class], to=1)
+    
+        node_unsqueeze_topk_scores = graph.make_node(
+            'Unsqueeze', inputs=[node_gather_topk_scores], axes=[0, 2])
+    
+        inputs_concat_final_results = [node_cast_topk_class, node_unsqueeze_topk_scores, \
+            node_gather_select_boxes]
+        node_sort_by_socre_results = graph.make_node(
+            'Concat', inputs=inputs_concat_final_results, axis=2)
+    
+        # select topk classes indices
+        node_squeeze_cast_topk_class = graph.make_node(
+            'Squeeze', inputs=[node_cast_topk_class], axes=[0, 2])
+        node_neg_squeeze_cast_topk_class = graph.make_node(
+            'Neg', inputs=[node_squeeze_cast_topk_class])
+    
+        outputs_topk_select_classes_indices = [result_name + "@topk_select_topk_classes_scores",\
+            result_name + "@topk_select_topk_classes_indices"]
+        node_topk_select_topk_indices = graph.make_node(
+            'TopK',
+            inputs=[node_neg_squeeze_cast_topk_class, node_cast_topk_indices],
+            outputs=outputs_topk_select_classes_indices)
+        node_concat_final_results = graph.make_node(
+            'Gather',
+            inputs=[
+                node_sort_by_socre_results,
+                outputs_topk_select_classes_indices[1]
+            ],
+            axis=1)
+        node_concat_final_results = graph.make_node(
+            'Squeeze',
+            inputs=[node_concat_final_results],
+            outputs=[node.output('Out', 0)],
+            axes=[0])
+    
+        if node.type == 'multiclass_nms2':
+            graph.make_node(
+                'Squeeze',
+                inputs=[node_gather_2],
+                outputs=node.output('Index'),
+                axes=[0])
+    
+
+def export_onnx_model(model, save_file, opset_version=10):
     if model.__class__.__name__ == "FastSCNN" or (
             model.model_type == "detector" and
             model.__class__.__name__ != "YOLOv3"):
@@ -46,401 +273,11 @@ def export_onnx_model(model, save_dir, opset_version=10):
         logging.warning(
             "Export for openVINO by default, the output of multiclass_nms exported to onnx will contains background. If you need onnx completely consistent with paddle, please use paddle2onnx to export"
         )
-        p2o.op_mapper.opset9.paddle_custom_layer.multiclass_nms.multiclass_nms = multiclass_nms_for_openvino
-    mapper = p2o.PaddleOpMapper()
-    mapper.convert(
+    
+    p2o.register_op_mapper('multiclass_nms', MultiClassNMS4OpenVINO)
+    
+    p2o.program2onnx(
         model.test_prog,
-        save_dir,
         scope=model.scope,
+        save_file=save_file,
         opset_version=opset_version)
-
-
-def multiclass_nms_for_openvino(op, block):
-    """
-    Convert the paddle multiclass_nms to onnx op.
-    This op is get the select boxes from origin boxes.
-    This op is for OpenVINO, which donn't support dynamic shape).
-    """
-    import math
-    import sys
-    import numpy as np
-    import paddle.fluid.core as core
-    import paddle.fluid as fluid
-    import onnx
-    import warnings
-    from onnx import helper, onnx_pb
-    inputs = dict()
-    outputs = dict()
-    attrs = dict()
-    for name in op.input_names:
-        inputs[name] = op.input(name)
-    for name in op.output_names:
-        outputs[name] = op.output(name)
-    for name in op.attr_names:
-        attrs[name] = op.attr(name)
-
-    result_name = outputs['Out'][0]
-    background = attrs['background_label']
-    normalized = attrs['normalized']
-    if normalized == False:
-        warnings.warn(
-            'The parameter normalized of multiclass_nms OP of Paddle is False, which has diff with ONNX. \
-                         Please set normalized=True in multiclass_nms of Paddle'
-        )
-
-    #convert the paddle attribute to onnx tensor
-    name_score_threshold = [outputs['Out'][0] + "@score_threshold"]
-    name_iou_threshold = [outputs['Out'][0] + "@iou_threshold"]
-    name_keep_top_k = [outputs['Out'][0] + '@keep_top_k']
-    name_keep_top_k_2D = [outputs['Out'][0] + '@keep_top_k_1D']
-
-    node_score_threshold = onnx.helper.make_node(
-        'Constant',
-        inputs=[],
-        outputs=name_score_threshold,
-        value=onnx.helper.make_tensor(
-            name=name_score_threshold[0] + "@const",
-            data_type=onnx.TensorProto.FLOAT,
-            dims=(),
-            vals=[float(attrs['score_threshold'])]))
-
-    node_iou_threshold = onnx.helper.make_node(
-        'Constant',
-        inputs=[],
-        outputs=name_iou_threshold,
-        value=onnx.helper.make_tensor(
-            name=name_iou_threshold[0] + "@const",
-            data_type=onnx.TensorProto.FLOAT,
-            dims=(),
-            vals=[float(attrs['nms_threshold'])]))
-
-    node_keep_top_k = onnx.helper.make_node(
-        'Constant',
-        inputs=[],
-        outputs=name_keep_top_k,
-        value=onnx.helper.make_tensor(
-            name=name_keep_top_k[0] + "@const",
-            data_type=onnx.TensorProto.INT64,
-            dims=(),
-            vals=[np.int64(attrs['keep_top_k'])]))
-
-    node_keep_top_k_2D = onnx.helper.make_node(
-        'Constant',
-        inputs=[],
-        outputs=name_keep_top_k_2D,
-        value=onnx.helper.make_tensor(
-            name=name_keep_top_k_2D[0] + "@const",
-            data_type=onnx.TensorProto.INT64,
-            dims=[1, 1],
-            vals=[np.int64(attrs['keep_top_k'])]))
-
-    # the paddle data format is x1,y1,x2,y2
-    kwargs = {'center_point_box': 0}
-
-    name_select_nms = [outputs['Out'][0] + "@select_index"]
-    node_select_nms= onnx.helper.make_node(
-        'NonMaxSuppression',
-        inputs=inputs['BBoxes'] + inputs['Scores'] + name_keep_top_k +\
-            name_iou_threshold + name_score_threshold,
-        outputs=name_select_nms)
-    # step 1 nodes select the nms class
-    node_list = [
-        node_score_threshold, node_iou_threshold, node_keep_top_k,
-        node_keep_top_k_2D, node_select_nms
-    ]
-
-    # create some const value to use
-    name_const_value = [result_name+"@const_0",
-        result_name+"@const_1",\
-        result_name+"@const_2",\
-        result_name+"@const_-1"]
-    value_const_value = [0, 1, 2, -1]
-    for name, value in zip(name_const_value, value_const_value):
-        node = onnx.helper.make_node(
-            'Constant',
-            inputs=[],
-            outputs=[name],
-            value=onnx.helper.make_tensor(
-                name=name + "@const",
-                data_type=onnx.TensorProto.INT64,
-                dims=[1],
-                vals=[value]))
-        node_list.append(node)
-
-    # In this code block, we will deocde the raw score data, reshape N * C * M to 1 * N*C*M
-    # and the same time, decode the select indices to 1 * D, gather the select_indices
-    outputs_gather_1_ = [result_name + "@gather_1_"]
-    node_gather_1_ = onnx.helper.make_node(
-        'Gather',
-        inputs=name_select_nms + [result_name + "@const_1"],
-        outputs=outputs_gather_1_,
-        axis=1)
-    node_list.append(node_gather_1_)
-    outputs_gather_1 = [result_name + "@gather_1"]
-    node_gather_1 = onnx.helper.make_node(
-        'Unsqueeze',
-        inputs=outputs_gather_1_,
-        outputs=outputs_gather_1,
-        axes=[0])
-    node_list.append(node_gather_1)
-
-    outputs_gather_2_ = [result_name + "@gather_2_"]
-    node_gather_2_ = onnx.helper.make_node(
-        'Gather',
-        inputs=name_select_nms + [result_name + "@const_2"],
-        outputs=outputs_gather_2_,
-        axis=1)
-    node_list.append(node_gather_2_)
-
-    outputs_gather_2 = [result_name + "@gather_2"]
-    node_gather_2 = onnx.helper.make_node(
-        'Unsqueeze',
-        inputs=outputs_gather_2_,
-        outputs=outputs_gather_2,
-        axes=[0])
-    node_list.append(node_gather_2)
-
-    # reshape scores N * C * M to (N*C*M) * 1
-    outputs_reshape_scores_rank1 = [result_name + "@reshape_scores_rank1"]
-    node_reshape_scores_rank1 = onnx.helper.make_node(
-        "Reshape",
-        inputs=inputs['Scores'] + [result_name + "@const_-1"],
-        outputs=outputs_reshape_scores_rank1)
-    node_list.append(node_reshape_scores_rank1)
-
-    # get the shape of scores
-    outputs_shape_scores = [result_name + "@shape_scores"]
-    node_shape_scores = onnx.helper.make_node(
-        'Shape', inputs=inputs['Scores'], outputs=outputs_shape_scores)
-    node_list.append(node_shape_scores)
-
-    # gather the index: 2 shape of scores
-    outputs_gather_scores_dim1 = [result_name + "@gather_scores_dim1"]
-    node_gather_scores_dim1 = onnx.helper.make_node(
-        'Gather',
-        inputs=outputs_shape_scores + [result_name + "@const_2"],
-        outputs=outputs_gather_scores_dim1,
-        axis=0)
-    node_list.append(node_gather_scores_dim1)
-
-    # mul class * M
-    outputs_mul_classnum_boxnum = [result_name + "@mul_classnum_boxnum"]
-    node_mul_classnum_boxnum = onnx.helper.make_node(
-        'Mul',
-        inputs=outputs_gather_1 + outputs_gather_scores_dim1,
-        outputs=outputs_mul_classnum_boxnum)
-    node_list.append(node_mul_classnum_boxnum)
-
-    # add class * M * index
-    outputs_add_class_M_index = [result_name + "@add_class_M_index"]
-    node_add_class_M_index = onnx.helper.make_node(
-        'Add',
-        inputs=outputs_mul_classnum_boxnum + outputs_gather_2,
-        outputs=outputs_add_class_M_index)
-    node_list.append(node_add_class_M_index)
-
-    # Squeeze the indices to 1 dim
-    outputs_squeeze_select_index = [result_name + "@squeeze_select_index"]
-    node_squeeze_select_index = onnx.helper.make_node(
-        'Squeeze',
-        inputs=outputs_add_class_M_index,
-        outputs=outputs_squeeze_select_index,
-        axes=[0, 2])
-    node_list.append(node_squeeze_select_index)
-
-    # gather the data from flatten scores
-    outputs_gather_select_scores = [result_name + "@gather_select_scores"]
-    node_gather_select_scores = onnx.helper.make_node('Gather',
-        inputs=outputs_reshape_scores_rank1 + \
-            outputs_squeeze_select_index,
-        outputs=outputs_gather_select_scores,
-        axis=0)
-    node_list.append(node_gather_select_scores)
-
-    # get nums to input TopK
-    outputs_shape_select_num = [result_name + "@shape_select_num"]
-    node_shape_select_num = onnx.helper.make_node(
-        'Shape',
-        inputs=outputs_gather_select_scores,
-        outputs=outputs_shape_select_num)
-    node_list.append(node_shape_select_num)
-
-    outputs_gather_select_num = [result_name + "@gather_select_num"]
-    node_gather_select_num = onnx.helper.make_node(
-        'Gather',
-        inputs=outputs_shape_select_num + [result_name + "@const_0"],
-        outputs=outputs_gather_select_num,
-        axis=0)
-    node_list.append(node_gather_select_num)
-
-    outputs_unsqueeze_select_num = [result_name + "@unsqueeze_select_num"]
-    node_unsqueeze_select_num = onnx.helper.make_node(
-        'Unsqueeze',
-        inputs=outputs_gather_select_num,
-        outputs=outputs_unsqueeze_select_num,
-        axes=[0])
-    node_list.append(node_unsqueeze_select_num)
-
-    outputs_concat_topK_select_num = [result_name + "@conat_topK_select_num"]
-    node_conat_topK_select_num = onnx.helper.make_node(
-        'Concat',
-        inputs=outputs_unsqueeze_select_num + name_keep_top_k_2D,
-        outputs=outputs_concat_topK_select_num,
-        axis=0)
-    node_list.append(node_conat_topK_select_num)
-
-    outputs_cast_concat_topK_select_num = [
-        result_name + "@concat_topK_select_num"
-    ]
-    node_outputs_cast_concat_topK_select_num = onnx.helper.make_node(
-        'Cast',
-        inputs=outputs_concat_topK_select_num,
-        outputs=outputs_cast_concat_topK_select_num,
-        to=6)
-    node_list.append(node_outputs_cast_concat_topK_select_num)
-    # get min(topK, num_select)
-    outputs_compare_topk_num_select = [
-        result_name + "@compare_topk_num_select"
-    ]
-    node_compare_topk_num_select = onnx.helper.make_node(
-        'ReduceMin',
-        inputs=outputs_cast_concat_topK_select_num,
-        outputs=outputs_compare_topk_num_select,
-        keepdims=0)
-    node_list.append(node_compare_topk_num_select)
-
-    # unsqueeze the indices to 1D tensor
-    outputs_unsqueeze_topk_select_indices = [
-        result_name + "@unsqueeze_topk_select_indices"
-    ]
-    node_unsqueeze_topk_select_indices = onnx.helper.make_node(
-        'Unsqueeze',
-        inputs=outputs_compare_topk_num_select,
-        outputs=outputs_unsqueeze_topk_select_indices,
-        axes=[0])
-    node_list.append(node_unsqueeze_topk_select_indices)
-
-    # cast the indices to INT64
-    outputs_cast_topk_indices = [result_name + "@cast_topk_indices"]
-    node_cast_topk_indices = onnx.helper.make_node(
-        'Cast',
-        inputs=outputs_unsqueeze_topk_select_indices,
-        outputs=outputs_cast_topk_indices,
-        to=7)
-    node_list.append(node_cast_topk_indices)
-
-    # select topk scores  indices
-    outputs_topk_select_topk_indices = [result_name + "@topk_select_topk_values",\
-        result_name + "@topk_select_topk_indices"]
-    node_topk_select_topk_indices = onnx.helper.make_node(
-        'TopK',
-        inputs=outputs_gather_select_scores + outputs_cast_topk_indices,
-        outputs=outputs_topk_select_topk_indices)
-    node_list.append(node_topk_select_topk_indices)
-
-    # gather topk label, scores, boxes
-    outputs_gather_topk_scores = [result_name + "@gather_topk_scores"]
-    node_gather_topk_scores = onnx.helper.make_node(
-        'Gather',
-        inputs=outputs_gather_select_scores +
-        [outputs_topk_select_topk_indices[1]],
-        outputs=outputs_gather_topk_scores,
-        axis=0)
-    node_list.append(node_gather_topk_scores)
-
-    outputs_gather_topk_class = [result_name + "@gather_topk_class"]
-    node_gather_topk_class = onnx.helper.make_node(
-        'Gather',
-        inputs=outputs_gather_1 + [outputs_topk_select_topk_indices[1]],
-        outputs=outputs_gather_topk_class,
-        axis=1)
-    node_list.append(node_gather_topk_class)
-
-    # gather the boxes need to gather the boxes id, then get boxes
-    outputs_gather_topk_boxes_id = [result_name + "@gather_topk_boxes_id"]
-    node_gather_topk_boxes_id = onnx.helper.make_node(
-        'Gather',
-        inputs=outputs_gather_2 + [outputs_topk_select_topk_indices[1]],
-        outputs=outputs_gather_topk_boxes_id,
-        axis=1)
-    node_list.append(node_gather_topk_boxes_id)
-
-    # squeeze the gather_topk_boxes_id to 1 dim
-    outputs_squeeze_topk_boxes_id = [result_name + "@squeeze_topk_boxes_id"]
-    node_squeeze_topk_boxes_id = onnx.helper.make_node(
-        'Squeeze',
-        inputs=outputs_gather_topk_boxes_id,
-        outputs=outputs_squeeze_topk_boxes_id,
-        axes=[0, 2])
-    node_list.append(node_squeeze_topk_boxes_id)
-
-    outputs_gather_select_boxes = [result_name + "@gather_select_boxes"]
-    node_gather_select_boxes = onnx.helper.make_node(
-        'Gather',
-        inputs=inputs['BBoxes'] + outputs_squeeze_topk_boxes_id,
-        outputs=outputs_gather_select_boxes,
-        axis=1)
-    node_list.append(node_gather_select_boxes)
-
-    # concat the final result
-    # before concat need to cast the class to float
-    outputs_cast_topk_class = [result_name + "@cast_topk_class"]
-    node_cast_topk_class = onnx.helper.make_node(
-        'Cast',
-        inputs=outputs_gather_topk_class,
-        outputs=outputs_cast_topk_class,
-        to=1)
-    node_list.append(node_cast_topk_class)
-
-    outputs_unsqueeze_topk_scores = [result_name + "@unsqueeze_topk_scores"]
-    node_unsqueeze_topk_scores = onnx.helper.make_node(
-        'Unsqueeze',
-        inputs=outputs_gather_topk_scores,
-        outputs=outputs_unsqueeze_topk_scores,
-        axes=[0, 2])
-    node_list.append(node_unsqueeze_topk_scores)
-
-    inputs_concat_final_results = outputs_cast_topk_class + outputs_unsqueeze_topk_scores +\
-        outputs_gather_select_boxes
-    outputs_sort_by_socre_results = [result_name + "@concat_topk_scores"]
-    node_sort_by_socre_results = onnx.helper.make_node(
-        'Concat',
-        inputs=inputs_concat_final_results,
-        outputs=outputs_sort_by_socre_results,
-        axis=2)
-    node_list.append(node_sort_by_socre_results)
-
-    # select topk classes indices
-    outputs_squeeze_cast_topk_class = [
-        result_name + "@squeeze_cast_topk_class"
-    ]
-    node_squeeze_cast_topk_class = onnx.helper.make_node(
-        'Squeeze',
-        inputs=outputs_cast_topk_class,
-        outputs=outputs_squeeze_cast_topk_class,
-        axes=[0, 2])
-    node_list.append(node_squeeze_cast_topk_class)
-    outputs_neg_squeeze_cast_topk_class = [
-        result_name + "@neg_squeeze_cast_topk_class"
-    ]
-    node_neg_squeeze_cast_topk_class = onnx.helper.make_node(
-        'Neg',
-        inputs=outputs_squeeze_cast_topk_class,
-        outputs=outputs_neg_squeeze_cast_topk_class)
-    node_list.append(node_neg_squeeze_cast_topk_class)
-    outputs_topk_select_classes_indices = [result_name + "@topk_select_topk_classes_scores",\
-        result_name + "@topk_select_topk_classes_indices"]
-    node_topk_select_topk_indices = onnx.helper.make_node(
-        'TopK',
-        inputs=outputs_neg_squeeze_cast_topk_class + outputs_cast_topk_indices,
-        outputs=outputs_topk_select_classes_indices)
-    node_list.append(node_topk_select_topk_indices)
-    outputs_concat_final_results = outputs['Out']
-    node_concat_final_results = onnx.helper.make_node(
-        'Gather',
-        inputs=outputs_sort_by_socre_results +
-        [outputs_topk_select_classes_indices[1]],
-        outputs=outputs_concat_final_results,
-        axis=1)
-    node_list.append(node_concat_final_results)
-    return node_list

+ 5 - 7
paddlex/cv/datasets/analysis.py

@@ -134,8 +134,7 @@ class Seg:
             self.label_value_list[id] = unique
             self.label_value_num_list[id] = counts
 
-    def _get_clipped_mean_std(self, start, end, clip_min_value,
-                              clip_max_value):
+    def _get_clipped_mean_std(self, start, end, clip_min_value, clip_max_value):
         for id in range(start, end):
             full_path_im, full_path_label = self.file_list[id]
             image, label = Compose.decode_image(full_path_im, full_path_label)
@@ -158,9 +157,9 @@ class Seg:
         self.im_std_list = [[] for i in range(len(self.file_list))]
         self.im_value_list = [[] for i in range(len(self.file_list))]
         self.im_value_num_list = [[] for i in range(len(self.file_list))]
-        self.im_height_list = np.zeros(len(self.file_list), dtype='int32')
-        self.im_width_list = np.zeros(len(self.file_list), dtype='int32')
-        self.im_channel_list = np.zeros(len(self.file_list), dtype='int32')
+        self.im_height_list = np.zeros(len(self.file_list), dtype='int64')
+        self.im_width_list = np.zeros(len(self.file_list), dtype='int64')
+        self.im_channel_list = np.zeros(len(self.file_list), dtype='int64')
         self.label_value_list = [[] for i in range(len(self.file_list))]
         self.label_value_num_list = [[] for i in range(len(self.file_list))]
 
@@ -171,8 +170,7 @@ class Seg:
             start = one_worker_file * i
             end = one_worker_file * (
                 i + 1) if i < num_workers - 1 else len(self.file_list)
-            t = threading.Thread(
-                target=self._get_image_info, args=(start, end))
+            t = threading.Thread(target=self._get_image_info, args=(start, end))
             threads.append(t)
         for t in threads:
             t.start()

+ 8 - 5
paddlex/cv/datasets/voc.py

@@ -143,12 +143,15 @@ class VOCDetection(Dataset):
                     cname = obj.find(name_tag).text.strip()
                     gt_class[i][0] = cname2cid[cname]
                     pattern = re.compile('<difficult>', re.IGNORECASE)
-                    diff_tag = pattern.findall(str(ET.tostringlist(obj)))[0][1:
-                                                                             -1]
-                    try:
-                        _difficult = int(obj.find(diff_tag).text)
-                    except Exception:
+                    diff_tag = pattern.findall(str(ET.tostringlist(obj)))
+                    if len(diff_tag) == 0:
                         _difficult = 0
+                    else:
+                        diff_tag = diff_tag[0][1:-1]
+                        try:
+                            _difficult = int(obj.find(diff_tag).text)
+                        except Exception:
+                            _difficult = 0
                     pattern = re.compile('<bndbox>', re.IGNORECASE)
                     box_tag = pattern.findall(str(ET.tostringlist(obj)))
                     if len(box_tag) == 0:

+ 7 - 0
paddlex/cv/models/load_model.py

@@ -131,6 +131,13 @@ def fix_input_shape(info, fixed_input_shape=None):
         padding = {'Padding': {}}
         if info['_Attributes']['model_type'] == 'classifier':
             pass
+        elif info['Model'].count('YOLO') > 0:
+            resize_op_index = None
+            for i in range(len(info['Transforms'])):
+                if list(info['Transforms'][i].keys())[0] == 'Resize':
+                    resize_op_index = i
+            if resize_op_index is not None:
+                info['Transforms'][resize_op_index]['Resize']['target_size'] = fixed_input_shape[0]
         else:
             resize['ResizeByShort']['short_size'] = min(fixed_input_shape)
             resize['ResizeByShort']['max_size'] = max(fixed_input_shape)

+ 2 - 2
paddlex/cv/models/utils/seg_eval.py

@@ -169,8 +169,8 @@ class ConfusionMatrix(object):
                 recall = 0
             else:
                 recall = self.confusion_matrix[c][c] / vij[c]
-            if vji[c] == 0 and vij[c] == 0:
-                f1score = 0
+            if recall + precision <= 1e-06:
+                f1_score = 0
             else:
                 f1score = 2 * precision * recall / (recall + precision)
             f1score_list.append(f1score)

+ 15 - 3
paddlex/deploy.py

@@ -35,7 +35,8 @@ class Predictor:
                  mkl_thread_num=4,
                  use_trt=False,
                  use_glog=False,
-                 memory_optimize=True):
+                 memory_optimize=True,
+                 max_trt_batch_size=1):
         """ 创建Paddle Predictor
 
             Args:
@@ -47,6 +48,7 @@ class Predictor:
                 use_trt: 是否使用TensorRT,默认False
                 use_glog: 是否启用glog日志, 默认False
                 memory_optimize: 是否启动内存优化,默认True
+                max_trt_batch_size: 在使用TensorRT时配置的最大batch size,默认1
         """
         if not osp.isdir(model_dir):
             raise Exception("[ERROR] Path {} not exist.".format(model_dir))
@@ -80,7 +82,8 @@ class Predictor:
                                            self.info['Transforms'], to_rgb)
         self.predictor = self.create_predictor(use_gpu, gpu_id, use_mkl,
                                                mkl_thread_num, use_trt,
-                                               use_glog, memory_optimize)
+                                               use_glog, memory_optimize,
+                                               max_trt_batch_size)
         # 线程池,在模型在预测时用于对输入数据以图片为单位进行并行处理
         # 主要用于batch_predict接口
         thread_num = mp.cpu_count() if mp.cpu_count() < 8 else 8
@@ -98,7 +101,8 @@ class Predictor:
                          mkl_thread_num=4,
                          use_trt=False,
                          use_glog=False,
-                         memory_optimize=True):
+                         memory_optimize=True,
+                         max_trt_batch_size=1):
         config = fluid.core.AnalysisConfig(
             os.path.join(self.model_dir, '__model__'),
             os.path.join(self.model_dir, '__params__'))
@@ -106,6 +110,14 @@ class Predictor:
         if use_gpu:
             # 设置GPU初始显存(单位M)和Device ID
             config.enable_use_gpu(100, gpu_id)
+            if use_trt:
+                config.enable_tensorrt_engine(
+                            workspace_size=1<<10,
+                            max_batch_size=max_trt_batch_size,
+                            min_subgraph_size=3,
+                            precision_mode=fluid.core.AnalysisConfig.Precision.Float32,
+                            use_static=False,
+                            use_calib_mode=False)
         else:
             config.disable_gpu()
         if use_mkl and not use_gpu:

+ 2 - 1
paddlex/tools/x2coco.py

@@ -66,7 +66,8 @@ class X2COCO(object):
         """
         assert osp.exists(image_dir), "he image folder does not exist!"
         assert osp.exists(json_dir), "The json folder does not exist!"
-        assert osp.exists(dataset_save_dir), "The save folder does not exist!"
+        if not osp.exists(dataset_save_dir):
+            os.makedirs(dataset_save_dir)
         # Convert the image files.
         new_image_dir = osp.join(dataset_save_dir, "JPEGImages")
         if osp.exists(new_image_dir):

+ 2 - 1
paddlex/tools/x2imagenet.py

@@ -36,7 +36,8 @@ class X2ImageNet(object):
         """
         assert osp.exists(image_dir), "The image folder does not exist!"
         assert osp.exists(json_dir), "The json folder does not exist!"
-        assert osp.exists(dataset_save_dir), "The save folder does not exist!"
+        if not osp.exists(dataset_save_dir):
+            os.makedirs(dataset_save_dir)
         assert len(os.listdir(dataset_save_dir)) == 0, "The save folder must be empty!"
         for img_name in os.listdir(image_dir):
             img_name_part = osp.splitext(img_name)[0]

+ 2 - 1
paddlex/tools/x2seg.py

@@ -105,7 +105,8 @@ class X2Seg(object):
         """
         assert osp.exists(image_dir), "The image folder does not exist!"
         assert osp.exists(json_dir), "The json folder does not exist!"
-        assert osp.exists(dataset_save_dir), "The save folder does not exist!"
+        if not osp.exists(dataset_save_dir):
+            os.makedirs(dataset_save_dir)
         # Convert the image files.
         new_image_dir = osp.join(dataset_save_dir, "JPEGImages")
         if osp.exists(new_image_dir):

+ 3 - 2
paddlex/tools/x2voc.py

@@ -35,7 +35,8 @@ class X2VOC(object):
         """
         assert osp.exists(image_dir), "The image folder does not exist!"
         assert osp.exists(json_dir), "The json folder does not exist!"
-        assert osp.exists(dataset_save_dir), "The save folder does not exist!"
+        if not osp.exists(dataset_save_dir):
+            os.makedirs(dataset_save_dir)
         # Convert the image files.
         new_image_dir = osp.join(dataset_save_dir, "JPEGImages")
         if osp.exists(new_image_dir):
@@ -66,7 +67,7 @@ class LabelMe2VOC(X2VOC):
             img_name_part = osp.splitext(img_name)[0]
             json_file = osp.join(json_dir, img_name_part + ".json")
             if not osp.exists(json_file):
-                os.remove(os.remove(osp.join(image_dir, img_name)))
+                os.remove(osp.join(image_dir, img_name))
                 continue
             xml_doc = minidom.Document() 
             root = xml_doc.createElement("annotation") 

+ 1 - 1
setup.py

@@ -19,7 +19,7 @@ long_description = "PaddlePaddle Entire Process Development Toolkit"
 
 setuptools.setup(
     name="paddlex",
-    version='1.2.3',
+    version='1.2.6',
     author="paddlex",
     author_email="paddlex@baidu.com",
     description=long_description,

+ 1 - 1
tutorials/slim/quant/image_classification/mobilenetv2_quant.py

@@ -28,4 +28,4 @@ dataset = pdx.datasets.ImageNet(
                 transforms=model.test_transforms)
 
 # 开始量化
-pdx.slim.export_quant_model(model, dataset, save_dir='./quant_mobilenet', cache_dir='./tmp')
+pdx.slim.export_quant_model(model, dataset, batch_size=4, batch_num=10, save_dir='./quant_mobilenet', cache_dir='./tmp')