Last active
October 31, 2019 01:12
-
-
Save paulomann/60003d4a7f54c8b92e467e1f0a2d7208 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
class Bottleneck(nn.Module): | |
''' Standard bottleneck block | |
input = inplanes * H * W | |
middle = planes * H/stride * W/stride | |
output = 4*planes * H/stride * W/stride | |
''' | |
expansion = 4 | |
def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None): | |
super(Bottleneck, 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=dilation, dilation=dilation, 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.relu = nn.ReLU(inplace=True) | |
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) | |
if self.downsample is not None: | |
residual = self.downsample(x) | |
out += residual | |
out = self.relu(out) | |
return out | |
class QuantizableBottleneck(Bottleneck): | |
expansion = 4 | |
def __init__(self, *args, **kwargs): | |
super(QuantizableBottleneck, self).__init__(*args, **kwargs) | |
self.skip_add_relu = nn.quantized.FloatFunctional() | |
self.relu1 = nn.ReLU(inplace=False) | |
self.relu2 = nn.ReLU(inplace=False) | |
def forward(self, x): | |
identity = x | |
out = self.conv1(x) | |
out = self.bn1(out) | |
out = self.relu1(out) | |
out = self.conv2(out) | |
out = self.bn2(out) | |
out = self.relu2(out) | |
out = self.conv3(out) | |
out = self.bn3(out) | |
if self.downsample is not None: | |
identity = self.downsample(x) | |
out = self.skip_add_relu.add_relu(out, identity) | |
return out | |
def fuse_model(self): | |
fuse_modules(self, [['conv1', 'bn1', 'relu1'], | |
['conv2', 'bn2', 'relu2'], | |
['conv3', 'bn3']], inplace=True) | |
if self.downsample: | |
torch.quantization.fuse_modules(self.downsample, ['0', '1'], inplace=True) | |
class ResNet(nn.Module): | |
""" A standard ResNet. | |
""" | |
def __init__(self, block, layers, fc_out, model_name, self_similarity_radius=None, self_similarity_version=2): | |
nn.Module.__init__(self) | |
self.model_name = model_name | |
# default values for a network pre-trained on imagenet | |
self.rgb_means = [0.485, 0.456, 0.406] | |
self.rgb_stds = [0.229, 0.224, 0.225] | |
self.input_size = (3, 224, 224) | |
self.inplanes = 64 | |
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, | |
bias=False) | |
self.bn1 = nn.BatchNorm2d(64) | |
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], self_similarity_radius=self_similarity_radius, self_similarity_version=self_similarity_version) | |
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, self_similarity_radius=self_similarity_radius, self_similarity_version=self_similarity_version) | |
self.layer3 = self._make_layer(block, 256, layers[2], stride=2, self_similarity_radius=self_similarity_radius, self_similarity_version=self_similarity_version) | |
self.layer4 = self._make_layer(block, 512, layers[3], stride=2, self_similarity_radius=self_similarity_radius, self_similarity_version=self_similarity_version) | |
reset_weights(self) | |
self.fc = None | |
self.fc_out = fc_out | |
if self.fc_out > 0: | |
self.avgpool = nn.AdaptiveAvgPool2d(1) | |
self.fc = nn.Linear(512 * block.expansion, fc_out) | |
self.fc_name = 'fc' | |
def _make_layer(self, block, planes, blocks, stride=1, self_similarity_radius=None, self_similarity_version=1): | |
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=planes, stride=stride, downsample=downsample)) | |
self.inplanes = planes * block.expansion | |
for i in range(1, blocks): | |
layers.append(block(self.inplanes, planes)) | |
if self_similarity_radius: | |
if self_similarity_version == 1: | |
from . self_sim import SelfSimilarity1 | |
layers.append(SelfSimilarity1(self_similarity_radius, self.inplanes)) | |
else: | |
from . self_sim import SelfSimilarity2 | |
layers.append(SelfSimilarity2(self_similarity_radius, self.inplanes)) | |
return nn.Sequential(*layers) | |
def forward(self, x, out_layer=0): | |
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) | |
if out_layer==-1: | |
return x, self.layer4(x) | |
x = self.layer4(x) | |
if self.fc_out > 0: | |
x = self.avgpool(x) | |
x = x.view(x.size(0), -1) | |
x = self.fc(x) | |
return x | |
def fuse_model(self): | |
r"""Fuse conv/bn/relu modules in resnet models | |
Fuse conv+bn+relu/ Conv+relu/conv+Bn modules to prepare for quantization. | |
Model is modified in place. Note that this operation does not change numerics | |
and the model after modification is in floating point | |
""" | |
fuse_modules(self, ['conv1', 'bn1', 'relu'], inplace=True) | |
for m in self.modules(): | |
if type(m) == QuantizableBottleneck: | |
m.fuse_model() | |
class ResNet_RMAC(ResNet): | |
""" ResNet for RMAC (without ROI pooling) | |
""" | |
def __init__(self, block, layers, model_name, out_dim=2048, norm_features=False, | |
pooling='gem', gemp=3, center_bias=0, | |
dropout_p=None, without_fc=False, **kwargs): | |
ResNet.__init__(self, block, layers, 0, model_name, **kwargs) | |
self.norm_features = norm_features | |
self.without_fc = without_fc | |
self.pooling = pooling | |
self.center_bias = center_bias | |
if pooling == 'max': | |
self.adpool = nn.AdaptiveMaxPool2d(output_size=1) | |
elif pooling == 'avg': | |
self.adpool = nn.AdaptiveAvgPool2d(output_size=1) | |
elif pooling.startswith('gem'): | |
self.adpool = GeneralizedMeanPoolingP(norm=gemp) | |
else: | |
raise ValueError(pooling) | |
self.dropout = nn.Dropout(dropout_p) if dropout_p is not None else None | |
self.fc = nn.Linear(512 * block.expansion, out_dim) | |
self.fc_name = 'fc' | |
self.feat_dim = out_dim | |
self.detach = False | |
def _forward(self, x): | |
bs, _, H, W = x.shape | |
x = ResNet.forward(self, x) | |
if self.dropout is not None: | |
x = self.dropout(x) | |
if self.detach: | |
# stop the back-propagation here, if needed | |
x = Variable(x.detach()) | |
x = self.id(x) # fake transformation | |
if self.center_bias > 0: | |
b = self.center_bias | |
bias = 1 + torch.FloatTensor([[[[0,0,0,0],[0,b,b,0],[0,b,b,0],[0,0,0,0]]]]).to(x.device) | |
bias = torch.nn.functional.interpolate(bias, size=x.shape[-2:], mode='bilinear', align_corners=True) | |
x = x*bias | |
# global pooling | |
x = self.adpool(x) | |
if self.norm_features: | |
x = l2_normalize(x, axis=1) | |
x.squeeze_() | |
if not self.without_fc: | |
x = self.fc(x) | |
x = l2_normalize(x, axis=-1) | |
return x | |
forward = _forward | |
class QuantizableRMAC(ResNet_RMAC): | |
def __init__(self, *args, **kwargs): | |
super(QuantizableRMAC, self).__init__(*args, **kwargs) | |
self.quant = torch.quantization.QuantStub() | |
self.dequant = torch.quantization.DeQuantStub() | |
def forward(self, x): | |
x = self.quant(x) | |
x = self._forward(x) | |
x = self.dequant(x) | |
return x | |
def fuse_model(self): | |
super().fuse_model() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment