def __init__(self, n_classes=19, n_blocks=[3, 4, 23, 3], pyramids=[6, 3, 2, 1], input_size=[713, 713], use_bn=True, output_features=False, output_all=False, use_original_base=False):
super(PSPNetClustering, self).__init__()
self.input_size = input_size
self.n_classes = n_classes
self.output_features = output_features
self.output_for_cluster = False
self.output_all = output_all
if use_original_base:
self.fcn = ResNetForPsp(Bottleneck, n_blocks)
else:
self.fcn = _DilatedFCN(n_blocks=n_blocks, use_bn=use_bn)
self.ppm = _PyramidPoolModule(
in_channels=2048, pyramids=pyramids, use_bn=use_bn)
self.final = nn.Sequential(OrderedDict([
('conv5_4', _ConvBatchNormReLU(4096, 512,
3, 1, 1, 1, relu=False, use_bn=use_bn)),
('drop5_4', nn.Dropout2d(p=0.1)),
]))
self.conv6relu = nn.ReLU()
self.conv6 = nn.Conv2d(512, n_classes, 1, stride=1, padding=0)
self.aux = nn.Sequential(OrderedDict([
('conv4_aux', _ConvBatchNormReLU(
1024, 256, 3, 1, 1, 1, relu=False, use_bn=use_bn)),
('drop4_aux', nn.Dropout2d(p=0.1)),
]))
self.conv6_1relu = nn.ReLU()
self.conv6_1 = nn.Conv2d(256, n_classes, 1, stride=1, padding=0)
def __init__(self, n_blocks):
super(_DilatedFCN, self).__init__()
self.layer1 = nn.Sequential(
OrderedDict([
("conv1", _ConvBatchNormReLU(3, 64, 3, 2, 1, 1)),
("conv2", _ConvBatchNormReLU(64, 64, 3, 1, 1, 1)),
("conv3", _ConvBatchNormReLU(64, 128, 3, 1, 1, 1)),
("pool", nn.MaxPool2d(3, 2, 1)),
]))
self.layer2 = _ResBlock(n_blocks[0], 128, 64, 256, 1, 1)
self.layer3 = _ResBlock(n_blocks[1], 256, 128, 512, 2, 1)
self.layer4 = _ResBlock(n_blocks[2], 512, 256, 1024, 1, 2)
self.layer5 = _ResBlock(n_blocks[3], 1024, 512, 2048, 1, 4)
def __init__(self, n_blocks, use_bn=True):
super(_DilatedFCN, self).__init__()
self.layer1 = nn.Sequential(OrderedDict([
('conv1', _ConvBatchNormReLU(3, 64, 3, 2, 1, 1, use_bn=use_bn)),
('conv2', _ConvBatchNormReLU(64, 64, 3, 1, 1, 1, use_bn=use_bn)),
('conv3', _ConvBatchNormReLU(64, 128, 3, 1, 1, 1, use_bn=use_bn)),
('pool', nn.MaxPool2d(3, 2, 1))
]))
self.layer2 = _ResBlock(n_blocks[0], 128, 64, 256, 1, 1, use_bn=use_bn)
self.layer3 = _ResBlock(
n_blocks[1], 256, 128, 512, 2, 1, use_bn=use_bn)
self.layer4 = _ResBlock(
n_blocks[2], 512, 256, 1024, 1, 2, use_bn=use_bn)
self.layer5 = _ResBlock(
n_blocks[3], 1024, 512, 2048, 1, 4, use_bn=use_bn)
def __init__(self, in_channels, pyramids=[6, 3, 2, 1], use_bn=True):
super(_PyramidPoolModule, self).__init__()
out_channels = in_channels // len(pyramids)
self.stages = nn.ModuleList([
nn.Sequential(OrderedDict([
('pool', nn.AdaptiveAvgPool2d(output_size=p)),
('conv', _ConvBatchNormReLU(in_channels, out_channels, 1, 1, 0, 1, use_bn=use_bn)), ]))
for p in pyramids
])
def __init__(self, n_classes, n_blocks, pyramids):
super(PSPNet, self).__init__()
self.n_classes = n_classes
self.fcn = _DilatedFCN(n_blocks=n_blocks)
self.ppm = _PyramidPoolModule(in_channels=2048, pyramids=pyramids)
# Main branch
self.final = nn.Sequential(
OrderedDict([
("conv5_4", _ConvBatchNormReLU(4096, 512, 3, 1, 1, 1)),
("drop5_4", nn.Dropout2d(p=0.1)),
("conv6", nn.Conv2d(512, n_classes, 1, stride=1, padding=0)),
]))
# Auxiliary branch
self.aux = nn.Sequential(
OrderedDict([
("conv4_aux", _ConvBatchNormReLU(1024, 256, 3, 1, 1, 1)),
("drop4_aux", nn.Dropout2d(p=0.1)),
("conv6_1", nn.Conv2d(256, n_classes, 1, stride=1, padding=0)),
]))
def __init__(self, in_channels, pyramids=[6, 3, 2, 1]):
super(_PyramidPoolModule, self).__init__()
out_channels = in_channels // len(pyramids)
self.stages = nn.Module()
for i, p in enumerate(pyramids):
self.stages.add_module(
"s{}".format(i),
nn.Sequential(
OrderedDict([
("pool", nn.AdaptiveAvgPool2d(output_size=p)),
(
"conv",
_ConvBatchNormReLU(in_channels, out_channels, 1, 1,
0, 1),
),
])),
)
