-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
48 lines (43 loc) · 1.7 KB
/
model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import torch.nn as nn
class ContextPredictionModel(nn.Module):
def __init__(self):
super(ContextPredictionModel, self).__init__() # vgg 16 base
self.base = nn.Sequential(
# 3 224 128
nn.Conv2d(3, 64, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(64, 64, 3, padding=1), nn.LeakyReLU(0.2),
nn.MaxPool2d(2, 2),
# 64 112 64
nn.Conv2d(64, 128, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(128, 128, 3, padding=1), nn.LeakyReLU(0.2),
nn.MaxPool2d(2, 2),
# 128 56 32
nn.Conv2d(128, 256, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(256, 256, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(256, 256, 3, padding=1), nn.LeakyReLU(0.2),
nn.MaxPool2d(2, 2),
# 256 28 16
nn.Conv2d(256, 512, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(512, 512, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(512, 512, 3, padding=1), nn.LeakyReLU(0.2),
nn.MaxPool2d(2, 2),
# 512 14 8
nn.Conv2d(512, 512, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(512, 512, 3, padding=1), nn.LeakyReLU(0.2),
nn.Conv2d(512, 512, 3, padding=1), nn.LeakyReLU(0.2),
nn.MaxPool2d(2, 2)
)
self.avgpool = nn.AdaptiveAvgPool2d(1)
# 512 1 1
self.classifier = nn.Linear(512, 8)
def forward(self, x_1, x_2):
# center
bs = x_1.shape[0]
x_1 = self.base(x_1)
x_2 = self.base(x_2)
x_1 = self.avgpool(x_1)
x_2 = self.avgpool(x_2)
x = x_1 + x_2
x = x.view(bs, -1)
x = self.classifier(x)
return x