def forward(self, input, mask):
'''
:param input: [batch, N, F]
:return:
'''
batch, N, input_dim = input.size()
if self.conv:
input = torch.cat(
(input, torch.zeros(
(batch, 2, input_dim), device=input.device)),
dim=1)
h = self.W(input.transpose(1, 2)).transpose(1, 2)
else:
# print(input.device, self.W.device)
#print(input.shape, input.device, self.W.weight.device)
h = torch.cat([
self.W1(input[:, :get_Parameter('taxi_size')]),
self.W2(input[:, get_Parameter('taxi_size'):])
],
dim=1)
#h = self.W(input)
# h = torch.matmul(input, self.W)
# a_input = torch.cat([h.repeat(1, 1, N).view(batch, N * N, -1), h.repeat(1, N, 1)], dim=2).view(batch, N, -1,
# 2 * self.outfeatures)
A1 = torch.matmul(h, self.a[:self.outfeatures])
A2 = torch.matmul(h, self.a[self.outfeatures:])
sim = torch.add(A1.expand(-1, -1, N),
torch.transpose(A2.expand(-1, -1, N), 1, 2))
del A1, A2
# print(sim)
e = self.leakyrelu(sim)
if mask is not None:
n = e.shape[0]
#mask = (mask>0).unsqueeze(0).repeat(n, 1, 1)
#print(mask)
e = mask.unsqueeze(0).repeat(n, 1, 1) * e
#e = e.masked_fill(mask, -np.inf)
#attention = F.softmax(e, dim=2)
# print(attention)
# zero_vec = -9e15 * torch.ones_like(e)
# print(attention)
# print(torch.gt(attention, 0.1).sum(dim=2))
# if mask is not None:
# mask = mask.repeat(156, 1, 1)
# e = torch.where(mask>0, e, zero_vec)
# e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(3))
# print(attention)
attention = F.softmax(e, dim=2)
attention = F.dropout(attention, self.dropout, training=self.training)
h_prime = torch.bmm(attention, h)
return F.elu(h_prime, inplace=True), attention
def __init__(self, embedding_size, hidden_dim):
super(FeatureEmbedding, self).__init__()
self.field_size = 7
self.embedding_size = embedding_size
self.hidden_dim = hidden_dim
# feature : traffic_mode, row, column, station, day, hour, work
'''
init fm part
'''
self.feature_sizes = [
2, 1, 1, 200 + get_Parameter('bike_size'), 7, 24, 2
]
self.fm_first_order_embeddings = nn.ModuleList([
nn.Embedding(feature_size, 1)
for feature_size in self.feature_sizes
])
self.fm_second_order_embeddings = nn.ModuleList([
nn.Embedding(feature_size, embedding_size)
for feature_size in self.feature_sizes
])
'''
init deep part
'''
all_dims = [self.field_size * embedding_size] + hidden_dim
deep_layers = []
for i in range(1, len(hidden_dim) + 1):
deep_layers.append(
nn.Sequential(nn.Linear(all_dims[i - 1], all_dims[i]),
nn.BatchNorm1d(all_dims[i]), nn.ReLU()))
self.deep_module = nn.ModuleList(deep_layers)
self.concat = nn.Linear(self.field_size + self.embedding_size * 2,
embedding_size)
def forward(self, graph, feat):
'''
:param graph: DGLGraph
:param feat: <N, b, F>
:return:
'''
with graph.local_scope():
N, b, _ = feat.size()
graph = graph.local_var()
graph = graph.to(feat.device)
feat = torch.cat([self.fc1(feat[:get_Parameter('taxi_size')]), self.fc2(feat[get_Parameter('taxi_size'):])], dim=0)
feat_src = feat_dst = feat.view(N, b, self._num_heads, self._out_feats)
#feat_src = feat_dst = self.fc(feat).view(N, b, self._num_heads, self._out_feats)
el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(-1)
er = (feat_dst * self.attn_l).sum(dim=-1).unsqueeze(-1)
graph.srcdata.update({'ft': feat_src, 'el': el})
graph.dstdata.update({'er': er})
graph.apply_edges(fn.u_add_v('el', 'er', 'e'))
#graph.apply_edges(fn.u_mul_e('e', 'w', 'e'))
e = self.leaky_relu(graph.edata.pop('e'))
graph.edata['a'] = self.attn_drop(edge_softmax(graph, e))
#print(graph.edata['a'].size())
graph.update_all(fn.u_mul_e('ft', 'a', 'm'), fn.sum('m', 'ft'))
rst = graph.dstdata['ft']
rst = rst.reshape(N, -1, self._num_heads*self._out_feats)
return rst, graph.edata['a']
def forward(self, input, features_input, features_target):
'''
:param input: <batch, T, N, F>
:param features_input: <batch, T, N, F>
:param features_target: <batch, 1, N, F>
:return: <batch, 1, N, 2>
'''
batch, T, N, input_dim = input.size()
attn_trans = torch.matmul(
features_input.transpose(1, 2),
features_target.transpose(1, 2).transpose(2, 3))
attn = torch.softmax(attn_trans, dim=2)
out = torch.matmul(input.permute(0, 2, 3, 1), attn).transpose(2, 3)
out = torch.cat([
self.out1(out[:, :get_Parameter('taxi_size')]),
self.out2(out[:, get_Parameter('taxi_size'):])
],
dim=1)
return out.transpose(1, 2)
def __init__(self, field_size, embedding_size):
super(SingleEmbedding, self).__init__()
self.field_size = field_size
self.embedding_size = embedding_size
self.feature_sizes = [
2, 1, 1, 200 + get_Parameter('bike_size'), 7, 24, 2
]
self.field_embeddings = nn.ModuleList([
nn.Embedding(feature_size, embedding_size)
for feature_size in self.feature_sizes
])
def forward(self, embedding, input):
'''
:param input: [batch, T, input_dim]
:param embedding: [batch, T, input_dim]
:return:
'''
batch, T, _ = embedding.size()
h = torch.matmul(embedding, self.W)
a_input = torch.cat(
[h.repeat(1, 1, T).view(batch, T * T, -1),
h.repeat(1, T, 1)],
dim=2).view(batch, T, -1, 2 * self.outfeatures)
e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(3))
attn = F.softmax(e, dim=2)
attn = attn.repeat(get_Parameter('input_size'), 1, 1)
attention = F.dropout(attn, self.dropout, training=self.training)
out = torch.bmm(attention, input)
return out
