-
Notifications
You must be signed in to change notification settings - Fork 26
/
Copy pathModel.py
61 lines (47 loc) · 1.53 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
import torch
from torch import nn
class Date2VecConvert:
def __init__(self, model_path="./d2v_model/d2v_98291_17.169918439404636.pth"):
self.model = torch.load(model_path, map_location='cpu').eval()
def __call__(self, x):
with torch.no_grad():
return self.model.encode(torch.Tensor(x).unsqueeze(0)).squeeze(0).cpu()
class Date2Vec(nn.Module):
def __init__(self, k=32, act="sin"):
super(Date2Vec, self).__init__()
if k % 2 == 0:
k1 = k // 2
k2 = k // 2
else:
k1 = k // 2
k2 = k // 2 + 1
self.fc1 = nn.Linear(6, k1)
self.fc2 = nn.Linear(6, k2)
self.d2 = nn.Dropout(0.3)
if act == 'sin':
self.activation = torch.sin
else:
self.activation = torch.cos
self.fc3 = nn.Linear(k, k // 2)
self.d3 = nn.Dropout(0.3)
self.fc4 = nn.Linear(k // 2, 6)
self.fc5 = torch.nn.Linear(6, 6)
def forward(self, x):
out1 = self.fc1(x)
out2 = self.d2(self.activation(self.fc2(x)))
out = torch.cat([out1, out2], 1)
out = self.d3(self.fc3(out))
out = self.fc4(out)
out = self.fc5(out)
return out
def encode(self, x):
out1 = self.fc1(x)
out2 = self.activation(self.fc2(x))
out = torch.cat([out1, out2], 1)
return out
if __name__ == "__main__":
model = Date2Vec()
inp = torch.randn(1, 6)
out = model(inp)
print(out)
print(out.shape)