diff --git a/test/test_models.py b/test/test_models.py index ecd24039..462787ca 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -2,7 +2,7 @@ import torch from yolort.models.backbone_utils import darknet_pan_backbone -from yolort.models.yolotr import darknet_pan_tr_backbone +from yolort.models.transformer import darknet_tan_backbone from yolort.models.anchor_utils import AnchorGenerator from yolort.models.box_head import YoloHead, PostProcess, SetCriterion @@ -112,7 +112,7 @@ def _init_test_backbone_with_pan_tr(self): backbone_name = 'darknet_s_r4_0' depth_multiple = 0.33 width_multiple = 0.5 - backbone_with_fpn_tr = darknet_pan_tr_backbone(backbone_name, depth_multiple, width_multiple) + backbone_with_fpn_tr = darknet_tan_backbone(backbone_name, depth_multiple, width_multiple) return backbone_with_fpn_tr def test_backbone_with_pan_tr(self): diff --git a/yolort/models/transformer.py b/yolort/models/transformer.py index dbeabe2b..f97a6705 100644 --- a/yolort/models/transformer.py +++ b/yolort/models/transformer.py @@ -24,7 +24,7 @@ def darknet_tan_backbone( version: str = 'v4.0', ): """ - Constructs a specified DarkNet backbone with PAN on top. Freezes the specified number of + Constructs a specified DarkNet backbone with TAN on top. Freezes the specified number of layers in the backbone. Examples:: diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py index 43c5a7a8..3b7bb95f 100644 --- a/yolort/models/yolo.py +++ b/yolort/models/yolo.py @@ -7,7 +7,7 @@ from torchvision.models.utils import load_state_dict_from_url from .backbone_utils import darknet_pan_backbone -from .transformer import darknet_pan_tr_backbone +from .transformer import darknet_tan_backbone from .anchor_utils import AnchorGenerator from .box_head import YoloHead, SetCriterion, PostProcess @@ -15,7 +15,7 @@ __all__ = ['YOLO', 'yolov5_darknet_pan_s_r31', 'yolov5_darknet_pan_m_r31', 'yolov5_darknet_pan_l_r31', 'yolov5_darknet_pan_s_r40', 'yolov5_darknet_pan_m_r40', 'yolov5_darknet_pan_l_r40', - 'yolov5_darknet_pan_s_tr'] + 'yolov5_darknet_tan_s_r40'] class YOLO(nn.Module): @@ -303,8 +303,8 @@ def yolov5_darknet_pan_l_r40(pretrained: bool = False, progress: bool = True, nu pretrained=pretrained, progress=progress, num_classes=num_classes, **kwargs) -def yolov5_darknet_pan_s_tr(pretrained: bool = False, progress: bool = True, num_classes: int = 80, - **kwargs: Any) -> YOLO: +def yolov5_darknet_tan_s_r40(pretrained: bool = False, progress: bool = True, num_classes: int = 80, + **kwargs: Any) -> YOLO: r"""yolov5 small with a transformer block model from `"dingyiwei/yolov5" `_. Args: @@ -317,7 +317,7 @@ def yolov5_darknet_pan_s_tr(pretrained: bool = False, progress: bool = True, num width_multiple = 0.5 version = 'v4.0' - backbone = darknet_pan_tr_backbone(backbone_name, depth_multiple, width_multiple, version=version) + backbone = darknet_tan_backbone(backbone_name, depth_multiple, width_multiple, version=version) anchor_grids = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119],