135 lines
5.3 KiB
Python
135 lines
5.3 KiB
Python
"ESPNetv2: A Light-weight, Power Efficient, and General Purpose for Semantic Segmentation"
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from core.models.base_models import eespnet, EESP
|
|
from core.nn import _ConvBNPReLU, _BNPReLU
|
|
|
|
|
|
class ESPNetV2(nn.Module):
|
|
r"""ESPNetV2
|
|
|
|
Parameters
|
|
----------
|
|
nclass : int
|
|
Number of categories for the training dataset.
|
|
backbone : string
|
|
Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
|
|
'resnet101' or 'resnet152').
|
|
norm_layer : object
|
|
Normalization layer used in backbone network (default: :class:`nn.BatchNorm`;
|
|
for Synchronized Cross-GPU BachNormalization).
|
|
aux : bool
|
|
Auxiliary loss.
|
|
|
|
Reference:
|
|
Sachin Mehta, et al. "ESPNetv2: A Light-weight, Power Efficient, and General Purpose Convolutional Neural Network."
|
|
arXiv preprint arXiv:1811.11431 (2018).
|
|
"""
|
|
|
|
def __init__(self, nclass, backbone='', aux=False, jpu=False, pretrained_base=False, **kwargs):
|
|
super(ESPNetV2, self).__init__()
|
|
self.pretrained = eespnet(pretrained=pretrained_base, **kwargs)
|
|
self.proj_L4_C = _ConvBNPReLU(256, 128, 1, **kwargs)
|
|
self.pspMod = nn.Sequential(
|
|
EESP(256, 128, stride=1, k=4, r_lim=7, **kwargs),
|
|
_PSPModule(128, 128, **kwargs))
|
|
self.project_l3 = nn.Sequential(
|
|
nn.Dropout2d(0.1),
|
|
nn.Conv2d(128, nclass, 1, bias=False))
|
|
self.act_l3 = _BNPReLU(nclass, **kwargs)
|
|
self.project_l2 = _ConvBNPReLU(64 + nclass, nclass, 1, **kwargs)
|
|
self.project_l1 = nn.Sequential(
|
|
nn.Dropout2d(0.1),
|
|
nn.Conv2d(32 + nclass, nclass, 1, bias=False))
|
|
|
|
self.aux = aux
|
|
|
|
self.__setattr__('exclusive', ['proj_L4_C', 'pspMod', 'project_l3', 'act_l3', 'project_l2', 'project_l1'])
|
|
|
|
def forward(self, x):
|
|
size = x.size()[2:]
|
|
out_l1, out_l2, out_l3, out_l4 = self.pretrained(x, seg=True)
|
|
out_l4_proj = self.proj_L4_C(out_l4)
|
|
up_l4_to_l3 = F.interpolate(out_l4_proj, scale_factor=2, mode='bilinear', align_corners=True)
|
|
merged_l3_upl4 = self.pspMod(torch.cat([out_l3, up_l4_to_l3], 1))
|
|
proj_merge_l3_bef_act = self.project_l3(merged_l3_upl4)
|
|
proj_merge_l3 = self.act_l3(proj_merge_l3_bef_act)
|
|
out_up_l3 = F.interpolate(proj_merge_l3, scale_factor=2, mode='bilinear', align_corners=True)
|
|
merge_l2 = self.project_l2(torch.cat([out_l2, out_up_l3], 1))
|
|
out_up_l2 = F.interpolate(merge_l2, scale_factor=2, mode='bilinear', align_corners=True)
|
|
merge_l1 = self.project_l1(torch.cat([out_l1, out_up_l2], 1))
|
|
|
|
outputs = list()
|
|
merge1_l1 = F.interpolate(merge_l1, scale_factor=2, mode='bilinear', align_corners=True)
|
|
outputs.append(merge1_l1)
|
|
if self.aux:
|
|
# different from paper
|
|
auxout = F.interpolate(proj_merge_l3_bef_act, size, mode='bilinear', align_corners=True)
|
|
outputs.append(auxout)
|
|
|
|
#return tuple(outputs)
|
|
return outputs[0]
|
|
|
|
# different from PSPNet
|
|
class _PSPModule(nn.Module):
|
|
def __init__(self, in_channels, out_channels=1024, sizes=(1, 2, 4, 8), **kwargs):
|
|
super(_PSPModule, self).__init__()
|
|
self.stages = nn.ModuleList(
|
|
[nn.Conv2d(in_channels, in_channels, 3, 1, 1, groups=in_channels, bias=False) for _ in sizes])
|
|
self.project = _ConvBNPReLU(in_channels * (len(sizes) + 1), out_channels, 1, 1, **kwargs)
|
|
|
|
def forward(self, x):
|
|
size = x.size()[2:]
|
|
feats = [x]
|
|
for stage in self.stages:
|
|
x = F.avg_pool2d(x, kernel_size=3, stride=2, padding=1)
|
|
upsampled = F.interpolate(stage(x), size, mode='bilinear', align_corners=True)
|
|
feats.append(upsampled)
|
|
return self.project(torch.cat(feats, dim=1))
|
|
|
|
|
|
def get_espnet(dataset='pascal_voc', backbone='', pretrained=False, root='~/.torch/models',
|
|
pretrained_base=False, **kwargs):
|
|
acronyms = {
|
|
'pascal_voc': 'pascal_voc',
|
|
'pascal_aug': 'pascal_aug',
|
|
'ade20k': 'ade',
|
|
'coco': 'coco',
|
|
'citys': 'citys',
|
|
}
|
|
from core.data.dataloader import datasets
|
|
model = ESPNetV2(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('espnet_%s_%s' % (backbone, acronyms[dataset]), root=root),
|
|
map_location=device))
|
|
return model
|
|
|
|
|
|
def get_espnet_citys(**kwargs):
|
|
return get_espnet('citys', **kwargs)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
#model = get_espnet_citys()
|
|
input = torch.rand(2, 3, 224, 224)
|
|
model =ESPNetV2(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))
|