Bladeren bron

add faster_rcnn.py and yolov3.py code

sunyanfang01 5 jaren geleden
bovenliggende
commit
9efb01006c
2 gewijzigde bestanden met toevoegingen van 8 en 6 verwijderingen
  1. 4 3
      paddlex/cv/models/faster_rcnn.py
  2. 4 3
      paddlex/cv/models/yolo_v3.py

+ 4 - 3
paddlex/cv/models/faster_rcnn.py

@@ -176,7 +176,7 @@ class FasterRCNN(BaseAPI):
             log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为20。
             save_dir (str): 模型保存路径。默认值为'output'。
             pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',
-                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为None
+                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'
             optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:
                 fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
             learning_rate (float): 默认优化器的初始学习率。默认为0.0025。
@@ -194,11 +194,12 @@ class FasterRCNN(BaseAPI):
         if metric is None:
             if isinstance(train_dataset, paddlex.datasets.CocoDetection):
                 metric = 'COCO'
-            elif isinstance(train_dataset, paddlex.datasets.VOCDetection):
+            elif isinstance(train_dataset, paddlex.datasets.VOCDetection) or \
+                    isinstance(train_dataset, paddlex.datasets.EasyDataDet):
                 metric = 'VOC'
             else:
                 raise ValueError(
-                    "train_dataset should be datasets.VOCDetection or datasets.COCODetection."
+                    "train_dataset should be datasets.VOCDetection or datasets.COCODetection or datasets.EasyDataDet."
                 )
         assert metric in ['COCO', 'VOC'], "Metric only support 'VOC' or 'COCO'"
         self.metric = metric

+ 4 - 3
paddlex/cv/models/yolo_v3.py

@@ -175,7 +175,7 @@ class YOLOv3(BaseAPI):
             log_interval_steps (int): 训练日志输出间隔(单位:迭代次数)。默认为10。
             save_dir (str): 模型保存路径。默认值为'output'。
             pretrain_weights (str): 若指定为路径时,则加载路径下预训练模型;若为字符串'IMAGENET',
-                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为None
+                则自动下载在ImageNet图片数据上预训练的模型权重;若为None,则不使用预训练模型。默认为'IMAGENET'
             optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时,使用默认优化器:
                 fluid.layers.piecewise_decay衰减策略,fluid.optimizer.Momentum优化方法。
             learning_rate (float): 默认优化器的学习率。默认为1.0/8000。
@@ -198,11 +198,12 @@ class YOLOv3(BaseAPI):
         if metric is None:
             if isinstance(train_dataset, paddlex.datasets.CocoDetection):
                 metric = 'COCO'
-            elif isinstance(train_dataset, paddlex.datasets.VOCDetection):
+            elif isinstance(train_dataset, paddlex.datasets.VOCDetection) or \
+                    isinstance(train_dataset, paddlex.datasets.EasyDataDet):
                 metric = 'VOC'
             else:
                 raise ValueError(
-                    "train_dataset should be datasets.VOCDetection or datasets.COCODetection."
+                    "train_dataset should be datasets.VOCDetection or datasets.COCODetection or datasets.EasyDataDet."
                 )
         assert metric in ['COCO', 'VOC'], "Metric only support 'VOC' or 'COCO'"
         self.metric = metric