-
Notifications
You must be signed in to change notification settings - Fork 1
/
model.py
66 lines (52 loc) · 1.71 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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import torch
import torch.nn as nn
import torch.nn.functional as F
class GNet(nn.Module):
def __init__(self):
super(GNet, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 64, 3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(3, stride=1),
nn.Conv2d(128, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(3, stride=1),
nn.Conv2d(64, 1, 3, padding=1),
nn.Tanh(),
)
def forward(self, x, y):
x = self.conv(x)
# applied
y = y.view(-1, 1, 3, 1)
x = torch.matmul(x, y)
x = x.view(x.size(0), -1)
return x
class Net(nn.Module):
def __init__(self, nfeat, outdims=3):
super(Net, self).__init__()
self.nfeat = nfeat
self.outdims = outdims
self.GNetList = nn.ModuleList([GNet() for _ in range(nfeat)])
self.fc = nn.Sequential(
nn.BatchNorm1d(3*nfeat),
nn.Linear(3*nfeat, 64),
nn.BatchNorm1d(64),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(64, 32),
nn.BatchNorm1d(32),
nn.ReLU(),
nn.Dropout(p=0.3),
nn.Linear(32, outdims),
)
def forward(self, x, y):
x = torch.cat([subnet(x[:,i:i+1,:,:], y[:,i*3:i*3+3]) for i, subnet in enumerate(self.GNetList)], dim=1)
x = self.fc(x)
x = F.normalize(x)
return x
def load_model(path_model, nfeatures):
model = Net(nfeatures)
model.load_state_dict(torch.load(path_model))
return model