-
Notifications
You must be signed in to change notification settings - Fork 113
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
rm .tar in filename, add se_resnet_ibn_b
- Loading branch information
1 parent
57dc45e
commit ff1d2bc
Showing
4 changed files
with
194 additions
and
3 deletions.
There are no files selected for viewing
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
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
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,190 @@ | ||
from .se_module import SELayer | ||
import torch.nn as nn | ||
import torch | ||
import math | ||
|
||
__all__ = ['se_resnet50_ibn_b', 'se_resnet101_ibn_b', 'se_resnet152_ibn_b'] | ||
|
||
def conv3x3(in_planes, out_planes, stride=1): | ||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) | ||
|
||
|
||
class SEBasicBlock(nn.Module): | ||
expansion = 1 | ||
|
||
def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16, IN=False): | ||
super(SEBasicBlock, self).__init__() | ||
self.conv1 = conv3x3(inplanes, planes, stride) | ||
self.bn1 = nn.BatchNorm2d(planes) | ||
self.relu = nn.ReLU(inplace=True) | ||
self.conv2 = conv3x3(planes, planes, 1) | ||
self.bn2 = nn.BatchNorm2d(planes) | ||
self.IN = None | ||
if IN: | ||
self.IN = nn.InstanceNorm2d(planes * 4, affine=True) | ||
self.se = SELayer(planes, reduction) | ||
self.downsample = downsample | ||
self.stride = stride | ||
|
||
def forward(self, x): | ||
residual = x | ||
out = self.conv1(x) | ||
out = self.bn1(out) | ||
out = self.relu(out) | ||
|
||
out = self.conv2(out) | ||
out = self.bn2(out) | ||
out = self.se(out) | ||
|
||
if self.downsample is not None: | ||
residual = self.downsample(x) | ||
|
||
out += residual | ||
if self.IN is not None: | ||
out = self.IN(out) | ||
out = self.relu(out) | ||
|
||
return out | ||
|
||
|
||
class SEBottleneck(nn.Module): | ||
expansion = 4 | ||
|
||
def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=16, IN=False): | ||
super(SEBottleneck, self).__init__() | ||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) | ||
self.bn1 = nn.BatchNorm2d(planes) | ||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, | ||
padding=1, bias=False) | ||
self.bn2 = nn.BatchNorm2d(planes) | ||
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) | ||
self.bn3 = nn.BatchNorm2d(planes * 4) | ||
self.IN = None | ||
if IN: | ||
self.IN = nn.InstanceNorm2d(planes * 4, affine=True) | ||
self.relu = nn.ReLU(inplace=True) | ||
self.se = SELayer(planes * 4, reduction) | ||
self.downsample = downsample | ||
self.stride = stride | ||
|
||
def forward(self, x): | ||
residual = 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) | ||
out = self.se(out) | ||
|
||
if self.downsample is not None: | ||
residual = self.downsample(x) | ||
|
||
out += residual | ||
if self.IN is not None: | ||
out = self.IN(out) | ||
out = self.relu(out) | ||
|
||
return out | ||
|
||
class ResNet(nn.Module): | ||
|
||
def __init__(self, block, layers, num_classes=1000): | ||
self.inplanes = 64 | ||
super(ResNet, self).__init__() | ||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, | ||
bias=False) | ||
self.bn1 = nn.InstanceNorm2d(64, affine=True) | ||
self.relu = nn.ReLU(inplace=True) | ||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) | ||
self.layer1 = self._make_layer(block, 64, layers[0], stride=1, IN=True) | ||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, IN=True) | ||
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.AvgPool2d(7) | ||
self.fc = nn.Linear(512 * block.expansion, num_classes) | ||
|
||
self.conv1.weight.data.normal_(0, math.sqrt(2. / (7 * 7 * 64))) | ||
for m in self.modules(): | ||
if isinstance(m, nn.Conv2d): | ||
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels | ||
m.weight.data.normal_(0, math.sqrt(2. / n)) | ||
elif isinstance(m, nn.BatchNorm2d): | ||
m.weight.data.fill_(1) | ||
m.bias.data.zero_() | ||
elif isinstance(m, nn.InstanceNorm2d): | ||
m.weight.data.fill_(1) | ||
m.bias.data.zero_() | ||
|
||
def _make_layer(self, block, planes, blocks, stride=1, IN=False): | ||
downsample = None | ||
if stride != 1 or self.inplanes != planes * block.expansion: | ||
downsample = nn.Sequential( | ||
nn.Conv2d(self.inplanes, planes * block.expansion, | ||
kernel_size=1, stride=stride, bias=False), | ||
nn.BatchNorm2d(planes * block.expansion), | ||
) | ||
|
||
layers = [] | ||
layers.append(block(self.inplanes, planes, stride, downsample)) | ||
self.inplanes = planes * block.expansion | ||
for i in range(1, blocks-1): | ||
layers.append(block(self.inplanes, planes)) | ||
layers.append(block(self.inplanes, planes, IN=IN)) | ||
|
||
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 se_resnet50_ibn_b(num_classes=1000): | ||
"""Constructs a ResNet-50 model. | ||
Args: | ||
pretrained (bool): If True, returns a model pre-trained on ImageNet | ||
""" | ||
model = ResNet(SEBottleneck, [3, 4, 6, 3], num_classes=num_classes) | ||
model.avgpool = nn.AdaptiveAvgPool2d(1) | ||
return model | ||
|
||
|
||
def se_resnet101_ibn_b(num_classes=1000): | ||
"""Constructs a ResNet-101 model. | ||
Args: | ||
pretrained (bool): If True, returns a model pre-trained on ImageNet | ||
""" | ||
model = ResNet(SEBottleneck, [3, 4, 23, 3], num_classes=num_classes) | ||
model.avgpool = nn.AdaptiveAvgPool2d(1) | ||
return model | ||
|
||
|
||
def se_resnet152_ibn_b(num_classes): | ||
"""Constructs a ResNet-152 model. | ||
Args: | ||
pretrained (bool): If True, returns a model pre-trained on ImageNet | ||
""" | ||
model = ResNet(SEBottleneck, [3, 8, 36, 3], num_classes=num_classes) | ||
model.avgpool = nn.AdaptiveAvgPool2d(1) | ||
return model |
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