180 lines
6.3 KiB
Python
180 lines
6.3 KiB
Python
"""Image Cascade Network"""
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from core.models.segbase import SegBaseModel
|
|
|
|
__all__ = ['ICNet', 'get_icnet', 'get_icnet_resnet50_citys',
|
|
'get_icnet_resnet101_citys', 'get_icnet_resnet152_citys']
|
|
|
|
|
|
class ICNet(SegBaseModel):
|
|
"""Image Cascade Network"""
|
|
|
|
def __init__(self, nclass, backbone='resnet50', aux=False, jpu=False, pretrained_base=True, **kwargs):
|
|
super(ICNet, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs)
|
|
self.conv_sub1 = nn.Sequential(
|
|
_ConvBNReLU(3, 32, 3, 2, **kwargs),
|
|
_ConvBNReLU(32, 32, 3, 2, **kwargs),
|
|
_ConvBNReLU(32, 64, 3, 2, **kwargs)
|
|
)
|
|
|
|
self.ppm = PyramidPoolingModule()
|
|
|
|
self.head = _ICHead(nclass, **kwargs)
|
|
|
|
self.__setattr__('exclusive', ['conv_sub1', 'head'])
|
|
|
|
def forward(self, x):
|
|
# sub 1
|
|
x_sub1 = self.conv_sub1(x)
|
|
|
|
# sub 2
|
|
x_sub2 = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
|
|
_, x_sub2, _, _ = self.base_forward(x_sub2)
|
|
|
|
# sub 4
|
|
x_sub4 = F.interpolate(x, scale_factor=0.25, mode='bilinear', align_corners=True)
|
|
_, _, _, x_sub4 = self.base_forward(x_sub4)
|
|
# add PyramidPoolingModule
|
|
x_sub4 = self.ppm(x_sub4)
|
|
outputs = self.head(x_sub1, x_sub2, x_sub4)
|
|
|
|
return tuple(outputs)
|
|
|
|
class PyramidPoolingModule(nn.Module):
|
|
def __init__(self, pyramids=[1,2,3,6]):
|
|
super(PyramidPoolingModule, self).__init__()
|
|
self.pyramids = pyramids
|
|
|
|
def forward(self, input):
|
|
feat = input
|
|
height, width = input.shape[2:]
|
|
for bin_size in self.pyramids:
|
|
x = F.adaptive_avg_pool2d(input, output_size=bin_size)
|
|
x = F.interpolate(x, size=(height, width), mode='bilinear', align_corners=True)
|
|
feat = feat + x
|
|
return feat
|
|
|
|
class _ICHead(nn.Module):
|
|
def __init__(self, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
|
|
super(_ICHead, self).__init__()
|
|
#self.cff_12 = CascadeFeatureFusion(512, 64, 128, nclass, norm_layer, **kwargs)
|
|
self.cff_12 = CascadeFeatureFusion(128, 64, 128, nclass, norm_layer, **kwargs)
|
|
self.cff_24 = CascadeFeatureFusion(2048, 512, 128, nclass, norm_layer, **kwargs)
|
|
|
|
self.conv_cls = nn.Conv2d(128, nclass, 1, bias=False)
|
|
|
|
def forward(self, x_sub1, x_sub2, x_sub4):
|
|
outputs = list()
|
|
x_cff_24, x_24_cls = self.cff_24(x_sub4, x_sub2)
|
|
outputs.append(x_24_cls)
|
|
#x_cff_12, x_12_cls = self.cff_12(x_sub2, x_sub1)
|
|
x_cff_12, x_12_cls = self.cff_12(x_cff_24, x_sub1)
|
|
outputs.append(x_12_cls)
|
|
|
|
up_x2 = F.interpolate(x_cff_12, scale_factor=2, mode='bilinear', align_corners=True)
|
|
up_x2 = self.conv_cls(up_x2)
|
|
outputs.append(up_x2)
|
|
up_x8 = F.interpolate(up_x2, scale_factor=4, mode='bilinear', align_corners=True)
|
|
outputs.append(up_x8)
|
|
# 1 -> 1/4 -> 1/8 -> 1/16
|
|
outputs.reverse()
|
|
|
|
return outputs
|
|
|
|
|
|
class _ConvBNReLU(nn.Module):
|
|
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1,
|
|
groups=1, norm_layer=nn.BatchNorm2d, bias=False, **kwargs):
|
|
super(_ConvBNReLU, self).__init__()
|
|
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
|
|
self.bn = norm_layer(out_channels)
|
|
self.relu = nn.ReLU(True)
|
|
|
|
def forward(self, x):
|
|
x = self.conv(x)
|
|
x = self.bn(x)
|
|
x = self.relu(x)
|
|
return x
|
|
|
|
|
|
class CascadeFeatureFusion(nn.Module):
|
|
"""CFF Unit"""
|
|
|
|
def __init__(self, low_channels, high_channels, out_channels, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
|
|
super(CascadeFeatureFusion, self).__init__()
|
|
self.conv_low = nn.Sequential(
|
|
nn.Conv2d(low_channels, out_channels, 3, padding=2, dilation=2, bias=False),
|
|
norm_layer(out_channels)
|
|
)
|
|
self.conv_high = nn.Sequential(
|
|
nn.Conv2d(high_channels, out_channels, 1, bias=False),
|
|
norm_layer(out_channels)
|
|
)
|
|
self.conv_low_cls = nn.Conv2d(out_channels, nclass, 1, bias=False)
|
|
|
|
def forward(self, x_low, x_high):
|
|
x_low = F.interpolate(x_low, size=x_high.size()[2:], mode='bilinear', align_corners=True)
|
|
x_low = self.conv_low(x_low)
|
|
x_high = self.conv_high(x_high)
|
|
x = x_low + x_high
|
|
x = F.relu(x, inplace=True)
|
|
x_low_cls = self.conv_low_cls(x_low)
|
|
|
|
return x, x_low_cls
|
|
|
|
|
|
def get_icnet(dataset='citys', backbone='resnet50', pretrained=False, root='~/.torch/models',
|
|
pretrained_base=True, **kwargs):
|
|
acronyms = {
|
|
'pascal_voc': 'pascal_voc',
|
|
'pascal_aug': 'pascal_aug',
|
|
'ade20k': 'ade',
|
|
'coco': 'coco',
|
|
'citys': 'citys',
|
|
}
|
|
from ..data.dataloader import datasets
|
|
model = ICNet(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
|
|
if pretrained:
|
|
from .model_store import get_model_file
|
|
device = torch.device(kwargs['local_rank'])
|
|
model.load_state_dict(torch.load(get_model_file('icnet_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
|
map_location=device))
|
|
return model
|
|
|
|
|
|
def get_icnet_resnet50_citys(**kwargs):
|
|
return get_icnet('citys', 'resnet50', **kwargs)
|
|
|
|
|
|
def get_icnet_resnet101_citys(**kwargs):
|
|
return get_icnet('citys', 'resnet101', **kwargs)
|
|
|
|
|
|
def get_icnet_resnet152_citys(**kwargs):
|
|
return get_icnet('citys', 'resnet152', **kwargs)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# img = torch.randn(1, 3, 256, 256)
|
|
# model = get_icnet_resnet50_citys()
|
|
# outputs = model(img)
|
|
input = torch.rand(2, 3, 224, 224)
|
|
model = ICNet(4, pretrained_base=False)
|
|
# target = torch.zeros(4, 512, 512).cuda()
|
|
# model.eval()
|
|
# print(model)
|
|
loss = model(input)
|
|
#print(loss, loss.shape)
|
|
|
|
# from torchsummary import summary
|
|
#
|
|
# summary(model, (3, 224, 224)) # 打印表格，按顺序输出每层的输出形状和参数
|
|
import torch
|
|
from thop import profile
|
|
from torchsummary import summary
|
|
|
|
flop, params = profile(model, input_size=(1, 3, 512, 512))
|
|
print('flops:{:.3f}G\nparams:{:.3f}M'.format(flop / 1e9, params / 1e6)) |