Răsfoiți Sursa

Merge pull request #142 from FlyingQianMM/develop_qh

add coco pretrained weights  for detection and segmentation (including cityscapes)
Jason 5 ani în urmă
părinte
comite
8391440bf4

+ 2 - 2
docs/apis/models/detection.md

@@ -42,7 +42,7 @@ train(self, num_epochs, train_dataset, train_batch_size=8, eval_dataset=None, sa
 > > - **save_interval_epochs** (int): 模型保存间隔(单位:迭代轮数)。默认为20。
 > > - **log_interval_steps** (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认值为'output'。
-> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为None。
+> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',则自动下载在COCO数据集上预训练的模型权重;若为None,则不使用预训练模型。默认为None。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
 > > - **learning_rate** (float): 默认优化器的学习率。默认为1.0/8000。
 > > - **warmup_steps** (int):  默认优化器进行warmup过程的步数。默认为1000。
@@ -129,7 +129,7 @@ train(self, num_epochs, train_dataset, train_batch_size=2, eval_dataset=None, sa
 > > - **save_interval_epochs** (int): 模型保存间隔(单位:迭代轮数)。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认值为'output'。
-> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为None。
+> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',则自动下载在COCO数据集上预训练的模型权重(注意:暂未提供ResNet18的COCO预训练模型);为None,则不使用预训练模型。默认为None。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认为0.0025。
 > > - **warmup_steps** (int):  默认优化器进行warmup过程的步数。默认为500。

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

@@ -34,7 +34,7 @@ train(self, num_epochs, train_dataset, train_batch_size=1, eval_dataset=None, sa
 > > - **save_interval_epochs** (int): 模型保存间隔(单位:迭代轮数)。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认值为'output'。
-> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为None。
+> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',则自动下载在COCO数据集上预训练的模型权重(注意:暂未提供ResNet18和HRNet_W18的COCO预训练模型);若为None,则不使用预训练模型。默认为None。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认为0.00125。
 > > - **warmup_steps** (int):  默认优化器进行warmup过程的步数。默认为500。

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

@@ -12,7 +12,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'],'MobileNetV2_x1.0'。
+> > - **backbone** (str): DeepLabv3+的backbone网络,实现特征图的计算,取值范围为['Xception65', 'Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5', 'MobileNetV2_x1.0', 'MobileNetV2_x1.5', 'MobileNetV2_x2.0'],默认值为'MobileNetV2_x1.0'。
 > > - **output_stride** (int): backbone 输出特征图相对于输入的下采样倍数,一般取值为8或16。默认16。
 > > - **aspp_with_sep_conv** (bool):  decoder模块是否采用separable convolutions。默认True。
 > > - **decoder_use_sep_conv** (bool): decoder模块是否采用separable convolutions。默认True。
@@ -40,7 +40,7 @@ train(self, num_epochs, train_dataset, train_batch_size=2, eval_dataset=None, ev
 > > - **save_interval_epochs** (int): 模型保存间隔(单位:迭代轮数)。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认'output'
-> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认'IMAGENET'。
+> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',则自动下载在COCO数据集上预训练的模型权重(注意:暂未提供Xception41、MobileNetV2_x0.25、MobileNetV2_x0.5、MobileNetV2_x1.5、MobileNetV2_x2.0的COCO预训练模型);若为字符串'CITYSCAPES',则自动下载在CITYSCAPES数据集上预训练的模型权重(注意:暂未提供Xception41、MobileNetV2_x0.25、MobileNetV2_x0.5、MobileNetV2_x1.5、MobileNetV2_x2.0的CITYSCAPES预训练模型);为None,则不使用预训练模型。默认'IMAGENET'。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认的优化器:使用fluid.optimizer.Momentum优化方法,polynomial的学习率衰减策略。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认0.01。
 > > - **lr_decay_power** (float): 默认优化器学习率衰减指数。默认0.9。
@@ -209,7 +209,7 @@ train(self, num_epochs, train_dataset, train_batch_size=2, eval_dataset=None, ev
 > > - **save_interval_epochs** (int): 模型保存间隔(单位:迭代轮数)。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认'output'
-> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet数据集上预训练的模型权重;若为None,则不使用预训练模型。默认'IMAGENET'。
+> > - **pretrain_weights** (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',则自动下载在ImageNet数据集上预训练的模型权重;若为字符串'CITYSCAPES',则自动下载在CITYSCAPES图片数据上预训练的模型权重(注意:目前仅提供`width`取值为18的CITYSCAPES预训练模型);为None,则不使用预训练模型。默认'IMAGENET'。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认的优化器:使用fluid.optimizer.Momentum优化方法,polynomial的学习率衰减策略。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认0.01。
 > > - **lr_decay_power** (float): 默认优化器学习率衰减指数。默认0.9。

+ 25 - 1
paddlex/cv/models/base.py

@@ -194,14 +194,38 @@ class BaseAPI:
                 if os.path.exists(pretrain_dir):
                     os.remove(pretrain_dir)
                 os.makedirs(pretrain_dir)
+            if pretrain_weights is not None and \
+                not os.path.isdir(pretrain_weights) \
+                and not os.path.isfile(pretrain_weights):
+                if self.model_type == 'classifier':
+                    if pretrain_weights not in ['IMAGENET']:
+                        logging.warning(
+                            "Pretrain_weights for classifier should be defined as directory path or parameter file or 'IMAGENET' or None, but it is {}, so we force to set it as 'IMAGENET'".
+                            format(pretrain_weights))
+                        pretrain_weights = 'IMAGENET'
+                elif self.model_type == 'detector':
+                    if pretrain_weights not in ['IMAGENET', 'COCO']:
+                        logging.warning(
+                            "Pretrain_weights for detector should be defined as directory path or parameter file or 'IMAGENET' or 'COCO' or None, but it is {}, so we force to set it as 'IMAGENET'".
+                            format(pretrain_weights))
+                        pretrain_weights = 'IMAGENET'
+                elif self.model_type == 'segmenter':
+                    if pretrain_weights not in [
+                            'IMAGENET', 'COCO', 'CITYSCAPES'
+                    ]:
+                        logging.warning(
+                            "Pretrain_weights for segmenter should be defined as directory path or parameter file or 'IMAGENET' or 'COCO' or 'CITYSCAPES', but it is {}, so we force to set it as 'IMAGENET'".
+                            format(pretrain_weights))
+                        pretrain_weights = 'IMAGENET'
             if hasattr(self, 'backbone'):
                 backbone = self.backbone
             else:
                 backbone = self.__class__.__name__
                 if backbone == "HRNet":
                     backbone = backbone + "_W{}".format(self.width)
+            class_name = self.__class__.__name__
             pretrain_weights = get_pretrain_weights(
-                pretrain_weights, self.model_type, backbone, pretrain_dir)
+                pretrain_weights, class_name, backbone, pretrain_dir)
         if startup_prog is None:
             startup_prog = fluid.default_startup_program()
         self.exe.run(startup_prog)

+ 3 - 1
paddlex/cv/models/deeplabv3p.py

@@ -242,7 +242,9 @@ class DeepLabv3p(BaseAPI):
             log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
             save_dir (str): 模型保存路径。默认'output'。
             pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',
-                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认'IMAGENET。
+                则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',
+                则自动下载在COCO数据集上预训练的模型权重;若为字符串'CITYSCAPES',
+                则自动下载在CITYSCAPES数据集上预训练的模型权重;若为None,则不使用预训练模型。默认'IMAGENET。
             optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认的优化器:使用
                 fluid.optimizer.Momentum优化方法,polynomial的学习率衰减策略。
             learning_rate (float): 默认优化器的初始学习率。默认0.01。

+ 2 - 1
paddlex/cv/models/faster_rcnn.py

@@ -196,7 +196,8 @@ class FasterRCNN(BaseAPI):
             log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为20。
             save_dir (str): 模型保存路径。默认值为'output'。
             pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',
-                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'。
+                则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',
+                则自动下载在COCO数据集上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'。
             optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:
                 fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
             learning_rate (float): 默认优化器的初始学习率。默认为0.0025。

+ 2 - 1
paddlex/cv/models/hrnet.py

@@ -152,7 +152,8 @@ class HRNet(DeepLabv3p):
             log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为2。
             save_dir (str): 模型保存路径。默认'output'。
             pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',
-                则自动下载在IMAGENET图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'。
+                则自动下载在IMAGENET图片数据上预训练的模型权重;若为字符串'CITYSCAPES'
+                则自动下载在CITYSCAPES图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'。
             optimizer (paddle.fluid.optimizer): 优化器。当改参数为None时,使用默认的优化器:使用
                 fluid.optimizer.Momentum优化方法,polynomial的学习率衰减策略。
             learning_rate (float): 默认优化器的初始学习率。默认0.01。

+ 2 - 1
paddlex/cv/models/mask_rcnn.py

@@ -155,7 +155,8 @@ class MaskRCNN(FasterRCNN):
             log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为20。
             save_dir (str): 模型保存路径。默认值为'output'。
             pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',
-                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为None。
+                则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',
+                则自动下载在COCO数据集上预训练的模型权重;若为None,则不使用预训练模型。默认为None。
             optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:
                 fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
             learning_rate (float): 默认优化器的学习率。默认为1.0/800。

+ 88 - 10
paddlex/cv/models/utils/pretrain_weights.py

@@ -1,4 +1,5 @@
 import paddlex
+import paddlex.utils.logging as logging
 import paddlehub as hub
 import os
 import os.path as osp
@@ -75,16 +76,90 @@ image_pretrain = {
 }
 
 coco_pretrain = {
-    'UNet': 'https://paddleseg.bj.bcebos.com/models/unet_coco_v3.tgz'
+    'YOLOv3_DarkNet53_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_darknet.tar',
+    'YOLOv3_MobileNetV1_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_mobilenet_v1.tar',
+    'YOLOv3_MobileNetV3_large_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_mobilenet_v3.pdparams',
+    'YOLOv3_ResNet34_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_r34.tar',
+    'YOLOv3_ResNet50_vd_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/yolov3_r50vd_dcn.tar',
+    'FasterRCNN_ResNet50_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r50_fpn_2x.tar',
+    'FasterRCNN_ResNet50_vd_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r50_vd_fpn_2x.tar',
+    'FasterRCNN_ResNet101_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r101_fpn_2x.tar',
+    'FasterRCNN_ResNet101_vd_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r101_vd_fpn_2x.tar',
+    'FasterRCNN_HRNet_W18_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_hrnetv2p_w18_2x.tar',
+    'MaskRCNN_ResNet50_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r50_fpn_2x.tar',
+    'MaskRCNN_ResNet50_vd_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r50_vd_fpn_2x.tar',
+    'MaskRCNN_ResNet101_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r101_fpn_1x.tar',
+    'MaskRCNN_ResNet101_vd_COCO':
+    'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r101_vd_fpn_1x.tar',
+    'UNet_COCO': 'https://paddleseg.bj.bcebos.com/models/unet_coco_v3.tgz',
+    'DeepLabv3p_MobileNetV2_x1.0_COCO':
+    'https://bj.bcebos.com/v1/paddleseg/deeplab_mobilenet_x1_0_coco.tgz',
+    'DeepLabv3p_Xception65_COCO':
+    'https://paddleseg.bj.bcebos.com/models/xception65_coco.tgz'
+}
+
+cityscapes_pretrain = {
+    'DeepLabv3p_MobileNetV2_x1.0_CITYSCAPES':
+    'https://paddleseg.bj.bcebos.com/models/mobilenet_cityscapes.tgz',
+    'DeepLabv3p_Xception65_CITYSCAPES':
+    'https://paddleseg.bj.bcebos.com/models/xception65_bn_cityscapes.tgz',
+    'HRNet_W18_CITYSCAPES':
+    'https://paddleseg.bj.bcebos.com/models/hrnet_w18_bn_cityscapes.tgz'
 }
 
 
-def get_pretrain_weights(flag, model_type, backbone, save_dir):
+def get_pretrain_weights(flag, class_name, backbone, save_dir):
     if flag is None:
         return None
     elif osp.isdir(flag):
         return flag
-    elif flag == 'IMAGENET':
+    elif osp.isfile(flag):
+        return flag
+    warning_info = "{} does not support to be finetuned with weights pretrained on the {} dataset, so pretrain_weights is forced to be set to {}"
+    if flag == 'COCO':
+        if class_name == "FasterRCNN" and backbone in ['ResNet18'] or \
+            class_name == "MaskRCNN" and backbone in ['ResNet18', 'HRNet_W18'] or \
+            class_name == 'DeepLabv3p' and backbone in ['Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5', 'MobileNetV2_x1.5', 'MobileNetV2_x2.0']:
+            model_name = '{}_{}'.format(class_name, backbone)
+            logging.warning(warning_info.format(model_name, flag, 'IMAGENET'))
+            flag = 'IMAGENET'
+        elif class_name == 'HRNet':
+            logging.warning(warning_info.format(class_name, flag, 'IMAGENET'))
+            flag = 'IMAGENET'
+    elif flag == 'CITYSCAPES':
+        model_name = '{}_{}'.format(class_name, backbone)
+        if class_name == 'UNet':
+            logging.warning(warning_info.format(class_name, flag, 'COCO'))
+            flag = 'COCO'
+        if class_name == 'HRNet' and backbone.split('_')[
+                -1] in ['W30', 'W32', 'W40', 'W48', 'W60', 'W64']:
+            logging.warning(warning_info.format(backbone, flag, 'IMAGENET'))
+            flag = 'IMAGENET'
+        if class_name == 'DeepLabv3p' and backbone in [
+                'Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5',
+                'MobileNetV2_x1.5', 'MobileNetV2_x2.0'
+        ]:
+            model_name = '{}_{}'.format(class_name, backbone)
+            logging.warning(warning_info.format(model_name, flag, 'IMAGENET'))
+            flag = 'IMAGENET'
+    elif flag == 'IMAGENET' and class_name == 'UNet':
+        logging.warning(warning_info.format(class_name, flag, 'COCO'))
+        flag = 'COCO'
+
+    if flag == 'IMAGENET':
         new_save_dir = save_dir
         if hasattr(paddlex, 'pretrain_dir'):
             new_save_dir = paddlex.pretrain_dir
@@ -96,7 +171,7 @@ def get_pretrain_weights(flag, model_type, backbone, save_dir):
             backbone = 'MobileNetV3_small_x1_0_ssld'
         elif backbone == 'MobileNetV3_large_ssld':
             backbone = 'MobileNetV3_large_x1_0_ssld'
-        if model_type == 'detector':
+        if class_name in ['YOLOv3', 'FasterRCNN', 'MaskRCNN']:
             if backbone == 'ResNet50':
                 backbone = 'DetResNet50'
         assert backbone in image_pretrain, "There is not ImageNet pretrain weights for {}, you may try COCO.".format(
@@ -121,17 +196,20 @@ def get_pretrain_weights(flag, model_type, backbone, save_dir):
                 raise Exception(
                     "Unexpected error, please make sure paddlehub >= 1.6.2")
         return osp.join(new_save_dir, backbone)
-    elif flag == 'COCO':
+    elif flag in ['COCO', 'CITYSCAPES']:
         new_save_dir = save_dir
         if hasattr(paddlex, 'pretrain_dir'):
             new_save_dir = paddlex.pretrain_dir
-        url = coco_pretrain[backbone]
+        if class_name in ['YOLOv3', 'FasterRCNN', 'MaskRCNN', 'DeepLabv3p']:
+            backbone = '{}_{}'.format(class_name, backbone)
+        backbone = "{}_{}".format(backbone, flag)
+        if flag == 'COCO':
+            url = coco_pretrain[backbone]
+        elif flag == 'CITYSCAPES':
+            url = cityscapes_pretrain[backbone]
         fname = osp.split(url)[-1].split('.')[0]
         #        paddlex.utils.download_and_decompress(url, path=new_save_dir)
         #        return osp.join(new_save_dir, fname)
-
-        assert backbone in coco_pretrain, "There is not COCO pretrain weights for {}, you may try ImageNet.".format(
-            backbone)
         try:
             hub.download(backbone, save_path=new_save_dir)
         except Exception as e:
@@ -148,5 +226,5 @@ def get_pretrain_weights(flag, model_type, backbone, save_dir):
         return osp.join(new_save_dir, backbone)
     else:
         raise Exception(
-            "pretrain_weights need to be defined as directory path or `IMAGENET` or 'COCO' (download pretrain weights automatically)."
+            "pretrain_weights need to be defined as directory path or 'IMAGENET' or 'COCO' or 'Cityscapes' (download pretrain weights automatically)."
         )

+ 2 - 1
paddlex/cv/models/yolo_v3.py

@@ -188,7 +188,8 @@ class YOLOv3(BaseAPI):
             log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为10。
             save_dir (str): 模型保存路径。默认值为'output'。
             pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',
-                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'。
+                则自动下载在ImageNet图片数据上预训练的模型权重;若为字符串'COCO',
+                则自动下载在COCO数据集上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'。
             optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:
                 fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
             learning_rate (float): 默认优化器的学习率。默认为1.0/8000。

+ 28 - 21
paddlex/utils/utils.py

@@ -67,8 +67,8 @@ def parse_param_file(param_file, return_shape=True):
         f.close()
         return tuple(tensor_desc.dims)
     if tensor_desc.data_type != 5:
-        raise Exception(
-            "Unexpected data type while parse {}".format(param_file))
+        raise Exception("Unexpected data type while parse {}".format(
+            param_file))
     data_size = 4
     for i in range(len(tensor_shape)):
         data_size *= tensor_shape[i]
@@ -139,7 +139,12 @@ def load_pdparams(exe, main_prog, model_dir):
 
     vars_to_load = list()
     import pickle
-    with open(osp.join(model_dir, 'model.pdparams'), 'rb') as f:
+
+    if osp.isfile(model_dir):
+        params_file = model_dir
+    else:
+        params_file = osp.join(model_dir, 'model.pdparams')
+    with open(params_file, 'rb') as f:
         params_dict = pickle.load(f) if six.PY2 else pickle.load(
             f, encoding='latin1')
     unused_vars = list()
@@ -185,8 +190,8 @@ def is_belong_to_optimizer(var):
     import paddle.fluid as fluid
     from paddle.fluid.proto.framework_pb2 import VarType
 
-    if not (isinstance(var, fluid.framework.Parameter)
-            or var.desc.need_check_feed()):
+    if not (isinstance(var, fluid.framework.Parameter) or
+            var.desc.need_check_feed()):
         return is_persistable(var)
     return False
 
@@ -206,9 +211,8 @@ def load_pdopt(exe, main_prog, model_dir):
     if len(optimizer_var_list) > 0:
         for var in optimizer_var_list:
             if var.name not in opt_dict:
-                raise Exception(
-                    "{} is not in saved paddlex optimizer, {}".format(
-                        var.name, exception_message))
+                raise Exception("{} is not in saved paddlex optimizer, {}".
+                                format(var.name, exception_message))
             if var.shape != opt_dict[var.name].shape:
                 raise Exception(
                     "Shape of optimizer variable {} doesn't match.(Last: {}, Now: {}), {}"
@@ -227,9 +231,8 @@ def load_pdopt(exe, main_prog, model_dir):
             "There is no optimizer parameters in the model, please set the optimizer!"
         )
     else:
-        logging.info(
-            "There are {} optimizer parameters in {} are loaded.".format(
-                len(optimizer_var_list), model_dir))
+        logging.info("There are {} optimizer parameters in {} are loaded.".
+                     format(len(optimizer_var_list), model_dir))
 
 
 def load_pretrain_weights(exe,
@@ -239,6 +242,12 @@ def load_pretrain_weights(exe,
                           resume=False):
     if not osp.exists(weights_dir):
         raise Exception("Path {} not exists.".format(weights_dir))
+    if osp.isfile(weights_dir):
+        if not weights_dir.endswith('.pdparams'):
+            raise Exception("File {} is not a paddle parameter file".format(
+                weights_dir))
+        load_pdparams(exe, main_prog, weights_dir)
+        return
     if osp.exists(osp.join(weights_dir, "model.pdparams")):
         load_pdparams(exe, main_prog, weights_dir)
         if resume:
@@ -255,9 +264,8 @@ def load_pretrain_weights(exe,
         if not isinstance(var, fluid.framework.Parameter):
             continue
         if not osp.exists(osp.join(weights_dir, var.name)):
-            logging.debug(
-                "[SKIP] Pretrained weight {}/{} doesn't exist".format(
-                    weights_dir, var.name))
+            logging.debug("[SKIP] Pretrained weight {}/{} doesn't exist".
+                          format(weights_dir, var.name))
             continue
         pretrained_shape = parse_param_file(osp.join(weights_dir, var.name))
         actual_shape = tuple(var.shape)
@@ -317,9 +325,8 @@ def load_pretrain_weights(exe,
                 "There is no optimizer parameters in the model, please set the optimizer!"
             )
         else:
-            logging.info(
-                "There are {} optimizer parameters in {} are loaded.".format(
-                    len(optimizer_var_list), weights_dir))
+            logging.info("There are {} optimizer parameters in {} are loaded.".
+                         format(len(optimizer_var_list), weights_dir))
 
 
 class EarlyStop:
@@ -342,12 +349,12 @@ class EarlyStop:
             self.max = current_score
             return False
         else:
-            if (abs(self.score - current_score) < self.thresh
-                    or current_score < self.score):
+            if (abs(self.score - current_score) < self.thresh or
+                    current_score < self.score):
                 self.counter += 1
                 self.score = current_score
-                logging.debug(
-                    "EarlyStopping: %i / %i" % (self.counter, self.patience))
+                logging.debug("EarlyStopping: %i / %i" %
+                              (self.counter, self.patience))
                 if self.counter >= self.patience:
                     logging.info("EarlyStopping: Stop training")
                     return True