def get_feat(self, graph_list):
node_feat, edge_feat, edge_from_idx, edge_to_idx, g_idx = prepare_gnn(
graph_list, self.is_cuda())
input_node_linear = self.w_n2l(node_feat)
input_message = input_node_linear
if edge_feat is not None:
input_edge_linear = self.w_e2l(edge_feat)
e2npool_input = scatter_add(input_edge_linear,
edge_to_idx,
dim=0,
dim_size=node_feat.shape[0])
input_message += e2npool_input
input_potential = self.act_func(input_message)
cur_message_layer = input_potential
all_embeds = [cur_message_layer]
edge_index = [edge_from_idx, edge_to_idx]
edge_index = torch.stack(edge_index)
for lv in range(self.max_lv):
node_linear = self.conv_layers[lv](cur_message_layer, edge_index)
merged_linear = node_linear + input_message
cur_message_layer = self.act_func(merged_linear)
all_embeds.append(cur_message_layer)
return self.readout_net(all_embeds, g_idx, len(graph_list))
def get_feat(self, graph_list):
node_feat, edge_feat, edge_from_idx, edge_to_idx, g_idx = prepare_gnn(
graph_list, self.is_cuda())
input_node_linear = self.w_n2l(node_feat)
input_message = input_node_linear
if edge_feat is not None:
input_edge_linear = self.w_e2l[0](edge_feat)
e2npool_input = scatter_add(input_edge_linear,
edge_to_idx,
dim=0,
dim_size=node_feat.shape[0])
input_message += e2npool_input
input_potential = self.act_func(input_message)
input_potential = self.msg_bn[0](input_potential)
cur_message_layer = input_potential
all_embeds = [cur_message_layer]
edge_index = [edge_from_idx, edge_to_idx]
for lv in range(self.max_lv):
node_linear = self.conv_layers[lv](cur_message_layer, edge_index)
edge_linear = self.w_e2l[lv + 1](edge_feat)
e2npool_input = scatter_add(edge_linear,
edge_to_idx,
dim=0,
dim_size=node_linear.shape[0])
merged_hidden = self.act_func(node_linear + e2npool_input)
merged_hidden = self.hidden_bn[lv](merged_hidden)
residual_out = self.conv_l2[lv](merged_hidden) + cur_message_layer
cur_message_layer = self.act_func(residual_out)
cur_message_layer = self.msg_bn[lv + 1](cur_message_layer)
all_embeds.append(cur_message_layer)
return self.readout_net(all_embeds, g_idx, len(graph_list))
def get_feat(self, graph_list):
node_feat, edge_feat, edge_from_idx, edge_to_idx, g_idx = prepare_gnn(
graph_list, self.is_cuda())
out = self.act_func(self.lin0(node_feat))
h = out.unsqueeze(0)
edge_index = [edge_from_idx, edge_to_idx]
edge_index = torch.stack(edge_index)
for lv in range(self.max_lv):
m = self.act_func(self.conv(out, edge_index, edge_feat))
out, h = self.gru(m.unsqueeze(0), h)
out = out.squeeze(0)
out = self.set2set(out, g_idx)
out = self.readout(out)
return out, None
def get_feat(self, graph_list):
node_feat, edge_feat, edge_from_idx, edge_to_idx, g_idx = prepare_gnn(
graph_list, self.is_cuda())
edge_index = [edge_from_idx, edge_to_idx]
node_states = self.node2hidden(node_feat)
init_embed = self.readout_funcs[-1](node_states)
outs = self.readout_agg(init_embed,
g_idx,
dim=0,
dim_size=len(graph_list))
for i in range(self.max_lv):
layer = self.layers[i]
new_states = layer(node_states, edge_index, edge_feat)
node_states = new_states
if self.out_method == 'last':
continue
out_states = self.readout_funcs[i](node_states)
graph_embed = self.readout_agg(out_states,
g_idx,
dim=0,
dim_size=len(graph_list))
if self.out_method == 'gru':
outs = self.final_cell(graph_embed, outs)
else:
outs += graph_embed
if self.out_method == 'last':
out_states = self.readout_funcs[0](node_states)
graph_embed = self.readout_agg(out_states,
g_idx,
dim=0,
dim_size=len(graph_list))
return graph_embed, (g_idx, out_states)
else:
if self.out_method == 'mean':
outs /= self.max_lv + 1
return outs, None