|
- """Criss-Cross Network"""
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
-
- from core.nn import CrissCrossAttention
- from core.models.segbase import SegBaseModel
- from core.models.fcn import _FCNHead
-
- #失败:NameError: name '_C' is not defined
-
- __all__ = ['CCNet', 'get_ccnet', 'get_ccnet_resnet50_citys', 'get_ccnet_resnet101_citys',
- 'get_ccnet_resnet152_citys', 'get_ccnet_resnet50_ade', 'get_ccnet_resnet101_ade',
- 'get_ccnet_resnet152_ade']
-
-
- class CCNet(SegBaseModel):
- r"""CCNet
-
- 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:
- Zilong Huang, et al. "CCNet: Criss-Cross Attention for Semantic Segmentation."
- arXiv preprint arXiv:1811.11721 (2018).
- """
-
- def __init__(self, nclass, backbone='resnet50', aux=False, pretrained_base=True, **kwargs):
- super(CCNet, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs)
- self.head = _CCHead(nclass, **kwargs)
- if aux:
- self.auxlayer = _FCNHead(1024, nclass, **kwargs)
-
- self.__setattr__('exclusive', ['head', 'auxlayer'] if aux else ['head'])
-
- def forward(self, x):
- size = x.size()[2:]
- _, _, c3, c4 = self.base_forward(x)
- outputs = list()
- x = self.head(c4)
- x = F.interpolate(x, size, mode='bilinear', align_corners=True)
- outputs.append(x)
-
- if self.aux:
- auxout = self.auxlayer(c3)
- auxout = F.interpolate(auxout, size, mode='bilinear', align_corners=True)
- outputs.append(auxout)
- return tuple(outputs)
-
-
- class _CCHead(nn.Module):
- def __init__(self, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
- super(_CCHead, self).__init__()
- self.rcca = _RCCAModule(2048, 512, norm_layer, **kwargs)
- self.out = nn.Conv2d(512, nclass, 1)
-
- def forward(self, x):
- x = self.rcca(x)
- x = self.out(x)
- return x
-
-
- class _RCCAModule(nn.Module):
- def __init__(self, in_channels, out_channels, norm_layer, **kwargs):
- super(_RCCAModule, self).__init__()
- inter_channels = in_channels // 4
- self.conva = nn.Sequential(
- nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
- norm_layer(inter_channels),
- nn.ReLU(True))
- self.cca = CrissCrossAttention(inter_channels)
- self.convb = nn.Sequential(
- nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
- norm_layer(inter_channels),
- nn.ReLU(True))
-
- self.bottleneck = nn.Sequential(
- nn.Conv2d(in_channels + inter_channels, out_channels, 3, padding=1, bias=False),
- norm_layer(out_channels),
- nn.Dropout2d(0.1))
-
- def forward(self, x, recurrence=1):
- out = self.conva(x)
- for i in range(recurrence):
- out = self.cca(out)
- out = self.convb(out)
- out = torch.cat([x, out], dim=1)
- out = self.bottleneck(out)
-
- return out
-
-
- def get_ccnet(dataset='pascal_voc', 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 = CCNet(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('ccnet_%s_%s' % (backbone, acronyms[dataset]), root=root),
- map_location=device))
- return model
-
-
- def get_ccnet_resnet50_citys(**kwargs):
- return get_ccnet('citys', 'resnet50', **kwargs)
-
-
- def get_ccnet_resnet101_citys(**kwargs):
- return get_ccnet('citys', 'resnet101', **kwargs)
-
-
- def get_ccnet_resnet152_citys(**kwargs):
- return get_ccnet('citys', 'resnet152', **kwargs)
-
-
- def get_ccnet_resnet50_ade(**kwargs):
- return get_ccnet('ade20k', 'resnet50', **kwargs)
-
-
- def get_ccnet_resnet101_ade(**kwargs):
- return get_ccnet('ade20k', 'resnet101', **kwargs)
-
-
- def get_ccnet_resnet152_ade(**kwargs):
- return get_ccnet('ade20k', 'resnet152', **kwargs)
-
-
- if __name__ == '__main__':
- # model = get_ccnet_resnet50_citys()
- # img = torch.randn(1, 3, 480, 480)
- # outputs = model(img)
- input = torch.rand(2, 3, 224, 224)
- model = CCNet(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))
|