mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
70 lines
2.2 KiB
70 lines
2.2 KiB
3 years ago
|
#!/usr/bin/env python
|
||
|
|
# -*- encoding: utf-8 -*-
|
||
|
|
|
||
|
|
from typing import Optional, Callable
|
||
|
|
|
||
|
|
import torch.nn as nn
|
||
|
|
from torch import Tensor
|
||
|
|
|
||
|
|
from colossalai.registry import LAYERS
|
||
|
|
from .conv import conv3x3, conv1x1
|
||
|
|
|
||
|
|
|
||
|
|
@LAYERS.register_module
|
||
|
|
class ResNetBottleneck(nn.Module):
|
||
|
|
# Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
|
||
|
|
# while original implementation places the stride at the first 1x1 convolution(self.conv1)
|
||
|
|
# according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
|
||
|
|
# This variant is also known as ResNet V1.5 and improves accuracy according to
|
||
|
|
# https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
|
||
|
|
|
||
|
|
expansion: int = 4
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
inplanes: int,
|
||
|
|
planes: int,
|
||
|
|
stride: int = 1,
|
||
|
|
downsample: Optional[nn.Module] = None,
|
||
|
|
groups: int = 1,
|
||
|
|
base_width: int = 64,
|
||
|
|
dilation: int = 1,
|
||
|
|
norm_layer: Optional[Callable[..., nn.Module]] = None
|
||
|
|
) -> None:
|
||
|
|
super().__init__()
|
||
|
|
if norm_layer is None:
|
||
|
|
norm_layer = nn.BatchNorm2d
|
||
|
|
width = int(planes * (base_width / 64.)) * groups
|
||
|
|
# Both self.conv2 and self.downsample layers downsample the input when stride != 1
|
||
|
|
self.conv1 = conv1x1(inplanes, width)
|
||
|
|
self.bn1 = norm_layer(width)
|
||
|
|
self.conv2 = conv3x3(width, width, stride, groups, dilation)
|
||
|
|
self.bn2 = norm_layer(width)
|
||
|
|
self.conv3 = conv1x1(width, planes * self.expansion)
|
||
|
|
self.bn3 = norm_layer(planes * self.expansion)
|
||
|
|
self.relu = nn.ReLU(inplace=True)
|
||
|
|
self.downsample = downsample
|
||
|
|
self.stride = stride
|
||
|
|
|
||
|
|
def forward(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.downsample is not None:
|
||
|
|
identity = self.downsample(x)
|
||
|
|
|
||
|
|
out += identity
|
||
|
|
out = self.relu(out)
|
||
|
|
|
||
|
|
return out
|