| import torch.nn as nn |
| from torchvision.models import vgg19, VGG19_Weights |
|
|
| class VGGEncoder(nn.Module): |
| def __init__(self): |
| super().__init__() |
| vgg = vgg19(weights=VGG19_Weights.DEFAULT).features |
| self.slice1 = nn.Sequential() |
| self.slice2 = nn.Sequential() |
| self.slice3 = nn.Sequential() |
| self.slice4 = nn.Sequential() |
|
|
| for x in range(2): |
| self.slice1.add_module(str(x), vgg[x]) |
| for x in range(2, 7): |
| self.slice2.add_module(str(x), vgg[x]) |
| for x in range(7, 12): |
| self.slice3.add_module(str(x), vgg[x]) |
| for x in range(12, 21): |
| self.slice4.add_module(str(x), vgg[x]) |
| |
| for param in self.parameters(): |
| param.requires_grad = False |
|
|
| def forward(self, x): |
| h1 = self.slice1(x) |
| h2 = self.slice2(h1) |
| h3 = self.slice3(h2) |
| h4 = self.slice4(h3) |
| return h1, h2, h3, h4 |
