-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathKonCept.py
35 lines (34 loc) · 1.15 KB
/
KonCept.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
from inceptionresnetv2 import inceptionresnetv2
import torch.nn as nn
class KonCept(nn.Module):
def __init__(self, config):
super(KonCept,self).__init__()
base_model = inceptionresnetv2(num_classes=1000, pretrained='imagenet')
self.base= nn.Sequential(*list(base_model.children())[:-2])
self.avg_pool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
self.fc = nn.Sequential(
nn.Linear(1536, 2048),
nn.ReLU(inplace=True),
nn.BatchNorm1d(2048),
nn.Dropout(p=0.25),
nn.Linear(2048, 1024),
nn.ReLU(inplace=True),
nn.BatchNorm1d(1024),
nn.Dropout(p=0.25),
nn.Linear(1024, 256),
nn.ReLU(inplace=True),
nn.BatchNorm1d(256),
nn.Dropout(p=0.5),
nn.Linear(256, 1),
)
if config.fc:
for param in self.base.parameters():
param.requires_grad = False
def forward(self, x):
x = self.base(x)
x = self.avg_pool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
output = []
output.append(x)
return output, x