| import torch | |
| import torch.nn as nn | |
| import torch.utils.model_zoo as model_zoo | |
| from torch.nn import functional as F | |
| class VGG(nn.Module): | |
| def __init__(self, features): | |
| super(VGG, self).__init__() | |
| self.features = features | |
| self.reg_layer = nn.Sequential( | |
| nn.Conv2d(512, 256, kernel_size=3, padding=1), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(256, 128, kernel_size=3, padding=1), | |
| nn.ReLU(inplace=True), | |
| ) | |
| self.density_layer = nn.Sequential(nn.Conv2d(128, 1, 1), nn.ReLU()) | |
| def forward(self, x): | |
| x = self.features(x) | |
| x = F.upsample_bilinear(x, scale_factor=2) | |
| x = self.reg_layer(x) | |
| mu = self.density_layer(x) | |
| B, C, H, W = mu.size() | |
| mu_sum = mu.view([B, -1]).sum(1).unsqueeze(1).unsqueeze(2).unsqueeze(3) | |
| mu_normed = mu / (mu_sum + 1e-6) | |
| return mu, mu_normed | |
| 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) | |
| cfg = { | |
| 'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512] | |
| } | |
| def vgg19(): | |
| """VGG 19-layer model (configuration "E") | |
| model pre-trained on ImageNet | |
| """ | |
| model = VGG(make_layers(cfg['E'])) | |
| return model |