def __init__(self, in_channel, channel, **kwargs): super(YOLODetectionBlockV3, self).__init__(**kwargs) assert channel % 2 == 0, "channel {} cannot be divided by 2".format( channel) self.body = list() for _ in range(2): # 1x1 reduce self.body.append(_conv2d(in_channel, channel, 1, 0, 1)) # 3x3 expand self.body.append(_conv2d(channel, channel * 2, 3, 1, 1)) in_channel = channel * 2 self.body.append(_conv2d(in_channel, channel, 1, 0, 1)) self.body = nn.Sequential(*self.body) self.tip = _conv2d(channel, channel * 2, 3, 1, 1)
def __init__(self, layers, channels, classes=1000, **kwargs): super(DarknetV3, self).__init__(**kwargs) assert len(layers) == len(channels) - 1, ( "len(channels) should equal to len(layers) + 1, given {} vs {}".format( len(channels), len(layers))) self.features = list() # first 3x3 conv self.features.append(_conv2d(3, channels[0], 3, 1, 1)) for i, (nlayer, channel) in enumerate(zip(layers, channels[1:])): assert channel % 2 == 0, "channel {} cannot be divided by 2".format(channel) # add downsample conv with stride=2 self.features.append(_conv2d(channels[i], channel, 3, 1, 2)) # add nlayer basic blocks for i in range(nlayer): self.features.append(DarknetBasicBlockV3(channel, channel // 2)) self.features = nn.Sequential(*self.features) # output self.output = nn.Linear(channels[-1], classes)
def __init__(self, stages, out_channels, block_channels, channels, anchors, strides, classes, alloc_size=(128, 128), nms_thresh=0.45, nms_topk=400, post_nms=100, pos_iou_thresh=1.0, ignore_iou_thresh=0.7, **kwargs): super(YOLOV3, self).__init__(**kwargs) self._anchors = anchors self._classes = classes self._num_class = len(self._classes) self.nms_thresh = nms_thresh self.nms_topk = nms_topk self.post_nms = post_nms self._pos_iou_thresh = pos_iou_thresh self._ignore_iou_thresh = ignore_iou_thresh if pos_iou_thresh >= 1: self._target_generator = YOLOV3TargetMerger( len(classes), ignore_iou_thresh) else: raise NotImplementedError( "pos_iou_thresh({}) < 1.0 is not implemented!".format( pos_iou_thresh)) self._loss = YOLOV3Loss() self.stages = nn.ModuleList() self.transitions = nn.ModuleList() self.yolo_blocks = nn.ModuleList() self.yolo_outputs = nn.ModuleList() # note that anchors and strides should be used in reverse order for i, stage, out_channel, block_channel, channel, anchor, stride in zip( range(len(stages)), stages, out_channels, block_channels, channels, anchors[::-1], strides[::-1]): self.stages.append(stage) block = YOLODetectionBlockV3(block_channel, channel) self.yolo_blocks.append(block) output = YOLOOutputV3(out_channel, len(classes), anchor, stride, alloc_size=alloc_size) self.yolo_outputs.append(output) if i > 0: self.transitions.append(_conv2d(out_channel, channel, 1, 0, 1))
def __init__(self, in_channel, channel, **kwargs): super(DarknetBasicBlockV3, self).__init__(**kwargs) self.body = nn.Sequential(_conv2d(in_channel, channel, 1, 0, 1), _conv2d(channel, channel * 2, 3, 1, 1))