Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| class VGG16(torch.nn.Module): | |
| def __init__(self, num_features, num_classes): | |
| super(VGG16, self).__init__() | |
| # calculate same padding: | |
| # (w - k + 2*p)/s + 1 = o | |
| # => p = (s(o-1) - w + k)/2 | |
| self.conv1_1 = nn.Sequential(nn.Conv2d(in_channels=3, | |
| out_channels=64, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1),) | |
| # (1(32-1)- 32 + 3)/2 = 1 | |
| self.block_1 = nn.Sequential( | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=64, | |
| out_channels=64, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| #nn.MaxPool2d(kernel_size=(2, 2),stride=(2, 2)) | |
| ) | |
| self.conv2_1 = nn.Sequential(nn.Conv2d(in_channels=64, | |
| out_channels=128, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1),) | |
| self.block_2 = nn.Sequential( | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=128, | |
| out_channels=128, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| #nn.MaxPool2d(kernel_size=(2, 2),stride=(2, 2)) | |
| ) | |
| self.conv3_1 = nn.Sequential(nn.Conv2d(in_channels=128, | |
| out_channels=256, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1),) | |
| self.block_3 = nn.Sequential( | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=256, | |
| out_channels=256, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=256, | |
| out_channels=256, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=256, | |
| out_channels=256, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| #nn.MaxPool2d(kernel_size=(2, 2),stride=(2, 2)) | |
| ) | |
| self.conv4_1 = nn.Sequential(nn.Conv2d(in_channels=256, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1),) | |
| self.block_4 = nn.Sequential( | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=512, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=512, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=512, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| #nn.MaxPool2d(kernel_size=(2, 2),stride=(2, 2)) | |
| ) | |
| self.conv5_1 = nn.Sequential(nn.Conv2d(in_channels=512, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1),) | |
| self.block_5 = nn.Sequential( | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=512, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=512, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| nn.Conv2d(in_channels=512, | |
| out_channels=512, | |
| kernel_size=(3, 3), | |
| stride=(1, 1), | |
| padding=1), | |
| nn.ReLU(), | |
| #nn.MaxPool2d(kernel_size=(2, 2),stride=(2, 2)) | |
| ) | |
| self.classifier = nn.Sequential( | |
| nn.Linear(512, 4096), | |
| nn.ReLU(True), | |
| nn.Linear(4096, 4096), | |
| nn.ReLU(True), | |
| nn.Linear(4096, num_classes) | |
| ) | |
| for m in self.modules(): | |
| if isinstance(m, torch.nn.Conv2d): | |
| #n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels | |
| #m.weight.data.normal_(0, np.sqrt(2. / n)) | |
| m.weight.detach().normal_(0, 0.05) | |
| if m.bias is not None: | |
| m.bias.detach().zero_() | |
| elif isinstance(m, torch.nn.Linear): | |
| m.weight.detach().normal_(0, 0.05) | |
| m.bias.detach().detach().zero_() | |
| def forward(self, x): | |
| x_conv = self.conv1_1(x) | |
| x = self.block_1(x_conv) | |
| x2_conv = self.conv2_1(x) | |
| x_2 = self.block_2(x2_conv) | |
| x3_conv = self.conv3_1(x_2) | |
| x_3 = self.block_3(x3_conv) | |
| x4_conv = self.conv4_1(x_3) | |
| x_4 = self.block_4(x4_conv) | |
| x5_conv = self.conv5_1(x_4) | |
| x_5 = self.block_5(x5_conv) | |
| result_dict = { | |
| "style" : [x_conv, x2_conv,x3_conv, x4_conv, x5_conv], | |
| "content" : [x,x_2,x_3,x_4,x_5] | |
| } | |
| #logits = self.classifier(x.view(-1, 512)) | |
| #probas = F.softmax(logits, dim=1) | |
| return result_dict | |