def __init__(self, cfg):
super().__init__()
self.cfg = cfg
# ----------------------- build backbone ------------------------ #
bottom_up = getattr(resnet, cfg.backbone)(
norm=layers.get_norm(cfg.backbone_norm),
pretrained=cfg.backbone_pretrained)
del bottom_up.fc
# ----------------------- build FPN ----------------------------- #
self.backbone = layers.FPN(
bottom_up=bottom_up,
in_features=cfg.fpn_in_features,
out_channels=cfg.fpn_out_channels,
norm=cfg.fpn_norm,
top_block=layers.FPNP6(),
strides=cfg.fpn_in_strides,
channels=cfg.fpn_in_channels,
)
# ----------------------- build RPN ----------------------------- #
self.rpn = layers.RPN(cfg)
# ----------------------- build RCNN head ----------------------- #
self.rcnn = layers.RCNN(cfg)
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
# ----------------------- build backbone ------------------------ #
bottom_up = getattr(resnet, cfg.backbone)(
norm=layers.get_norm(cfg.resnet_norm),
pretrained=cfg.backbone_pretrained)
del bottom_up.fc
# ----------------------- build FPN ----------------------------- #
out_channels = 256
self.backbone = layers.FPN(
bottom_up=bottom_up,
in_features=["res2", "res3", "res4", "res5"],
out_channels=out_channels,
norm=cfg.fpn_norm,
top_block=layers.FPNP6(),
strides=[4, 8, 16, 32],
channels=[256, 512, 1024, 2048],
)
# ----------------------- build RPN ----------------------------- #
self.rpn = layers.RPN(cfg)
# ----------------------- build RCNN head ----------------------- #
self.rcnn = layers.RCNN(cfg)
def __init__(self, cfg, batch_size):
super().__init__()
self.cfg = cfg
cfg.batch_per_gpu = batch_size
self.batch_size = batch_size
# ----------------------- build the backbone ------------------------ #
bottom_up = getattr(resnet, cfg.backbone)(
norm=layers.get_norm(cfg.resnet_norm),
pretrained=cfg.backbone_pretrained)
# ------------ freeze the weights of resnet stage1 and stage 2 ------ #
if self.cfg.backbone_freeze_at >= 1:
for p in bottom_up.conv1.parameters():
p.requires_grad = False
if self.cfg.backbone_freeze_at >= 2:
for p in bottom_up.layer1.parameters():
p.requires_grad = False
# ----------------------- build the FPN ----------------------------- #
out_channels = 256
self.backbone = layers.FPN(
bottom_up=bottom_up,
in_features=["res2", "res3", "res4", "res5"],
out_channels=out_channels,
norm=cfg.fpn_norm,
top_block=layers.FPNP6(),
strides=[4, 8, 16, 32],
channels=[256, 512, 1024, 2048],
)
# ----------------------- build the RPN ----------------------------- #
self.RPN = layers.RPN(cfg)
# ----------------------- build the RCNN head ----------------------- #
self.RCNN = layers.RCNN(cfg)
# -------------------------- input Tensor --------------------------- #
self.inputs = {
"image":
mge.tensor(
np.random.random([2, 3, 224, 224]).astype(np.float32),
dtype="float32",
),
"im_info":
mge.tensor(
np.random.random([2, 5]).astype(np.float32),
dtype="float32",
),
"gt_boxes":
mge.tensor(
np.random.random([2, 100, 5]).astype(np.float32),
dtype="float32",
),
}