File size: 7,943 Bytes
92f0e98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 |
'''
Customized version of pytorch resnet, alexnets.
'''
import numpy, torch, math, os
from torch import nn
from collections import OrderedDict
from torchvision.models import resnet
from torchvision.models.alexnet import model_urls as alexnet_model_urls
class CustomResNet(nn.Module):
'''
Customizable ResNet, compatible with pytorch's resnet, but:
* The top-level sequence of modules can be modified to add
or remove or alter layers.
* Extra outputs can be produced, to allow backprop and access
to internal features.
* Pooling is replaced by resizable GlobalAveragePooling so that
any size can be input (e.g., any multiple of 32 pixels).
* halfsize=True halves striding on the first pooling to
set the default size to 112x112 instead of 224x224.
'''
def __init__(self, size=None, block=None, layers=None, num_classes=1000,
extra_output=None, modify_sequence=None, halfsize=False):
standard_sizes = {
18: (resnet.BasicBlock, [2, 2, 2, 2]),
34: (resnet.BasicBlock, [3, 4, 6, 3]),
50: (resnet.Bottleneck, [3, 4, 6, 3]),
101: (resnet.Bottleneck, [3, 4, 23, 3]),
152: (resnet.Bottleneck, [3, 8, 36, 3])
}
assert (size in standard_sizes) == (block is None) == (layers is None)
if size in standard_sizes:
block, layers = standard_sizes[size]
if modify_sequence is None:
modify_sequence = lambda x: x
self.inplanes = 64
norm_layer = nn.BatchNorm2d
self._norm_layer = norm_layer # for recent resnet
self.dilation = 1
self.groups = 1
self.base_width = 64
sequence = modify_sequence([
('conv1', nn.Conv2d(3, 64, kernel_size=7, stride=2,
padding=3, bias=False)),
('bn1', norm_layer(64)),
('relu', nn.ReLU(inplace=True)),
('maxpool', nn.MaxPool2d(3, stride=1 if halfsize else 2,
padding=1)),
('layer1', self._make_layer(block, 64, layers[0])),
('layer2', self._make_layer(block, 128, layers[1], stride=2)),
('layer3', self._make_layer(block, 256, layers[2], stride=2)),
('layer4', self._make_layer(block, 512, layers[3], stride=2)),
('avgpool', GlobalAveragePool2d()),
('fc', nn.Linear(512 * block.expansion, num_classes))
])
super(CustomResNet, self).__init__()
for name, layer in sequence:
setattr(self, name, layer)
self.extra_output = extra_output
def _make_layer(self, block, channels, depth, stride=1):
return resnet.ResNet._make_layer(self, block, channels, depth, stride)
def forward(self, x):
extra = []
for name, module in self._modules.items():
x = module(x)
if self.extra_output and name in self.extra_output:
extra.append(x)
if self.extra_output:
return (x,) + tuple(extra)
return x
class CustomAlexNet(nn.Module):
'''
Customizable AlexNet, compatible with pytorch's alexnet, but:
* The top-level sequence of modules can be modified to add
or remove or alter layers.
* Extra outputs can be produced, to allow backprop and access
to internal features.
* halfsize=True halves striding on the first convolution to
allow 119x119 images to be processed rather than 227x227 only.
'''
def __init__(self, channels=None, num_classes=1000,
extra_output=None, modify_sequence=None, halfsize=False):
if channels is None:
channels = [3, 64, 192, 384, 256, 256, 4096, 4096]
if modify_sequence is None:
modify_sequence = lambda x: x
sequence = modify_sequence([
('conv1', nn.Conv2d(channels[0], channels[1],
kernel_size=11, stride=4, padding=2)),
('relu1', nn.ReLU(inplace=True)),
('pool1', nn.MaxPool2d(kernel_size=3, stride=1 if halfsize else 2)),
('conv2', nn.Conv2d(channels[1], channels[2],
kernel_size=5, padding=2)),
('relu2', nn.ReLU(inplace=True)),
('pool2', nn.MaxPool2d(kernel_size=3, stride=2)),
('conv3', nn.Conv2d(channels[2], channels[3],
kernel_size=3, padding=1)),
('relu3', nn.ReLU(inplace=True)),
('conv4', nn.Conv2d(channels[3], channels[4],
kernel_size=3, padding=1)),
('relu4', nn.ReLU(inplace=True)),
('conv5', nn.Conv2d(channels[4], channels[5],
kernel_size=3, padding=1)),
('relu5', nn.ReLU(inplace=True)),
('pool5', nn.MaxPool2d(kernel_size=3, stride=2)),
('flatten', Vectorize()),
('dropout6', nn.Dropout()),
('fc6', nn.Linear(channels[5] * 6 * 6, channels[6])),
('relu6', nn.ReLU(inplace=True)),
('dropout7', nn.Dropout()),
('fc7', nn.Linear(channels[6], channels[7])),
('relu7', nn.ReLU(inplace=True)),
('fc8', nn.Linear(channels[7], num_classes))
])
super(CustomAlexNet, self).__init__()
for name, layer in sequence:
setattr(self, name, layer)
self.extra_output = extra_output
def forward(self, x):
extra = []
for name, module in self._modules.items():
x = module(x)
if self.extra_output and name in self.extra_output:
extra.append(x)
if self.extra_output:
return (x,) + tuple(extra)
return x
def load_state_dict(self, state_dict, **kwargs):
'''
Translates from pytorch's AlexNet parameter names
into the custom parameter names.
'''
custom_names = [
('features.0.', 'conv1.'),
('features.3.', 'conv2.'),
('features.6.', 'conv3.'),
('features.8.', 'conv4.'),
('features.10.', 'conv5.'),
('classifier.1.', 'fc6.'),
('classifier.4.', 'fc7.'),
('classifier.6.', 'fc8.')
]
custom_state_dict = {}
for k, v in state_dict.items():
for op, np in custom_names:
if k.startswith(op):
k = np + k[len(op):]
break
custom_state_dict[k] = v
super(CustomAlexNet, self).load_state_dict(custom_state_dict, **kwargs)
class Vectorize(nn.Module):
def __init__(self):
super(Vectorize, self).__init__()
def forward(self, x):
x = x.view(x.size(0), int(numpy.prod(x.size()[1:])))
return x
class GlobalAveragePool2d(nn.Module):
def __init__(self):
super(GlobalAveragePool2d, self).__init__()
def forward(self, x):
x = torch.mean(x.view(x.size(0), x.size(1), -1), dim=2)
return x
if __name__ == '__main__':
import torch.utils.model_zoo as model_zoo
# Verify that at the default settings, pytorch standard pretrained
# models can be loaded into each of the custom nets.
print('Loading alexnet')
model = CustomAlexNet()
model.load_state_dict(model_zoo.load_url(alexnet_model_urls['alexnet']))
print('Loading resnet18')
model = CustomResNet(18)
model.load_state_dict(model_zoo.load_url(resnet.model_urls['resnet18']))
print('Loading resnet34')
model = CustomResNet(34)
model.load_state_dict(model_zoo.load_url(resnet.model_urls['resnet34']))
print('Loading resnet50')
model = CustomResNet(50)
model.load_state_dict(model_zoo.load_url(resnet.model_urls['resnet50']))
print('Loading resnet101')
model = CustomResNet(101)
model.load_state_dict(model_zoo.load_url(resnet.model_urls['resnet101']))
print('Loading resnet152')
model = CustomResNet(152)
model.load_state_dict(model_zoo.load_url(resnet.model_urls['resnet152']))
|