Created
June 16, 2020 03:13
-
-
Save gasvn/0c359e2d6a79760196c828f763879273 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import torch.nn as nn | |
import math | |
import torch.utils.model_zoo as model_zoo | |
import torch | |
import torch.nn.functional as F | |
__all__ = ['Res2Net', 'res2net50_v1b', 'res2net101_v1b'] | |
model_urls = { | |
'res2net50_v1b_26w_4s': 'https://shanghuagao.oss-cn-beijing.aliyuncs.com/res2net/res2net50_v1b_26w_4s-3cf99910.pth', | |
'res2net101_v1b_26w_4s': 'https://shanghuagao.oss-cn-beijing.aliyuncs.com/res2net/res2net101_v1b_26w_4s-0812c246.pth', | |
} | |
class Bottle2neck(nn.Module): | |
expansion = 4 | |
def __init__(self, inplanes, planes, stride=1, downsample=None, baseWidth=26, scale = 4, stype='normal'): | |
""" Constructor | |
Args: | |
inplanes: input channel dimensionality | |
planes: output channel dimensionality | |
stride: conv stride. Replaces pooling layer. | |
downsample: None when stride = 1 | |
baseWidth: basic width of conv3x3 | |
scale: number of scale. | |
type: 'normal': normal set. 'stage': first block of a new stage. | |
""" | |
super(Bottle2neck, self).__init__() | |
width = int(math.floor(planes * (baseWidth/64.0))) | |
self.conv1 = nn.Conv2d(inplanes, width*scale, kernel_size=1, bias=False) | |
self.bn1 = nn.BatchNorm2d(width*scale) | |
assert scale>1, 'Res2Net degenerate to ResNet when scale=1!' | |
if stype == 'stage': | |
self.pool = nn.AvgPool2d(kernel_size=3, stride = stride, padding=1) | |
convs = [] | |
bns = [] | |
for i in range(scale-1): | |
convs.append(nn.Conv2d(width, width, kernel_size=3, stride = stride, padding=1, bias=False)) | |
bns.append(nn.BatchNorm2d(width)) | |
self.convs = nn.ModuleList(convs) | |
self.bns = nn.ModuleList(bns) | |
self.conv3 = nn.Conv2d(width*scale, planes * self.expansion, kernel_size=1, bias=False) | |
self.bn3 = nn.BatchNorm2d(planes * self.expansion) | |
self.relu = nn.ReLU(inplace=True) | |
self.downsample = downsample | |
self.stype = stype | |
self.scale = scale | |
self.width = width | |
self.stride = stride | |
def forward(self, x): | |
residual = x | |
out = self.conv1(x) | |
out = self.bn1(out) | |
out = self.relu(out) | |
spx = torch.split(out, self.width, 1) | |
sp = self.convs[0](spx[0]) | |
sp = self.relu(self.bns[0](sp)) | |
out = [sp] | |
for i in range(1, self.scale-1): | |
if self.stype=='stage': | |
sp = spx[i] | |
else: | |
sp = sp + spx[i] | |
sp = self.convs[i](sp) | |
sp = self.relu(self.bns[i](sp)) | |
out.append(sp) | |
if self.stype=='normal': | |
out.append(spx[self.scale-1]) | |
elif self.stype=='stage': | |
out.append(self.pool(spx[self.scale-1])) | |
out = torch.cat(out,1) | |
out = self.conv3(out) | |
out = self.bn3(out) | |
if self.downsample is not None: | |
residual = self.downsample(x) | |
out += residual | |
out = self.relu(out) | |
return out | |
class Res2Net(nn.Module): | |
def __init__(self, block, layers, baseWidth = 26, scale = 4, num_classes=1000): | |
self.inplanes = 64 | |
super(Res2Net, self).__init__() | |
self.baseWidth = baseWidth | |
self.scale = scale | |
self.conv1 = nn.Sequential( | |
nn.Conv2d(3, 32, 3, 2, 1, bias=False), | |
nn.BatchNorm2d(32), | |
nn.ReLU(inplace=True), | |
nn.Conv2d(32, 32, 3, 1, 1, bias=False), | |
nn.BatchNorm2d(32), | |
nn.ReLU(inplace=True), | |
nn.Conv2d(32, 64, 3, 1, 1, bias=False) | |
) | |
self.bn1 = nn.BatchNorm2d(64) | |
self.relu = nn.ReLU() | |
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) | |
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.avgpool = nn.AdaptiveAvgPool2d(1) | |
self.fc = nn.Linear(512 * block.expansion, num_classes) | |
for m in self.modules(): | |
if isinstance(m, nn.Conv2d): | |
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') | |
elif isinstance(m, nn.BatchNorm2d): | |
nn.init.constant_(m.weight, 1) | |
nn.init.constant_(m.bias, 0) | |
def _make_layer(self, block, planes, blocks, stride=1): | |
downsample = None | |
if stride != 1 or self.inplanes != planes * block.expansion: | |
downsample = nn.Sequential( | |
nn.AvgPool2d(kernel_size=stride, stride=stride, | |
ceil_mode=True, count_include_pad=False), | |
nn.Conv2d(self.inplanes, planes * block.expansion, | |
kernel_size=1, stride=1, bias=False), | |
nn.BatchNorm2d(planes * block.expansion), | |
) | |
layers = [] | |
layers.append(block(self.inplanes, planes, stride, downsample=downsample, | |
stype='stage', baseWidth = self.baseWidth, scale=self.scale)) | |
self.inplanes = planes * block.expansion | |
for i in range(1, blocks): | |
layers.append(block(self.inplanes, planes, baseWidth = self.baseWidth, scale=self.scale)) | |
return nn.Sequential(*layers) | |
def forward(self, x): | |
x = self.conv1(x) | |
x = self.bn1(x) | |
x = self.relu(x) | |
x = self.maxpool(x) | |
x = self.layer1(x) | |
x = self.layer2(x) | |
x = self.layer3(x) | |
x = self.layer4(x) | |
x = self.avgpool(x) | |
x = x.view(x.size(0), -1) | |
x = self.fc(x) | |
return x | |
def res2net50_v1b(pretrained=False, **kwargs): | |
"""Constructs a Res2Net-50_v1b model. | |
Res2Net-50 refers to the Res2Net-50_v1b_26w_4s. | |
Args: | |
pretrained (bool): If True, returns a model pre-trained on ImageNet | |
""" | |
model = Res2Net(Bottle2neck, [3, 4, 6, 3], baseWidth = 26, scale = 4, **kwargs) | |
if pretrained: | |
model.load_state_dict(model_zoo.load_url(model_urls['res2net50_v1b_26w_4s'])) | |
return model | |
def res2net101_v1b(pretrained=False, **kwargs): | |
"""Constructs a Res2Net-50_v1b_26w_4s model. | |
Args: | |
pretrained (bool): If True, returns a model pre-trained on ImageNet | |
""" | |
model = Res2Net(Bottle2neck, [3, 4, 23, 3], baseWidth = 26, scale = 4, **kwargs) | |
if pretrained: | |
model.load_state_dict(model_zoo.load_url(model_urls['res2net101_v1b_26w_4s'])) | |
return model | |
def res2net50_v1b_26w_4s(pretrained=False, **kwargs): | |
"""Constructs a Res2Net-50_v1b_26w_4s model. | |
Args: | |
pretrained (bool): If True, returns a model pre-trained on ImageNet | |
""" | |
model = Res2Net(Bottle2neck, [3, 4, 6, 3], baseWidth = 26, scale = 4, **kwargs) | |
if pretrained: | |
model.load_state_dict(model_zoo.load_url(model_urls['res2net50_v1b_26w_4s'])) | |
return model | |
def res2net101_v1b_26w_4s(pretrained=False, **kwargs): | |
"""Constructs a Res2Net-50_v1b_26w_4s model. | |
Args: | |
pretrained (bool): If True, returns a model pre-trained on ImageNet | |
""" | |
model = Res2Net(Bottle2neck, [3, 4, 23, 3], baseWidth = 26, scale = 4, **kwargs) | |
if pretrained: | |
model.load_state_dict(model_zoo.load_url(model_urls['res2net101_v1b_26w_4s'])) | |
return model | |
def res2net152_v1b_26w_4s(pretrained=False, **kwargs): | |
"""Constructs a Res2Net-50_v1b_26w_4s model. | |
Args: | |
pretrained (bool): If True, returns a model pre-trained on ImageNet | |
""" | |
model = Res2Net(Bottle2neck, [3, 8, 36, 3], baseWidth = 26, scale = 4, **kwargs) | |
if pretrained: | |
model.load_state_dict(model_zoo.load_url(model_urls['res2net152_v1b_26w_4s'])) | |
return model | |
if __name__ == '__main__': | |
images = torch.rand(1, 3, 224, 224).cuda(0) | |
model = res2net50_v1b_26w_4s(pretrained=True) | |
model = model.cuda(0) | |
print(model(images).size()) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment