-
Notifications
You must be signed in to change notification settings - Fork 12
/
MeLU.py
55 lines (46 loc) · 2.03 KB
/
MeLU.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
import torch
import numpy as np
from copy import deepcopy
from torch.autograd import Variable
from torch.nn import functional as F
from collections import OrderedDict
from embeddings import item, user
import torch.nn as nn
class Linear(nn.Linear): #used in MAML to forward input with fast weight
def __init__(self, in_features, out_features):
super(Linear, self).__init__(in_features, out_features)
self.weight.fast = None #Lazy hack to add fast weight link
self.bias.fast = None
def forward(self, x):
if self.weight.fast is not None and self.bias.fast is not None:
out = F.linear(x, self.weight.fast, self.bias.fast) #weight.fast (fast weight) is the temporaily adapted weight
else:
out = super(Linear, self).forward(x)
return out
class user_preference_estimator(torch.nn.Module):
def __init__(self, config):
super(user_preference_estimator, self).__init__()
self.embedding_dim = config.embedding_dim
self.fc1_in_dim = config.embedding_dim * 8
self.fc2_in_dim = config.first_fc_hidden_dim
self.fc2_out_dim = config.second_fc_hidden_dim
self.item_emb = item(config)
self.user_emb = user(config)
self.fc1 = Linear(self.fc1_in_dim, self.fc2_in_dim)
self.fc2 = Linear(self.fc2_in_dim, self.fc2_out_dim)
self.linear_out = Linear(self.fc2_out_dim, 1)
self.final_part = nn.Sequential(self.fc1, nn.ReLU(), self.fc2, nn.ReLU(), self.linear_out)
def forward(self, x, training = True):
rate_idx = x[:, 0]
genre_idx = x[:, 1:26]
director_idx = x[:, 26:2212]
actor_idx = x[:, 2212:10242]
gender_idx = x[:, 10242]
age_idx = x[:, 10243]
occupation_idx = x[:, 10244]
area_idx = x[:, 10245]
item_emb = self.item_emb(rate_idx, genre_idx, director_idx, actor_idx)
user_emb = self.user_emb(gender_idx, age_idx, occupation_idx, area_idx)
x = torch.cat((item_emb, user_emb), 1)
x = self.final_part(x)
return x