def __init__(self, num_features_in, num_anchors=9, feature_size=256):
super(RegressionModel, self).__init__()
self.conv1 = MetaConv2d(num_features_in,
feature_size,
kernel_size=3,
padding=1)
self.act1 = nn.ReLU()
self.conv2 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
padding=1)
self.act2 = nn.ReLU()
self.conv3 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
padding=1)
self.act3 = nn.ReLU()
self.conv4 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
padding=1)
self.act4 = nn.ReLU()
self.output = MetaConv2d(feature_size,
num_anchors * 4,
kernel_size=3,
padding=1)
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
act=GroupSort(2, axis=1)):
super(Bottleneck, self).__init__()
self.conv1 = MetaConv2d(inplanes, planes, kernel_size=1, bias=False)
self.conv2 = MetaConv2d(planes,
planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
self.conv3 = MetaConv2d(planes, planes * 4, kernel_size=1, bias=False)
self.act = act
self.downsample = downsample
self.stride = stride
self.bn1 = MetaBatchNorm2d(planes)
self.bn2 = MetaBatchNorm2d(planes)
self.bn3 = MetaBatchNorm2d(planes * 4)
def __init__(self, C3_size, C4_size, C5_size, feature_size=256):
super(PyramidFeatures, self).__init__()
# upsample C5 to get P5 from the FPN paper
self.P5_1 = MetaConv2d(C5_size,
feature_size,
kernel_size=1,
stride=1,
padding=0)
self.P5_upsampled = nn.Upsample(scale_factor=2, mode='nearest')
self.P5_2 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
stride=1,
padding=1)
# add P5 elementwise to C4
self.P4_1 = MetaConv2d(C4_size,
feature_size,
kernel_size=1,
stride=1,
padding=0)
self.P4_upsampled = nn.Upsample(scale_factor=2, mode='nearest')
self.P4_2 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
stride=1,
padding=1)
# add P4 elementwise to C3
self.P3_1 = MetaConv2d(C3_size,
feature_size,
kernel_size=1,
stride=1,
padding=0)
self.P3_2 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
stride=1,
padding=1)
# "P6 is obtained via a 3x3 stride-2 conv on C5"
self.P6 = MetaConv2d(C5_size,
feature_size,
kernel_size=3,
stride=2,
padding=1)
# "P7 is computed by applying ReLU followed by a 3x3 stride-2 conv on P6"
self.P7_1 = nn.ReLU()
self.P7_2 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
stride=2,
padding=1)
def __init__(self,
num_features_in,
num_anchors=9,
num_classes=80,
prior=0.01,
feature_size=256):
super(ClassificationModel, self).__init__()
self.num_classes = num_classes
self.num_anchors = num_anchors
self.conv1 = MetaConv2d(num_features_in,
feature_size,
kernel_size=3,
padding=1)
self.act1 = nn.ReLU()
self.conv2 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
padding=1)
self.act2 = nn.ReLU()
self.conv3 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
padding=1)
self.act3 = nn.ReLU()
self.conv4 = MetaConv2d(feature_size,
feature_size,
kernel_size=3,
padding=1)
self.act4 = nn.ReLU()
self.output = MetaConv2d(feature_size,
num_anchors * num_classes,
kernel_size=3,
padding=1)
self.output_act = nn.Sigmoid()
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
act=nn.ReLU()):
super(BasicBlock, self).__init__()
self.conv1 = MetaConv2d(inplanes,
planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
self.act = act
self.conv2 = MetaConv2d(planes,
planes,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.downsample = downsample
self.stride = stride
def _make_layer(self,
block,
planes,
blocks,
stride=1,
act=nn.ReLU(inplace=True),
conv=nn.Conv2d):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
MetaConv2d(self.inplanes,
planes * block.expansion,
kernel_size=1,
stride=stride,
bias=False),
MetaBatchNorm2d(planes * block.expansion),
)
layers = [block(self.inplanes, planes, stride, downsample, act=act)]
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes, act=act))
return nn.Sequential(*layers)
def __init__(self,
num_classes,
block,
layers,
act=nn.ReLU(),
conv=nn.Conv2d):
self.inplanes = 4
super(ResNetReduced, self).__init__()
self.conv1 = MetaConv2d(3,
4,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.act = act
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 4, layers[0])
self.layer2 = self._make_layer(block, 8, layers[1], stride=2)
self.layer3 = self._make_layer(block, 16, layers[2], stride=2)
self.layer4 = self._make_layer(block, 32, layers[3], stride=2)
self.bn1 = MetaBatchNorm2d(4)
if block == BasicBlock:
fpn_sizes = [
self.layer2[layers[1] - 1].conv2.out_channels,
self.layer3[layers[2] - 1].conv2.out_channels,
self.layer4[layers[3] - 1].conv2.out_channels
]
elif block == Bottleneck:
fpn_sizes = [
self.layer2[layers[1] - 1].conv3.out_channels,
self.layer3[layers[2] - 1].conv3.out_channels,
self.layer4[layers[3] - 1].conv3.out_channels
]
else:
raise ValueError(f"Block type {block} not understood")
self.fpn = PyramidFeatures(fpn_sizes[0],
fpn_sizes[1],
fpn_sizes[2],
feature_size=32)
self.regressionModel = RegressionModel(32,
num_anchors=15,
feature_size=32)
self.classificationModel = ClassificationModel(32,
num_classes=num_classes,
num_anchors=15,
feature_size=32)
self.anchors = Anchors()
self.regressBoxes = BBoxTransform()
self.clipBoxes = ClipBoxes()
self.focalLoss = losses.FocalLoss()
for m in self.modules():
if isinstance(m, conv):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, MetaBatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
prior = 0.01
self.classificationModel.output.weight.data.fill_(0)
self.classificationModel.output.bias.data.fill_(-math.log(
(1.0 - prior) / prior))
self.regressionModel.output.weight.data.fill_(0)
self.regressionModel.output.bias.data.fill_(0)
self.freeze_bn()