|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from .build import BACKBONE_REGISTRY |
|
|
from .backbone import Backbone |
|
|
|
|
|
try: |
|
|
from torch.hub import load_state_dict_from_url |
|
|
except ImportError: |
|
|
from torch.utils.model_zoo import load_url as load_state_dict_from_url |
|
|
|
|
|
model_urls = { |
|
|
"vgg11": "https://download.pytorch.org/models/vgg11-bbd30ac9.pth", |
|
|
"vgg13": "https://download.pytorch.org/models/vgg13-c768596a.pth", |
|
|
"vgg16": "https://download.pytorch.org/models/vgg16-397923af.pth", |
|
|
"vgg19": "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth", |
|
|
"vgg11_bn": "https://download.pytorch.org/models/vgg11_bn-6002323d.pth", |
|
|
"vgg13_bn": "https://download.pytorch.org/models/vgg13_bn-abd245e5.pth", |
|
|
"vgg16_bn": "https://download.pytorch.org/models/vgg16_bn-6c64b313.pth", |
|
|
"vgg19_bn": "https://download.pytorch.org/models/vgg19_bn-c79401a0.pth", |
|
|
} |
|
|
|
|
|
|
|
|
class VGG(Backbone): |
|
|
|
|
|
def __init__(self, features, init_weights=True): |
|
|
super().__init__() |
|
|
self.features = features |
|
|
self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) |
|
|
|
|
|
self.classifier = nn.Sequential( |
|
|
nn.Linear(512 * 7 * 7, 4096), |
|
|
nn.ReLU(True), |
|
|
nn.Dropout(), |
|
|
nn.Linear(4096, 4096), |
|
|
nn.ReLU(True), |
|
|
nn.Dropout(), |
|
|
) |
|
|
|
|
|
self._out_features = 4096 |
|
|
|
|
|
if init_weights: |
|
|
self._initialize_weights() |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.features(x) |
|
|
x = self.avgpool(x) |
|
|
x = torch.flatten(x, 1) |
|
|
return self.classifier(x) |
|
|
|
|
|
def _initialize_weights(self): |
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
nn.init.kaiming_normal_( |
|
|
m.weight, mode="fan_out", nonlinearity="relu" |
|
|
) |
|
|
if m.bias is not None: |
|
|
nn.init.constant_(m.bias, 0) |
|
|
elif isinstance(m, nn.BatchNorm2d): |
|
|
nn.init.constant_(m.weight, 1) |
|
|
nn.init.constant_(m.bias, 0) |
|
|
elif isinstance(m, nn.Linear): |
|
|
nn.init.normal_(m.weight, 0, 0.01) |
|
|
nn.init.constant_(m.bias, 0) |
|
|
|
|
|
|
|
|
def make_layers(cfg, batch_norm=False): |
|
|
layers = [] |
|
|
in_channels = 3 |
|
|
for v in cfg: |
|
|
if v == "M": |
|
|
layers += [nn.MaxPool2d(kernel_size=2, stride=2)] |
|
|
else: |
|
|
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) |
|
|
if batch_norm: |
|
|
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] |
|
|
else: |
|
|
layers += [conv2d, nn.ReLU(inplace=True)] |
|
|
in_channels = v |
|
|
return nn.Sequential(*layers) |
|
|
|
|
|
|
|
|
cfgs = { |
|
|
"A": [64, "M", 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], |
|
|
"B": |
|
|
[64, 64, "M", 128, 128, "M", 256, 256, "M", 512, 512, "M", 512, 512, "M"], |
|
|
"D": [ |
|
|
64, |
|
|
64, |
|
|
"M", |
|
|
128, |
|
|
128, |
|
|
"M", |
|
|
256, |
|
|
256, |
|
|
256, |
|
|
"M", |
|
|
512, |
|
|
512, |
|
|
512, |
|
|
"M", |
|
|
512, |
|
|
512, |
|
|
512, |
|
|
"M", |
|
|
], |
|
|
"E": [ |
|
|
64, |
|
|
64, |
|
|
"M", |
|
|
128, |
|
|
128, |
|
|
"M", |
|
|
256, |
|
|
256, |
|
|
256, |
|
|
256, |
|
|
"M", |
|
|
512, |
|
|
512, |
|
|
512, |
|
|
512, |
|
|
"M", |
|
|
512, |
|
|
512, |
|
|
512, |
|
|
512, |
|
|
"M", |
|
|
], |
|
|
} |
|
|
|
|
|
|
|
|
def _vgg(arch, cfg, batch_norm, pretrained): |
|
|
init_weights = False if pretrained else True |
|
|
model = VGG( |
|
|
make_layers(cfgs[cfg], batch_norm=batch_norm), |
|
|
init_weights=init_weights |
|
|
) |
|
|
if pretrained: |
|
|
state_dict = load_state_dict_from_url(model_urls[arch], progress=True) |
|
|
model.load_state_dict(state_dict, strict=False) |
|
|
return model |
|
|
|
|
|
|
|
|
@BACKBONE_REGISTRY.register() |
|
|
def vgg16(pretrained=True, **kwargs): |
|
|
return _vgg("vgg16", "D", False, pretrained) |
|
|
|
