def __init__(self, cfg):
super().__init__()
self.mean, self.std = cfg.MODEL.PIXEL_MEAN, cfg.MODEL.PIXEL_STD
self.register_buffer("pixel_mean", torch.Tensor(cfg.MODEL.PIXEL_MEAN).view(-1, 1, 1))
self.register_buffer("pixel_std", torch.Tensor(cfg.MODEL.PIXEL_STD).view(-1, 1, 1))
self.backbone = build_backbone(cfg)
self.proposal_generator = build_proposal_generator(
cfg, self.backbone.output_shape()) # TODO: change to a more precise name
def __init__(self, cfg):
super().__init__()
self.device = torch.device(cfg.MODEL.DEVICE)
self.backbone = build_backbone(cfg)
self.proposal_generator = build_proposal_generator(
cfg, self.backbone.output_shape())
self.roi_heads = build_roi_heads(cfg, self.backbone.output_shape())
self.mask_threshold = cfg.MODEL.ROI_MASK_HEAD.MASK_THRESHOLD
self.nms = cfg.MODEL.ROI_MASK_HEAD.NMS
self.depth_head_on = cfg.MODEL.DEPTH_ON
if self.depth_head_on:
self.depth_head = build_depth_head(cfg)
self.camera_on = cfg.MODEL.CAMERA_ON
if self.camera_on:
self.camera_head = build_camera_head(cfg)
self.input_format = cfg.INPUT.FORMAT
assert len(cfg.MODEL.PIXEL_MEAN) == len(cfg.MODEL.PIXEL_STD)
num_channels = len(cfg.MODEL.PIXEL_MEAN)
pixel_mean = (torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(
num_channels, 1, 1))
pixel_std = (torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(
num_channels, 1, 1))
self.normalizer = lambda x: (x - pixel_mean) / pixel_std
self.embedding_on = cfg.MODEL.EMBEDDING_ON
if self.embedding_on:
self._asnet_on = (
cfg.MODEL.ROI_EMBEDDING_HEAD.NAME == "EmbeddingRCNNASNetHead")
self.embedding_loss_weight = cfg.MODEL.ROI_EMBEDDING_HEAD.LOSS_WEIGHT
if cfg.MODEL.ROI_EMBEDDING_HEAD.LOSS_TYPE == "TripletLoss":
if not self._asnet_on:
self.embedding_loss = OnlineTripletLoss(
cfg.MODEL.ROI_EMBEDDING_HEAD.MARGIN,
cfg.MODEL.DEVICE,
selector_type=cfg.MODEL.ROI_EMBEDDING_HEAD.
TRIPLET_SELECTOR_TYPE,
)
else:
self.embedding_loss = CooperativeTripletLoss(
cfg.MODEL.ROI_EMBEDDING_HEAD.MARGIN,
cfg.MODEL.DEVICE,
selector_type=cfg.MODEL.ROI_EMBEDDING_HEAD.
TRIPLET_SELECTOR_TYPE,
)
else:
raise NotImplementedError
self._eval_gt_box = cfg.TEST.EVAL_GT_BOX
self.to(self.device)
self._freeze = cfg.MODEL.FREEZE
for layers in self._freeze:
layer = layers.split(".")
final = self
for l in layer:
final = getattr(final, l)
for params in final.parameters():
params.requires_grad = False
def __init__(self, cfg):
super().__init__()
self.backbone = build_backbone(cfg)
self.proposal_generator = build_proposal_generator(
cfg, self.backbone.output_shape())
self.roi_heads = build_roi_heads(cfg, self.backbone.output_shape())
self.vis_period = cfg.VIS_PERIOD
self.input_format = cfg.INPUT.FORMAT
assert len(cfg.MODEL.PIXEL_MEAN) == len(cfg.MODEL.PIXEL_STD)
self.register_buffer("pixel_mean",
torch.Tensor(cfg.MODEL.PIXEL_MEAN).view(-1, 1, 1))
self.register_buffer("pixel_std",
torch.Tensor(cfg.MODEL.PIXEL_STD).view(-1, 1, 1))
def __init__(self, cfg):
super(ConPropNetProposalNetwork, self).__init__(cfg)
self.device = torch.device(cfg.MODEL.DEVICE)
self.backbone = build_backbone(cfg)
self.proposal_generator = build_proposal_generator(
cfg, self.backbone.output_shape())
pixel_mean = torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(
-1, 1, 1)
pixel_std = torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(
-1, 1, 1)
self.normalizer = lambda x: (x - pixel_mean) / pixel_std
self.to(self.device)
def __init__(self, cfg):
super().__init__()
self.device = torch.device(cfg.MODEL.DEVICE)
self.backbone = build_backbone(cfg)
self.proposal_generator = build_proposal_generator(
cfg, self.backbone.output_shape())
self.roi_heads = build_roi_heads(cfg, self.backbone.output_shape())
assert len(cfg.MODEL.PIXEL_MEAN) == len(cfg.MODEL.PIXEL_STD)
num_channels = len(cfg.MODEL.PIXEL_MEAN)
pixel_mean = torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(
num_channels, 1, 1)
pixel_std = torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(
num_channels, 1, 1)
self.normalizer = lambda x: (x - pixel_mean) / pixel_std
self.to(self.device)
def __init__(self, cfg):
super().__init__()
self.device = torch.device(cfg.MODEL.DEVICE)
self.backbone = build_backbone(cfg)
self.proposal_generator = build_proposal_generator(
cfg, self.backbone.output_shape())
pixel_mean = torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(
-1, 1, 1)
pixel_std = torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(
-1, 1, 1)
self.normalizer = lambda x: (x - pixel_mean) / pixel_std
self.to(self.device)
self.vis_period = cfg.VIS_PERIOD
self.input_format = cfg.INPUT.FORMAT
def __init__(self, cfg):
super().__init__()
self.device = torch.device(cfg.MODEL.DEVICE)
# build feature extraction backbone
self.backbone = build_backbone(cfg)
# build classification model
if cfg.MODEL.MULTI_TASK.CLASSIFICATION_ON:
self.classifier_in_features = cfg.MODEL.MULTI_TASK.CLASSIFICATION_IN_FEATURES
self.classifier = build_multilabel_classifier(cfg)
else:
self.classifier = None
# build segmentation model
if cfg.MODEL.MULTI_TASK.SEGMENTATION_ON:
self.metal_segmentation_in_features = cfg.MODEL.MULTI_TASK.SEGMENTATION_IN_FEATURES
self.metal_segmentation = build_metal_segmentation_model(cfg, self.backbone.out_feature_strides)
else:
self.metal_segmentation = None
# build object detection model
if cfg.MODEL.MULTI_TASK.DETECTION_ON:
self.proposal_generator = build_proposal_generator(cfg, self.backbone.output_shape())
self.roi_heads = build_roi_heads(cfg, self.backbone.output_shape())
else:
self.proposal_generator = None
self.roi_heads = None
# TODO: build multi-task layer
self.multi_loss_layer = build_multitask_loss_layer(cfg)
# other setting
self.vis_period = cfg.VIS_PERIOD
self.input_format = cfg.INPUT.FORMAT
assert len(cfg.MODEL.PIXEL_MEAN) == len(cfg.MODEL.PIXEL_STD)
num_channels = len(cfg.MODEL.PIXEL_MEAN)
pixel_mean = torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(num_channels, 1, 1)
pixel_std = torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(num_channels, 1, 1)
self.normalizer = lambda x: (x - pixel_mean) / pixel_std
self.to(self.device) # all to cuda
def __init__(self, args, cfg, device, num_max_regions):
super(RFGenerator, self).__init__()
self.device = device
self.cfg = cfg
self.backbone = build_backbone(self.cfg)
self.pooler_resolution = 14
self.canonical_level = 4
self.canonical_scale_factor = 2 ** self.canonical_level
self.pooler_scales = (1 / self.canonical_scale_factor,)
self.sampling_ratio = 0
self.proposal_generator = build_proposal_generator(self.cfg, self.backbone.output_shape())
self.roi_pooler = ROIPooler(
output_size=self.pooler_resolution,
scales=self.pooler_scales,
sampling_ratio=self.sampling_ratio,
pooler_type="ROIPool"
)
self.num_max_regions = num_max_regions
self.args = args
import torch from detectron2.modeling import build_backbone, build_proposal_generator, build_roi_heads from backbone import cfg from backbone import build_shufflenetv2_fpn_backbone backbone_model = build_backbone(cfg) print(backbone_model) torch.save(backbone_model.state_dict(), '1.backbone.pth') proposal_model = build_proposal_generator(cfg, backbone_model.output_shape()) print(proposal_model) torch.save(proposal_model.state_dict(), '2.rpn.pth') roi_model = build_roi_heads(cfg, backbone_model.output_shape()) print(roi_model) torch.save(roi_model.state_dict(), '3.roi.pth')