import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import vgg11


class VGG11Embedding(nn.Module):
    def __init__(self, embedding_size, weights=None):
        super(VGG11Embedding, self).__init__()
        vgg = vgg11(weights=weights)
        self.features = vgg.features
        self.linear = nn.Linear(512, embedding_size)
        
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.linear(x)
        x = F.normalize(x, p=2, dim=1)
        return x