Spaces:
Sleeping
Sleeping
import torch | |
import torch.nn as nn | |
from typing import Union, List, Dict, Any, cast | |
import torchvision | |
import torch.nn.functional as F | |
class VGG(torch.nn.Module): | |
def __init__(self, arch_type, pretrained, progress): | |
super().__init__() | |
self.layer1 = torch.nn.Sequential() | |
self.layer2 = torch.nn.Sequential() | |
self.layer3 = torch.nn.Sequential() | |
self.layer4 = torch.nn.Sequential() | |
self.layer5 = torch.nn.Sequential() | |
if arch_type == 'vgg11': | |
official_vgg = torchvision.models.vgg11(pretrained=pretrained, progress=progress) | |
blocks = [ [0,2], [2,5], [5,10], [10,15], [15,20] ] | |
last_idx = 20 | |
elif arch_type == 'vgg19': | |
official_vgg = torchvision.models.vgg19(pretrained=pretrained, progress=progress) | |
blocks = [ [0,4], [4,9], [9,18], [18,27], [27,36] ] | |
last_idx = 36 | |
else: | |
raise NotImplementedError | |
for x in range( *blocks[0] ): | |
self.layer1.add_module(str(x), official_vgg.features[x]) | |
for x in range( *blocks[1] ): | |
self.layer2.add_module(str(x), official_vgg.features[x]) | |
for x in range( *blocks[2] ): | |
self.layer3.add_module(str(x), official_vgg.features[x]) | |
for x in range( *blocks[3] ): | |
self.layer4.add_module(str(x), official_vgg.features[x]) | |
for x in range( *blocks[4] ): | |
self.layer5.add_module(str(x), official_vgg.features[x]) | |
self.max_pool = official_vgg.features[last_idx] | |
self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) | |
self.fc1 = official_vgg.classifier[0] | |
self.fc2 = official_vgg.classifier[3] | |
self.fc3 = official_vgg.classifier[6] | |
self.dropout = nn.Dropout() | |
def forward(self, x): | |
out = {} | |
x = self.layer1(x) | |
out['f0'] = x | |
x = self.layer2(x) | |
out['f1'] = x | |
x = self.layer3(x) | |
out['f2'] = x | |
x = self.layer4(x) | |
out['f3'] = x | |
x = self.layer5(x) | |
out['f4'] = x | |
x = self.max_pool(x) | |
x = self.avgpool(x) | |
x = x.view(-1,512*7*7) | |
x = self.fc1(x) | |
x = F.relu(x) | |
x = self.dropout(x) | |
x = self.fc2(x) | |
x = F.relu(x) | |
out['penultimate'] = x | |
x = self.dropout(x) | |
x = self.fc3(x) | |
out['logits'] = x | |
return out | |
def vgg11(pretrained=False, progress=True): | |
r"""VGG 11-layer model (configuration "A") from | |
`"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. | |
Args: | |
pretrained (bool): If True, returns a model pre-trained on ImageNet | |
progress (bool): If True, displays a progress bar of the download to stderr | |
""" | |
return VGG('vgg11', pretrained, progress) | |
def vgg19(pretrained=False, progress=True): | |
r"""VGG 19-layer model (configuration "E") | |
`"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_. | |
Args: | |
pretrained (bool): If True, returns a model pre-trained on ImageNet | |
progress (bool): If True, displays a progress bar of the download to stderr | |
""" | |
return VGG('vgg19', pretrained, progress) | |