From c5efc77470df749beded8dcda7a14d92db998c4e Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Tue, 26 Oct 2021 02:50:41 +0800 Subject: [PATCH] Upgrade the default setting to r6.0 (#213) * Upgrade the default setting of yolort.models to r6.0 * Revert setting about testing coco evaluator --- test/test_engine.py | 4 ++-- yolort/models/__init__.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/test/test_engine.py b/test/test_engine.py index d7302370..488d8962 100644 --- a/test/test_engine.py +++ b/test/test_engine.py @@ -94,7 +94,7 @@ def test_vanilla_coco_evaluator(): coco = data_helper.get_coco_api_from_dataset(val_dataloader.dataset) coco_evaluator = COCOEvaluator(coco) # Load model - model = yolov5s(pretrained=True) + model = yolov5s(upstream_version="r4.0", pretrained=True) model.eval() for images, targets in val_dataloader: preds = model(images) @@ -116,7 +116,7 @@ def test_test_epoch_end(): val_dataloader = data_helper.get_dataloader(data_root=data_path, mode="val") # Load model - model = yolov5s(pretrained=True, annotation_path=annotation_file) + model = yolov5s(upstream_version="r4.0", pretrained=True, annotation_path=annotation_file) # test step trainer = pl.Trainer(max_epochs=1) diff --git a/yolort/models/__init__.py b/yolort/models/__init__.py index 51d7fb26..148713f1 100644 --- a/yolort/models/__init__.py +++ b/yolort/models/__init__.py @@ -41,11 +41,11 @@ def yolov5n(upstream_version: str = "r6.0", export_friendly: bool = False, **kwa return model -def yolov5s(upstream_version: str = "r4.0", export_friendly: bool = False, **kwargs: Any): +def yolov5s(upstream_version: str = "r6.0", export_friendly: bool = False, **kwargs: Any): """ Args: upstream_version (str): model released by the upstream YOLOv5. Possible values - are ["r3.1", "r4.0", "r6.0"]. Default: "r4.0". + are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0". export_friendly (bool): Deciding whether to use (ONNX/TVM) export friendly mode. Default: False. """ @@ -64,11 +64,11 @@ def yolov5s(upstream_version: str = "r4.0", export_friendly: bool = False, **kwa return model -def yolov5m(upstream_version: str = "r4.0", export_friendly: bool = False, **kwargs: Any): +def yolov5m(upstream_version: str = "r6.0", export_friendly: bool = False, **kwargs: Any): """ Args: upstream_version (str): model released by the upstream YOLOv5. Possible values - are ["r3.1", "r4.0", "r6.0"]. Default: "r4.0". + are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0". export_friendly (bool): Deciding whether to use (ONNX/TVM) export friendly mode. Default: False. """ @@ -87,11 +87,11 @@ def yolov5m(upstream_version: str = "r4.0", export_friendly: bool = False, **kwa return model -def yolov5l(upstream_version: str = "r4.0", export_friendly: bool = False, **kwargs: Any): +def yolov5l(upstream_version: str = "r6.0", export_friendly: bool = False, **kwargs: Any): """ Args: upstream_version (str): model released by the upstream YOLOv5. Possible values - are ["r3.1", "r4.0", "r6.0"]. Default: "r4.0". + are ["r3.1", "r4.0", "r6.0"]. Default: "r6.0". export_friendly (bool): Deciding whether to use (ONNX/TVM) export friendly mode. Default: False. """