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))