"""
MindSpore implementation of `SKNet`.
Refer to Selective Kernel Networks.
"""
from typing import Optional, Type, List, Dict, Union
from mindspore import nn, Tensor
from .layers.selective_kernel import SelectiveKernel
from .utils import load_pretrained
from .registry import register_model
from .resnet import ResNet
__all__ = [
"SKNet",
"sk_resnet18",
"sk_resnet34",
"sk_resnet50",
"sk_resnext50_32x4d"
]
def _cfg(url='', **kwargs):
return {
'url': url,
'num_classes': 1000,
'first_conv': 'conv1', 'classifier': 'classifier',
**kwargs
}
default_cfgs = {
'sk_resnet18': _cfg(url=''),
'sk_resnet34': _cfg(url=''),
'sk_resnet50': _cfg(url=''),
'sk_resnext50_32X4d': _cfg(url='')
}
class SelectiveKernelBasic(nn.Cell):
"""build basic block of sknet"""
expansion = 1
def __init__(self,
in_channels: int,
out_channels: int,
stride: int = 1,
groups: int = 1,
down_sample: Optional[nn.Cell] = None,
base_width: int = 64,
norm: Optional[nn.Cell] = None,
sk_kwargs: Optional[Dict] = None
):
super().__init__()
if norm is None:
norm = nn.BatchNorm2d
if sk_kwargs is None:
sk_kwargs = {}
assert groups == 1, 'BasicBlock only supports cardinality of 1'
assert base_width == 64, 'BasicBlock doest not support changing base width'
self.conv1 = SelectiveKernel(
in_channels, out_channels, stride=stride, **sk_kwargs)
self.conv2 = nn.SequentialCell([
nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=3, padding=1, pad_mode='pad'),
norm(out_channels * self.expansion)
])
self.relu = nn.ReLU()
self.down_sample = down_sample
def construct(self, x: Tensor) -> Tensor:
identity = x
out = self.conv1(x)
out = self.conv2(out)
if self.down_sample is not None:
identity = self.down_sample(x)
out += identity
out = self.relu(out)
return out
class SelectiveKernelBottleneck(nn.Cell):
"""build the bottleneck of the sknet"""
expansion = 4
def __init__(self,
in_channels: int,
out_channels: int,
stride: int = 1,
down_sample: Optional[nn.Cell] = None,
groups: int = 1,
base_width: int = 64,
norm: Optional[nn.Cell] = None,
sk_kwargs: Optional[Dict] = None,
):
super().__init__()
if norm is None:
norm = nn.BatchNorm2d
if sk_kwargs is None:
sk_kwargs = {}
width = int(out_channels * (base_width / 64.0)) * groups
self.conv1 = nn.SequentialCell([
nn.Conv2d(in_channels, width, kernel_size=1),
norm(width)
])
self.conv2 = SelectiveKernel(
width, width, stride=stride, groups=groups, **sk_kwargs)
self.conv3 = nn.SequentialCell([
nn.Conv2d(width, out_channels * self.expansion, kernel_size=1),
norm(out_channels * self.expansion)
])
self.relu = nn.ReLU()
self.down_sample = down_sample
def construct(self, x: Tensor) -> Tensor:
identity = x
out = self.conv1(x)
out = self.conv2(out)
out = self.conv3(out)
if self.down_sample:
identity = self.down_sample(x)
out += identity
out = self.relu(out)
return out
[文档]class SKNet(ResNet):
r"""SKNet model class, based on
`"Selective Kernel Networks" <https://arxiv.org/abs/1903.06586>`_
Args:
block: block of sknet.
layers: number of layers of each stage.
num_classes: number of classification classes. Default: 1000.
in_channels: number the channels of the input. Default: 3.
groups: number of groups for group conv in blocks. Default: 1.
base_width: base width of pre group hidden channel in blocks. Default: 64.
norm: normalization layer in blocks. Default: None.
sk_kwargs: kwargs of selective kernel. Default: None.
"""
def __init__(self,
block: Type[nn.Cell],
layers: List[int],
num_classes: int = 1000,
in_channels: int = 3,
groups: int = 1,
base_width: int = 64,
norm: Optional[nn.Cell] = None,
sk_kwargs: Optional[Dict] = None
) -> None:
self.sk_kwargs: Optional[Dict] = sk_kwargs # make pylint happy
super().__init__(block, layers, num_classes, in_channels, groups, base_width, norm)
def _make_layer(self,
block: Type[Union[SelectiveKernelBasic, SelectiveKernelBottleneck]],
channels: int,
block_nums: int,
stride: int = 1
) -> nn.SequentialCell:
down_sample = None
if stride != 1 or self.input_channels != channels * block.expansion:
down_sample = nn.SequentialCell([
nn.Conv2d(self.input_channels, channels * block.expansion, kernel_size=1, stride=stride),
self.norm(channels * block.expansion)
])
layers = []
layers.append(
block(
self.input_channels,
channels,
stride=stride,
down_sample=down_sample,
groups=self.groups,
base_width=self.base_with,
norm=self.norm,
sk_kwargs=self.sk_kwargs
)
)
self.input_channels = channels * block.expansion
for _ in range(1, block_nums):
layers.append(
block(
self.input_channels,
channels,
groups=self.groups,
base_width=self.base_with,
norm=self.norm,
sk_kwargs=self.sk_kwargs
)
)
return nn.SequentialCell(layers)
@register_model
def sk_resnet18(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs) -> ResNet:
"""Get 18 layers SKNet model.
Refer to the base class `models.SKNet` for more details.
"""
default_cfg = default_cfgs['sk_resnet18']
sk_kwargs = dict(rd_ratio=1 / 8, rd_divisor=16, split_input=True)
model = SKNet(SelectiveKernelBasic, [2, 2, 2, 2], num_classes=num_classes, in_channels=in_channels,
sk_kwargs=sk_kwargs, **kwargs)
if pretrained:
load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
return model
@register_model
def sk_resnet34(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs) -> ResNet:
"""Get 34 layers SKNet model.
Refer to the base class `models.SKNet` for more details.
"""
default_cfg = default_cfgs['sk_resnet34']
sk_kwargs = dict(rd_ratio=1 / 8, rd_divisor=16, split_input=True)
model = SKNet(SelectiveKernelBasic, [3, 4, 6, 3], num_classes=num_classes, in_channels=in_channels,
sk_kwargs=sk_kwargs, **kwargs)
if pretrained:
load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
return model
@register_model
def sk_resnet50(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs) -> ResNet:
"""Get 50 layers SKNet model.
Refer to the base class `models.SKNet` for more details.
"""
default_cfg = default_cfgs['sk_resnet50']
sk_kwargs = dict(split_input=True)
model = SKNet(SelectiveKernelBottleneck, [3, 4, 6, 3], num_classes=num_classes, in_channels=in_channels,
sk_kwargs=sk_kwargs, **kwargs)
if pretrained:
load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
return model
@register_model
def sk_resnext50_32x4d(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs) -> ResNet:
"""Get 50 layers SKNeXt model with 32 groups of GPConv.
Refer to the base class `models.SKNet` for more details.
"""
default_cfg = default_cfgs['sk_resnext50_32X4d']
sk_kwargs = dict(rd_ratio=1 / 16, rd_divisor=32, split_input=False)
model = SKNet(SelectiveKernelBottleneck, [3, 4, 6, 3], num_classes=num_classes, in_channels=in_channels,
sk_kwargs=sk_kwargs, **kwargs)
if pretrained:
load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
return model