def __init__(self,
layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
num_classes=2,
zoom_factor=8,
use_ppm=True,
criterion=nn.CrossEntropyLoss(ignore_index=255),
pretrained=True):
super(PSPNet, self).__init__()
assert layers in [50, 101, 152]
assert 2048 % len(bins) == 0
assert num_classes > 1
assert zoom_factor in [1, 2, 4, 8]
self.zoom_factor = zoom_factor
self.use_ppm = use_ppm
self.criterion = criterion
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
# self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu, resnet.conv2, resnet.bn2, resnet.relu,
# resnet.conv3, resnet.bn3, resnet.relu, resnet.maxpool)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 2048
if use_ppm:
self.ppm = PPM(fea_dim, int(fea_dim / len(bins)), bins)
fea_dim *= 2
self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True), nn.Dropout2d(p=dropout),
nn.Conv2d(512, num_classes, kernel_size=1))
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout),
nn.Conv2d(256, num_classes, kernel_size=1))
def __init__(self,
layers=50,
dropout=0.1,
classes=2,
zoom_factor=8,
use_psa=True,
psa_type=2,
compact=False,
shrink_factor=2,
mask_h=59,
mask_w=59,
normalization_factor=1.0,
psa_softmax=True,
criterion=nn.CrossEntropyLoss(ignore_index=255),
pretrained=True):
super(PSANet, self).__init__()
assert layers in [50, 101, 152]
assert classes > 1
assert zoom_factor in [1, 2, 4, 8]
assert psa_type in [0, 1, 2]
self.zoom_factor = zoom_factor
self.use_psa = use_psa
self.criterion = criterion
if layers == 50:
resnet = models.resnet50(pretrained=pretrained)
elif layers == 101:
resnet = models.resnet101(pretrained=pretrained)
else:
resnet = models.resnet152(pretrained=pretrained)
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.conv2, resnet.bn2, resnet.relu,
resnet.conv3, resnet.bn3, resnet.relu,
resnet.maxpool)
self.layer1, self.layer2, self.layer3, self.layer4 = resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (2, 2), (2, 2), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.dilation, m.padding, m.stride = (4, 4), (4, 4), (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
fea_dim = 2048
if use_psa:
self.psa = PSA(fea_dim, 512, psa_type, compact, shrink_factor,
mask_h, mask_w, normalization_factor, psa_softmax)
fea_dim *= 2
self.cls = nn.Sequential(
nn.Conv2d(fea_dim, 512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(512, classes, kernel_size=1))
if self.training:
self.aux = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256), nn.ReLU(inplace=True),
nn.Dropout2d(p=dropout), nn.Conv2d(256, classes,
kernel_size=1))