def hc_multi_select(ids_from, ids_to, h_froms, c_froms):
h_vecs = multi_index_select(ids_from,
ids_to,
*h_froms)
c_vecs = multi_index_select(ids_from,
ids_to,
*c_froms)
return h_vecs, c_vecs
def forward_train(self, h_bot, c_bot, h_buf0, c_buf0, prev_rowsum_h,
prrev_rowsum_c):
# embed row tree
tree_agg_ids = TreeLib.PrepareRowEmbed()
row_embeds = [(self.init_h0, self.init_c0)]
if h_bot is not None:
row_embeds.append((h_bot, c_bot))
if prev_rowsum_h is not None:
row_embeds.append((prev_rowsum_h, prrev_rowsum_c))
if h_buf0 is not None:
row_embeds.append((h_buf0, c_buf0))
th_bot = h_bot
tc_bot = c_bot
for i, all_ids in enumerate(tree_agg_ids):
fn_ids = lambda x: all_ids[x]
if i:
th_bot = tc_bot = None
new_states = batch_tree_lstm3(th_bot, tc_bot, row_embeds[-1][0],
row_embeds[-1][1], prev_rowsum_h,
prrev_rowsum_c, fn_ids,
self.merge_cell)
row_embeds.append(new_states)
h_list, c_list = zip(*row_embeds)
joint_h = torch.cat(h_list, dim=0)
joint_c = torch.cat(c_list, dim=0)
# get history representation
init_select, all_ids, last_tos, next_ids, pos_info = TreeLib.PrepareRowSummary(
)
cur_state = (joint_h[init_select], joint_c[init_select])
ret_state = (joint_h[next_ids], joint_c[next_ids])
hist_rnn_states = []
hist_froms = []
hist_tos = []
for i, (done_from, done_to, proceed_from,
proceed_input) in enumerate(all_ids):
hist_froms.append(done_from)
hist_tos.append(done_to)
hist_rnn_states.append(cur_state)
next_input = joint_h[proceed_input], joint_c[proceed_input]
sub_state = cur_state[0][proceed_from], cur_state[1][proceed_from]
cur_state = self.summary_cell(sub_state, next_input)
hist_rnn_states.append(cur_state)
hist_froms.append(None)
hist_tos.append(last_tos)
hist_h_list, hist_c_list = zip(*hist_rnn_states)
pos_embed = self.pos_enc(pos_info)
row_h = multi_index_select(hist_froms, hist_tos, *
hist_h_list) + pos_embed
row_c = multi_index_select(hist_froms, hist_tos, *
hist_c_list) + pos_embed
return (row_h, row_c), ret_state
def tree_state_select(h_bot, c_bot, h_buf, c_buf, fn_all_ids):
bot_froms, bot_tos, prev_froms, prev_tos = fn_all_ids()
if h_buf is None or prev_tos is None:
h_vecs = multi_index_select([bot_froms], [bot_tos], h_bot)
c_vecs = multi_index_select([bot_froms], [bot_tos], c_bot)
elif h_bot is None or bot_tos is None:
h_vecs = multi_index_select([prev_froms], [prev_tos], h_buf)
c_vecs = multi_index_select([prev_froms], [prev_tos], c_buf)
else:
h_vecs, c_vecs = hc_multi_select([bot_froms, prev_froms],
[bot_tos, prev_tos],
[h_bot, h_buf], [c_bot, c_buf])
return h_vecs, c_vecs
def forward_train(self, graph_ids, list_node_starts=None, num_nodes=-1, prev_rowsum_states=[None, None], list_col_ranges=None):
fn_hc_bot, h_buf_list, c_buf_list = self.forward_row_trees(graph_ids, list_node_starts, num_nodes, list_col_ranges)
row_states, next_states = self.row_tree.forward_train(*(fn_hc_bot(0)), h_buf_list[0], c_buf_list[0], *prev_rowsum_states)
# make prediction
logit_has_edge = self.pred_has_ch(row_states[0])
has_ch, _ = TreeLib.GetChLabel(0, dtype=np.bool)
ll = self.binary_ll(logit_has_edge, has_ch)
# has_ch_idx
cur_states = (row_states[0][has_ch], row_states[1][has_ch])
lv = 0
while True:
is_nonleaf = TreeLib.QueryNonLeaf(lv)
if is_nonleaf is None or np.sum(is_nonleaf) == 0:
break
cur_states = (cur_states[0][is_nonleaf], cur_states[1][is_nonleaf])
left_logits = self.pred_has_left(cur_states[0], lv)
has_left, num_left = TreeLib.GetChLabel(-1, lv)
left_update = self.topdown_left_embed[has_left] + self.tree_pos_enc(num_left)
left_ll, float_has_left = self.binary_ll(left_logits, has_left, need_label=True, reduction='sum')
ll = ll + left_ll
cur_states = self.cell_topdown(left_update, cur_states, lv)
left_ids = TreeLib.GetLeftRootStates(lv)
h_bot, c_bot = fn_hc_bot(lv + 1)
if lv + 1 < len(h_buf_list):
h_next_buf, c_next_buf = h_buf_list[lv + 1], c_buf_list[lv + 1]
else:
h_next_buf = c_next_buf = None
left_subtree_states = tree_state_select(h_bot, c_bot,
h_next_buf, c_next_buf,
lambda: left_ids)
has_right, num_right = TreeLib.GetChLabel(1, lv)
right_pos = self.tree_pos_enc(num_right)
left_subtree_states = [x + right_pos for x in left_subtree_states]
topdown_state = self.l2r_cell(cur_states, left_subtree_states, lv)
right_logits = self.pred_has_right(topdown_state[0], lv)
right_update = self.topdown_right_embed[has_right]
topdown_state = self.cell_topright(right_update, topdown_state, lv)
right_ll = self.binary_ll(right_logits, has_right, reduction='none') * float_has_left
ll = ll + torch.sum(right_ll)
lr_ids = TreeLib.GetLeftRightSelect(lv, np.sum(has_left), np.sum(has_right))
new_states = []
for i in range(2):
new_s = multi_index_select([lr_ids[0], lr_ids[2]], [lr_ids[1], lr_ids[3]],
cur_states[i], topdown_state[i])
new_states.append(new_s)
cur_states = tuple(new_states)
lv += 1
return ll, next_states