|
|
import torch
|
|
|
from torch import nn
|
|
|
from torchvision import models
|
|
|
from einops import rearrange
|
|
|
from torchvision.models._utils import IntermediateLayerGetter
|
|
|
|
|
|
|
|
|
class Vgg(nn.Module):
|
|
|
def __init__(self, name, ss, ks, hidden, pretrained=True, dropout=0.5):
|
|
|
super(Vgg, self).__init__()
|
|
|
|
|
|
if name == 'vgg11_bn':
|
|
|
cnn = models.vgg11_bn(weights='DEFAULT')
|
|
|
elif name == 'vgg19_bn':
|
|
|
cnn = models.vgg19_bn(weights='DEFAULT')
|
|
|
|
|
|
pool_idx = 0
|
|
|
|
|
|
for i, layer in enumerate(cnn.features):
|
|
|
if isinstance(layer, torch.nn.MaxPool2d):
|
|
|
cnn.features[i] = torch.nn.AvgPool2d(kernel_size=ks[pool_idx], stride=ss[pool_idx], padding=0)
|
|
|
pool_idx += 1
|
|
|
|
|
|
self.features = cnn.features
|
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
self.last_conv_1x1 = nn.Conv2d(512, hidden, 1)
|
|
|
|
|
|
def forward(self, x):
|
|
|
"""
|
|
|
Shape:
|
|
|
- x: (N, C, H, W)
|
|
|
- output: (W, N, C)
|
|
|
"""
|
|
|
|
|
|
conv = self.features(x)
|
|
|
conv = self.dropout(conv)
|
|
|
conv = self.last_conv_1x1(conv)
|
|
|
|
|
|
|
|
|
conv = conv.transpose(-1, -2)
|
|
|
conv = conv.flatten(2)
|
|
|
conv = conv.permute(-1, 0, 1)
|
|
|
return conv
|
|
|
|
|
|
def vgg11_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
|
|
|
return Vgg('vgg11_bn', ss, ks, hidden, pretrained, dropout)
|
|
|
|
|
|
def vgg19_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
|
|
|
return Vgg('vgg19_bn', ss, ks, hidden, pretrained, dropout)
|
|
|
|
|
|
|
