def __call__(self, n_id):
data_flow = self.sampler(n_id)
num_layers = len(self.convs)
jk_hidden = []
x = self.to_x(data_flow[0].n_id)
for i, conv, block in zip(range(num_layers), self.convs, data_flow):
if block.e_id is None:
edge_attr = None
else:
edge_attr = self.get_edge_attr(block)
x_src = mp_ops.gather(x, block.res_n_id)
x_dst = None if self.whole_graph else x
x = self.calculate_conv(conv, [x_src, x_dst],
block.edge_index,
size=block.size,
edge_attr=edge_attr)
x = tf.nn.relu(x)
jk_hidden.append(x)
for j in range(i):
jk_hidden[j] = mp_ops.gather(jk_hidden[j], block.res_n_id)
if self.jk_mode == 'concat':
x = tf.concat(jk_hidden, axis=1)
elif self.jk_mode == 'maxpool':
x = tf.reduce_sum(tf.stack(jk_hidden, 1), 1)
x = self.fc(x)
return x
def get_edge_attr(self, block):
n_id_dst = tf.cast(tf.expand_dims(block.n_id, -1), dtype=tf.float32)
n_id_src = mp_ops.gather(n_id_dst, block.res_n_id)
n_id_src = mp_ops.gather(n_id_src, block.edge_index[0])
n_id_dst = mp_ops.gather(n_id_dst, block.edge_index[1])
n_id_src = tf.cast(tf.squeeze(n_id_src, -1), dtype=tf.int64)
n_id_dst = tf.cast(tf.squeeze(n_id_dst, -1), dtype=tf.int64)
edge_attr = self.to_edge(n_id_src, n_id_dst, block.e_id)
return edge_attr
def __call__(self, group_n_id):
group_x = []
for sampler, n_id in zip(self.group_sampler, group_n_id):
data_flow = sampler(n_id)
num_layers = len(self.convs)
x = self.to_x(data_flow[0].n_id)
for i, conv, block in zip(range(num_layers), self.convs,
data_flow):
if block.e_id is None:
edge_attr = None
else:
edge_attr = self.get_edge_attr(block)
if block.edge_weight is None:
edge_weight = None
else:
edge_weight = block.edge_weight
x_src = mp_ops.gather(x, block.res_n_id)
x_dst = x
x = self.calculate_conv(conv, (x_src, x_dst),
block.edge_index,
size=block.size,
edge_attr=edge_attr,
edge_weight=edge_weight)
x = tf.nn.relu(x)
x = self.fc(x)
group_x.append(x)
return group_x
def testGatherGrad(self):
x = tf.constant([[1., 2.], [3., 4.], [5., 6.]])
idx = tf.constant([1, 0, 1, 2])
out = mp_ops.gather(x, idx)
with self.test_session():
diff = tf.test.compute_gradient_error(x, [3, 2], out, [4, 2])
self.assertLess(diff, 1e-4)
def testGather(self):
x = tf.constant([[1., 2.], [3., 4.], [5., 6.]])
idx = tf.constant([1, 0, 1, 2])
out = mp_ops.gather(x, idx)
with self.test_session():
self.assertAllEqual([[3., 4.], [1., 2.], [3., 4.], [5., 6]],
out.eval())
def gather_feature(self, features, edge_index):
convert_features = []
for feature in features:
convert_feature = []
assert isinstance(feature, tuple) or isinstance(feature, list)
assert len(feature) == 2
if feature[1] is None:
feature[1] = feature[0]
for idx, tmp in enumerate(feature):
if tmp is not None:
tmp = mp_ops.gather(tmp, edge_index[idx])
convert_feature.append(tmp)
convert_features.append(convert_feature)
return convert_features
def __call__(self, n_id):
data_flow = self.sampler(n_id)
num_layers = len(self.convs)
x = self.to_x(data_flow[0].n_id)
for i, conv, block in zip(range(num_layers), self.convs, data_flow):
if block.e_id is None:
edge_attr = None
else:
edge_attr = self.get_edge_attr(block)
x_src = mp_ops.gather(x, block.res_n_id)
x_dst = None if self.whole_graph else x
x = self.calculate_conv(conv, (x_src, x_dst),
block.edge_index,
size=block.size,
edge_attr=edge_attr)
x = tf.nn.relu(x)
x = self.fc(x)
return x