def __init__(self, num_classes, network='efficientdet-d0', D_bifpn=3, W_bifpn=88): super(EfficientDetBiFPN, self).__init__() self.backbone = EfficientNet.get_network_from_name(MODEL_MAP[network]) self.neck = BIFPN(in_channels=self.backbone.get_list_features()[-5:], out_channels=W_bifpn, stack=D_bifpn, num_outs=5) self.bbox_head = RetinaHead(num_classes=num_classes, in_channels=W_bifpn) self.anchors = Anchors() self.regressBoxes = BBoxTransform() self.clipBoxes = ClipBoxes() for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels m.weight.data.normal_(0, math.sqrt(2. / n)) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() self.freeze_bn()
def __init__(self, num_classes, network='efficientdet-d0', D_bifpn=3, W_bifpn=88, D_class=3, is_training=True, threshold=0.01, iou_threshold=0.5): super(EfficientDet, self).__init__() # self.backbone = EfficientNet.from_pretrained(MODEL_MAP[network]) self.backbone = EfficientNet.get_network_from_name(MODEL_MAP[network]) # print backbone parameters # params = list(self.backbone.named_parameters()) # for param_key, param_value in params: # print("{}, {}".format(param_key, param_value.shape)) # # for features in self.backbone.get_list_features(): # print(features) self.is_training = is_training self.neck = BIFPN(in_channels=self.backbone.get_list_features()[-5:], out_channels=W_bifpn, stack=D_bifpn, num_outs=5) self.bbox_head = RetinaHead(num_classes=num_classes, in_channels=W_bifpn) self.anchors = Anchors() self.regressBoxes = BBoxTransform() self.clipBoxes = ClipBoxes() self.threshold = threshold self.iou_threshold = iou_threshold for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels m.weight.data.normal_(0, math.sqrt(2. / n)) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() self.freeze_bn() self.criterion = FocalLoss()