|
- """
- MindSpore implementation of `ResNet`.
- Refer to Deep Residual Learning for Image Recognition.
- """
-
- from typing import Optional, Type, List, Union
-
- from mindspore import nn, Tensor
-
- from .layers.pooling import GlobalAvgPooling
- from .utils import load_pretrained
- from .registry import register_model
-
- __all__ = [
- 'ResNet',
- 'resnet18',
- 'resnet34',
- 'resnet50',
- 'resnet101',
- 'resnet152',
- 'resnext50_32x4d',
- 'resnext101_32x4d',
- 'resnext101_64x4d',
- 'resnext152_64x4d'
- ]
-
-
- def _cfg(url='', **kwargs):
- return {
- 'url': url,
- 'num_classes': 1000,
- 'first_conv': 'conv1', 'classifier': 'classifier',
- **kwargs
- }
-
-
- default_cfgs = {
- 'resnet18': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnet/resnet18_224.ckpt'),
- 'resnet34': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnet/resnet34_224.ckpt'),
- 'resnet50': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnet/resnet50_224.ckpt'),
- 'resnet101': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnet/resnet101_224.ckpt'),
- 'resnet152': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnet/resnet152_224.ckpt'),
- 'resnext50_32x4d': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnext/resnext50_32x4d_224.ckpt'),
- 'resnext101_32x4d': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnext/resnext101_32x4d_224.ckpt'),
- 'resnext101_64x4d': _cfg(url=''),
- 'resnext152_64x4d': _cfg(url='https://download.mindspore.cn/toolkits/mindcv/resnext/resnext152_64x4d_224.ckpt')
- }
-
-
- class BasicBlock(nn.Cell):
- """define the basic block of resnet"""
- expansion: int = 1
-
- def __init__(self,
- in_channels: int,
- channels: int,
- stride: int = 1,
- groups: int = 1,
- base_width: int = 64,
- norm: Optional[nn.Cell] = None,
- down_sample: Optional[nn.Cell] = None
- ) -> None:
- super().__init__()
- if norm is None:
- norm = nn.BatchNorm2d
- assert groups == 1, 'BasicBlock only supports groups=1'
- assert base_width == 64, 'BasicBlock only supports base_width=64'
-
- self.conv1 = nn.Conv2d(in_channels, channels, kernel_size=3,
- stride=stride, padding=1, pad_mode='pad')
- self.bn1 = norm(channels)
- self.relu = nn.ReLU()
- self.conv2 = nn.Conv2d(channels, channels, kernel_size=3,
- stride=1, padding=1, pad_mode='pad')
- self.bn2 = norm(channels)
- self.down_sample = down_sample
-
- def construct(self, x: Tensor) -> Tensor:
- identity = x
-
- out = self.conv1(x)
- out = self.bn1(out)
- out = self.relu(out)
-
- out = self.conv2(out)
- out = self.bn2(out)
-
- if self.down_sample is not None:
- identity = self.down_sample(x)
-
- out += identity
- out = self.relu(out)
-
- return out
-
-
- class Bottleneck(nn.Cell):
- """
- Bottleneck here places the stride for downsampling at 3x3 convolution(self.conv2) as torchvision does,
- while original implementation places the stride at the first 1x1 convolution(self.conv1)
- """
- expansion: int = 4
-
- def __init__(self,
- in_channels: int,
- channels: int,
- stride: int = 1,
- groups: int = 1,
- base_width: int = 64,
- norm: Optional[nn.Cell] = None,
- down_sample: Optional[nn.Cell] = None
- ) -> None:
- super().__init__()
- if norm is None:
- norm = nn.BatchNorm2d
-
- width = int(channels * (base_width / 64.0)) * groups
-
- self.conv1 = nn.Conv2d(in_channels, width, kernel_size=1, stride=1)
- self.bn1 = norm(width)
- self.conv2 = nn.Conv2d(width, width, kernel_size=3, stride=stride,
- padding=1, pad_mode='pad', group=groups)
- self.bn2 = norm(width)
- self.conv3 = nn.Conv2d(width, channels * self.expansion,
- kernel_size=1, stride=1)
- self.bn3 = norm(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.bn1(out)
- out = self.relu(out)
-
- out = self.conv2(out)
- out = self.bn2(out)
- out = self.relu(out)
-
- out = self.conv3(out)
- out = self.bn3(out)
-
- if self.down_sample is not None:
- identity = self.down_sample(x)
-
- out += identity
- out = self.relu(out)
-
- return out
-
-
- class ResNet(nn.Cell):
- r"""ResNet model class, based on
- `"Deep Residual Learning for Image Recognition" <https://arxiv.org/abs/1512.03385>`_
-
- Args:
- block: block of resnet.
- 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.
- """
-
- def __init__(self,
- block: Type[Union[BasicBlock, Bottleneck]],
- layers: List[int],
- num_classes: int = 1000,
- in_channels: int = 3,
- groups: int = 1,
- base_width: int = 64,
- norm: Optional[nn.Cell] = None
- ) -> None:
- super().__init__()
- if norm is None:
- norm = nn.BatchNorm2d
-
- self.norm: nn.Cell = norm # add type hints to make pylint happy
- self.input_channels = 64
- self.groups = groups
- self.base_with = base_width
-
- self.conv1 = nn.Conv2d(in_channels, self.input_channels, kernel_size=7,
- stride=2, pad_mode='pad', padding=3)
- self.bn1 = norm(self.input_channels)
- self.relu = nn.ReLU()
- self.max_pool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='same')
- self.layer1 = self._make_layer(block, 64, layers[0])
- self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
- self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
- self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-
- self.pool = GlobalAvgPooling()
- self.num_features = 512 * block.expansion
- self.classifier = nn.Dense(self.num_features, num_classes)
-
- def _make_layer(self,
- block: Type[Union[BasicBlock, Bottleneck]],
- channels: int,
- block_nums: int,
- stride: int = 1
- ) -> nn.SequentialCell:
- """build model depending on cfgs"""
- 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
- )
- )
- 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
- )
- )
-
- return nn.SequentialCell(layers)
-
- def forward_features(self, x: Tensor) -> Tensor:
- """Network forward feature extraction."""
- x = self.conv1(x)
- x = self.bn1(x)
- x = self.relu(x)
- x = self.max_pool(x)
-
- x = self.layer1(x)
- x = self.layer2(x)
- x = self.layer3(x)
- x = self.layer4(x)
- return x
-
- def forward_head(self, x: Tensor) -> Tensor:
- x = self.pool(x)
- x = self.classifier(x)
- return x
-
- def construct(self, x: Tensor) -> Tensor:
- x = self.forward_features(x)
- x = self.forward_head(x)
- return x
-
-
- @register_model
- def resnet18(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 18 layers ResNet model.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnet18']
- model = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnet34(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 34 layers ResNet model.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnet34']
- model = ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes, in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnet50(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 50 layers ResNet model.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnet50']
- model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=num_classes, in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnet101(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 101 layers ResNet model.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnet101']
- model = ResNet(Bottleneck, [3, 4, 23, 3], num_classes=num_classes, in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnet152(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 152 layers ResNet model.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnet152']
- model = ResNet(Bottleneck, [3, 8, 36, 3], num_classes=num_classes, in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnext50_32x4d(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 50 layers ResNeXt model with 32 groups of GPConv.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnext50_32x4d']
- model = ResNet(Bottleneck, [3, 4, 6, 3], groups=32, base_width=4, num_classes=num_classes,
- in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnext101_32x4d(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 101 layers ResNeXt model with 32 groups of GPConv.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnext101_32x4d']
- model = ResNet(Bottleneck, [3, 4, 23, 3], groups=32, base_width=4, num_classes=num_classes,
- in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnext101_64x4d(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- """Get 101 layers ResNeXt model with 64 groups of GPConv.
- Refer to the base class `models.ResNet` for more details.
- """
- default_cfg = default_cfgs['resnext101_64x4d']
- model = ResNet(Bottleneck, [3, 4, 23, 3], groups=64, base_width=4, num_classes=num_classes,
- in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
-
-
- @register_model
- def resnext152_64x4d(pretrained: bool = False, num_classes: int = 1000, in_channels=3, **kwargs):
- default_cfg = default_cfgs['resnext101_64x4d']
- model = ResNet(Bottleneck, [3, 8, 36, 3], groups=64, base_width=4, num_classes=num_classes,
- in_channels=in_channels, **kwargs)
-
- if pretrained:
- load_pretrained(model, default_cfg, num_classes=num_classes, in_channels=in_channels)
-
- return model
|