diff --git a/d2go/export/logfiledb.py b/d2go/export/logfiledb.py index c462af69..a810a4d6 100644 --- a/d2go/export/logfiledb.py +++ b/d2go/export/logfiledb.py @@ -5,7 +5,7 @@ import logging import numpy as np -from mobile_cv.torch.utils_caffe2.ws_utils import ScopedWS +from detectron2.export.shared import ScopedWS logger = logging.getLogger(__name__) diff --git a/d2go/modeling/modeldef/modeldef.py b/d2go/modeling/modeldef/modeldef.py index 89044ffc..1c718d90 100644 --- a/d2go/modeling/modeldef/modeldef.py +++ b/d2go/modeling/modeldef/modeldef.py @@ -519,6 +519,54 @@ def _repeat_last(stage, n=None): "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, "basic_args": _BASIC_ARGS, }, + "FBNetV3_A_fpn": { + "trunk": FBNetV3_A[0:5], # FPN uses all 5 stages + "rpn": [[_repeat_last(FBNetV3_A[3], n=1)]], + "bbox": [FBNetV3_A[4]], + "mask": FPN_UPSAMPLE_HEAD_STAGES, + "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, + "basic_args": _BASIC_ARGS, + }, + "FBNetV3_B_fpn": { + "trunk": FBNetV3_B[0:5], # FPN uses all 5 stages + "rpn": [[_repeat_last(FBNetV3_B[3], n=1)]], + "bbox": [FBNetV3_B[4]], + "mask": FPN_UPSAMPLE_HEAD_STAGES, + "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, + "basic_args": _BASIC_ARGS, + }, + "FBNetV3_C_fpn": { + "trunk": FBNetV3_C[0:5], # FPN uses all 5 stages + "rpn": [[_repeat_last(FBNetV3_C[3], n=1)]], + "bbox": [FBNetV3_C[4]], + "mask": FPN_UPSAMPLE_HEAD_STAGES, + "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, + "basic_args": _BASIC_ARGS, + }, + "FBNetV3_D_fpn": { + "trunk": FBNetV3_D[0:5], # FPN uses all 5 stages + "rpn": [[_repeat_last(FBNetV3_D[3], n=1)]], + "bbox": [FBNetV3_D[4]], + "mask": FPN_UPSAMPLE_HEAD_STAGES, + "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, + "basic_args": _BASIC_ARGS, + }, + "FBNetV3_E_fpn": { + "trunk": FBNetV3_E[0:5], # FPN uses all 5 stages + "rpn": [[_repeat_last(FBNetV3_E[3], n=1)]], + "bbox": [FBNetV3_E[4]], + "mask": FPN_UPSAMPLE_HEAD_STAGES, + "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, + "basic_args": _BASIC_ARGS, + }, + "FBNetV3_F_fpn": { + "trunk": FBNetV3_F[0:5], # FPN uses all 5 stages + "rpn": [[_repeat_last(FBNetV3_F[3], n=1)]], + "bbox": [FBNetV3_F[4]], + "mask": FPN_UPSAMPLE_HEAD_STAGES, + "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, + "basic_args": _BASIC_ARGS, + }, "FBNetV3_G_fpn": { "trunk": FBNetV3_G[0:5], # FPN uses all 5 stages "rpn": [[_repeat_last(FBNetV3_G[3], n=1)]], @@ -527,6 +575,14 @@ def _repeat_last(stage, n=None): "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, "basic_args": _BASIC_ARGS, }, + "FBNetV3_H_fpn": { + "trunk": FBNetV3_H[0:5], # FPN uses all 5 stages + "rpn": [[_repeat_last(FBNetV3_H[3], n=1)]], + "bbox": [FBNetV3_H[4]], + "mask": FPN_UPSAMPLE_HEAD_STAGES, + "kpts": LARGE_UPSAMPLE_HEAD_D21_STAGES, + "basic_args": _BASIC_ARGS, + }, } FBNetV2ModelArch.add_archs(MODEL_ARCH_BUILTIN)