|
- "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))
|