-
Notifications
You must be signed in to change notification settings - Fork 3
/
train_TCQE_Query2box.py
177 lines (149 loc) · 7.35 KB
/
train_TCQE_Query2box.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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
"""
@date: 2021/10/26
@description: null
"""
from typing import Tuple
import click
import torch
import torch.nn as nn
import torch.nn.functional as F
from ComplexTemporalQueryData import ICEWS05_15, ICEWS14, ComplexTemporalQueryDatasetCachePath, TemporalComplexQueryData, GDELT
from toolbox.exp.OutputSchema import OutputSchema
from toolbox.nn.BetaE import BoxOffsetIntersection, CenterIntersection
from toolbox.utils.RandomSeeds import set_seeds
from train_TCQE_TFLEX import MyExperiment
from TCQE_static_QE import TYPE_token, TCQE
class EntityProjection(nn.Module):
def __init__(self, dim, hidden_dim=800, num_layers=2, drop=0.1):
super(EntityProjection, self).__init__()
def forward(self,
q_center, q_offset,
r_center, r_offset,
t_center, t_offset):
center = q_center + r_center
offset = q_offset + r_offset
return center, offset
class EntityIntersection(nn.Module):
def __init__(self, dim):
super(EntityIntersection, self).__init__()
self.dim = dim
self.entity_dim = dim
self.center_net = CenterIntersection(self.entity_dim)
self.offset_net = BoxOffsetIntersection(self.entity_dim)
def forward(self, center, offset):
# N x B x d
center = self.center_net(center)
offset = self.offset_net(offset)
return center, offset
class EntityUnion(nn.Module):
def __init__(self, dim):
super(EntityUnion, self).__init__()
self.dim = dim
def forward(self, center, offset):
return center, offset
class EntityNegation(nn.Module):
def __init__(self, dim):
super(EntityNegation, self).__init__()
self.dim = dim
def forward(self, center, offset):
raise NotImplementedError("box cannot handle queries with negation")
class TFLEX(TCQE):
def __init__(self, nentity, nrelation, ntimestamp, hidden_dim, gamma,
test_batch_size=1,
center_reg=None, drop: float = 0.):
super(TFLEX, self).__init__(nentity, nrelation, ntimestamp, hidden_dim, gamma, test_batch_size, center_reg, drop)
self.entity_projection = EntityProjection(hidden_dim, drop=drop)
self.entity_intersection = EntityIntersection(hidden_dim)
self.entity_union = EntityUnion(hidden_dim)
self.entity_negation = EntityNegation(hidden_dim)
def entity_feature(self, idx):
return self.entity_feature_embedding(idx)
def entity_token(self, idx) -> TYPE_token:
feature = self.entity_feature(idx)
logic = torch.zeros_like(feature).to(feature.device)
return feature, logic
def relation_token(self, idx) -> TYPE_token:
feature = self.relation_feature_embedding(idx)
logic = self.relation_logic_embedding(idx)
return feature, logic
def distance_between_entity_and_query(self, entity_embedding, query_center_embedding, query_offset_embedding):
"""
entity_embedding (B, 1, N, d)
query_center_embedding (B, 1, 1, dt) or (B, 2, 1, dt)
query_offset_embedding (B, 1, 1, dt) or (B, 2, 1, dt)
"""
delta = (entity_embedding - query_center_embedding).abs()
distance_out = F.relu(delta - query_offset_embedding)
distance_in = torch.min(delta, query_offset_embedding)
distance = torch.norm(distance_out, p=1, dim=-1) + self.cen * torch.norm(distance_in, p=1, dim=-1)
return distance
@click.command()
@click.option("--data_home", type=str, default="data", help="The folder path to dataset.")
@click.option("--dataset", type=str, default="ICEWS14", help="Which dataset to use: ICEWS14, ICEWS05_15, GDELT.")
@click.option("--name", type=str, default="TFLEX_base", help="Name of the experiment.")
@click.option("--start_step", type=int, default=0, help="start step.")
@click.option("--max_steps", type=int, default=200001, help="Number of steps.")
@click.option("--every_test_step", type=int, default=10000, help="Number of steps.")
@click.option("--every_valid_step", type=int, default=10000, help="Number of steps.")
@click.option("--batch_size", type=int, default=512, help="Batch size.")
@click.option("--test_batch_size", type=int, default=8, help="Test batch size.")
@click.option('--negative_sample_size', default=128, type=int, help="negative entities sampled per query")
@click.option("--train_device", type=str, default="cuda:0", help="choice: cuda:0, cuda:1, cpu.")
@click.option("--test_device", type=str, default="cuda:0", help="choice: cuda:0, cuda:1, cpu.")
@click.option("--resume", type=bool, default=False, help="Resume from output directory.")
@click.option("--resume_by_score", type=float, default=0.0, help="Resume by score from output directory. Resume best if it is 0. Default: 0")
@click.option("--lr", type=float, default=0.0001, help="Learning rate.")
@click.option('--cpu_num', type=int, default=1, help="used to speed up torch.dataloader")
@click.option('--hidden_dim', type=int, default=800, help="embedding dimension")
@click.option("--input_dropout", type=float, default=0.1, help="Input layer dropout.")
@click.option('--gamma', type=float, default=15.0, help="margin in the loss")
@click.option('--center_reg', type=float, default=0.02, help='center_reg for ConE, center_reg balances the in_cone dist and out_cone dist')
@click.option('--train_tasks', type=str, default="Pe,Pe2,Pe3,e2i,e3i", help='the tasks for training')
@click.option('--train_all', type=bool, default=False, help='if training all, it will use all tasks in data.train_queries_answers')
@click.option('--eval_tasks', type=str, default="e2i", help='the tasks for evaluation')
@click.option('--eval_all', type=bool, default=False, help='if evaluating all, it will use all tasks in data.test_queries_answers')
def main(data_home, dataset, name,
start_step, max_steps, every_test_step, every_valid_step,
batch_size, test_batch_size, negative_sample_size,
train_device, test_device,
resume, resume_by_score,
lr, cpu_num,
hidden_dim, input_dropout, gamma, center_reg, train_tasks, train_all, eval_tasks, eval_all
):
set_seeds(0)
output = OutputSchema(dataset + "-" + name)
if dataset == "ICEWS14":
dataset = ICEWS14(data_home)
elif dataset == "ICEWS05_15":
dataset = ICEWS05_15(data_home)
elif dataset == "GDELT":
dataset = GDELT(data_home)
cache = ComplexTemporalQueryDatasetCachePath(dataset.cache_path)
data = TemporalComplexQueryData(dataset, cache_path=cache)
data.preprocess_data_if_needed()
data.load_cache(["meta"])
entity_count = data.entity_count
relation_count = data.relation_count
timestamp_count = data.timestamp_count
max_relation_id = relation_count
model = TFLEX(
nentity=entity_count,
nrelation=relation_count + max_relation_id, # with reverse relations
ntimestamp=timestamp_count,
hidden_dim=hidden_dim,
gamma=gamma,
center_reg=center_reg,
test_batch_size=test_batch_size,
drop=input_dropout,
)
MyExperiment(
output, data, model,
start_step, max_steps, every_test_step, every_valid_step,
batch_size, test_batch_size, negative_sample_size,
train_device, test_device,
resume, resume_by_score,
lr, cpu_num,
hidden_dim, input_dropout, gamma, center_reg, train_tasks, train_all, eval_tasks, eval_all
)
if __name__ == '__main__':
main()