def test_homogeneous_graph(self):
spec = FeatureSpec(10)
for i in range(3):
spec.append_dense()
total_dim = 3
for i in range(7):
dim = random.randint(8, 10)
spec.append_sparse(20 + 10 * i, dim, False)
total_dim += dim
hops = [4, 5]
# the centric vertices share the same spec with 2-hop neighbors
schema = [("nodes", spec), ("nodes", spec), ("nodes", spec)]
# [f_num, batch_size] = [3, 2]
batch_floats = np.array([[1.0 * i, 2.0 * i] for i in range(3)])
batch_floats = tf.convert_to_tensor(batch_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2]
batch_ints = np.array([[i, 2 * i] for i in range(7)])
batch_ints = tf.convert_to_tensor(batch_ints, dtype=tf.int64)
vertices = Vertex(floats=batch_floats, ints=batch_ints)
# [f_num, batch_size] = [3, 2 * 4]
hop1_floats = np.array([[1.0 * i, 2.0 * i] * hops[0]
for i in range(3)])
hop1_floats = tf.convert_to_tensor(hop1_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * 4]
hop1_ints = np.array([[i, 2 * i] * hops[0] for i in range(7)])
hop1_ints = tf.convert_to_tensor(hop1_ints, dtype=tf.int64)
neighbor_hop_1 = Vertex(floats=hop1_floats, ints=hop1_ints)
# [f_num, batch_size] = [3, 2 * 4 * 5]
hop2_floats = np.array([[1.0 * i, 2.0 * i] * hops[0] * hops[1]
for i in range(3)])
hop2_floats = tf.convert_to_tensor(hop2_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * 4 * 5]
hop2_ints = np.array([[i, 2 * i] * hops[0] * hops[1]
for i in range(7)])
hop2_ints = tf.convert_to_tensor(hop2_ints, dtype=tf.int64)
neighbor_hop_2 = Vertex(floats=hop2_floats, ints=hop2_ints)
g = EgoGraph(vertices, [neighbor_hop_1, neighbor_hop_2], schema, hops)
dims = np.array([total_dim, 16, 8])
model = HomoEgoGIN(dims,
num_head=5,
bn_fn=None,
active_fn=tf.nn.relu,
droput=0.1)
embeddings = model.forward(g)
with tf.Session() as sess:
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
ret = sess.run(embeddings)
self.assertListEqual([2, 8],
list(ret.shape)) # [batch_size, output_dim]
def get_model_and_graph(self):
spec = FeatureSpec(10)
for i in range(3):
spec.append_dense()
total_dim = 3
for i in range(7):
dim = random.randint(8, 10)
spec.append_sparse(20 + 10 * i, dim, False)
total_dim += dim
hops = [4, 5]
# the centric vertices share the same spec with 2-hop neighbors
schema = [("nodes", spec), ("nodes", spec), ("nodes", spec)]
# [f_num, batch_size] = [3, 2]
batch_floats = np.array([[1.0 * i, 2.0 * i] for i in range(3)])
batch_floats = tf.convert_to_tensor(batch_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2]
batch_ints = np.array([[i, 2 * i] for i in range(7)])
batch_ints = tf.convert_to_tensor(batch_ints, dtype=tf.int64)
# [batch_size] = [2]
batch_labels = np.array([1, 0])
batch_labels = tf.convert_to_tensor(batch_labels, dtype=tf.int32)
vertices = Vertex(floats=batch_floats,
ints=batch_ints,
labels=batch_labels)
# [f_num, batch_size] = [3, 2 * 4]
hop1_floats = np.array([[1.0 * i, 2.0 * i] * hops[0]
for i in range(3)])
hop1_floats = tf.convert_to_tensor(hop1_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * 4]
hop1_ints = np.array([[i, 2 * i] * hops[0] for i in range(7)])
hop1_ints = tf.convert_to_tensor(hop1_ints, dtype=tf.int64)
neighbor_hop_1 = Vertex(floats=hop1_floats, ints=hop1_ints)
# [f_num, batch_size] = [3, 2 * 4 * 5]
hop2_floats = np.array([[1.0 * i, 2.0 * i] * hops[0] * hops[1]
for i in range(3)])
hop2_floats = tf.convert_to_tensor(hop2_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * 4 * 5]
hop2_ints = np.array([[i, 2 * i] * hops[0] * hops[1]
for i in range(7)])
hop2_ints = tf.convert_to_tensor(hop2_ints, dtype=tf.int64)
neighbor_hop_2 = Vertex(floats=hop2_floats, ints=hop2_ints)
g = EgoGraph(vertices, [neighbor_hop_1, neighbor_hop_2], schema, hops)
dims = np.array([total_dim, 16, 8])
model = HomoEgoGraphSAGE(dims,
bn_fn=None,
active_fn=tf.nn.relu,
droput=0.1)
return model, g, 2 # batch_size
def get_graph(cls, hops, neg=None):
spec = FeatureSpec(10)
for i in range(3):
spec.append_dense()
total_dim = 3
for i in range(7):
dim = i + 1
spec.append_sparse(20 + 10 * i, dim, False)
total_dim += dim
neg = 1 if not neg else int(neg)
hops[0] = int(hops[0] * neg)
# the centric vertices share the same spec with 2-hop neighbors
schema = [("nodes", spec), ("nodes", spec), ("nodes", spec)]
# batch_size = 2
# [f_num, batch_size] = [3, 2 * neg]
batch_floats = np.array([[1.0 * i, 2.0 * i] * neg for i in range(3)])
batch_floats = tf.convert_to_tensor(batch_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * neg]
batch_ints = np.array([[i, 2 * i] * neg for i in range(7)])
batch_ints = tf.convert_to_tensor(batch_ints, dtype=tf.int64)
# [batch_size] = [2]
batch_labels = np.array([1, 0])
batch_labels = tf.convert_to_tensor(batch_labels, dtype=tf.int32)
vertices = Vertex(floats=batch_floats,
ints=batch_ints,
labels=batch_labels)
# [f_num, batch_size] = [3, 2 * neg * hop0]
hop1_floats = np.array([[1.0 * i, 2.0 * i] * hops[0]
for i in range(3)])
hop1_floats = tf.convert_to_tensor(hop1_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * neg * hop0]
hop1_ints = np.array([[i, 2 * i] * hops[0] for i in range(7)])
hop1_ints = tf.convert_to_tensor(hop1_ints, dtype=tf.int64)
neighbor_hop_1 = Vertex(floats=hop1_floats, ints=hop1_ints)
# [f_num, batch_size] = [3, 2 * neg * hop0 * hop1]
hop2_floats = np.array([[1.0 * i, 2.0 * i] * hops[0] * hops[1]
for i in range(3)])
hop2_floats = tf.convert_to_tensor(hop2_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * neg * hop0 * hop1]
hop2_ints = np.array([[i, 2 * i] * hops[0] * hops[1]
for i in range(7)])
hop2_ints = tf.convert_to_tensor(hop2_ints, dtype=tf.int64)
neighbor_hop_2 = Vertex(floats=hop2_floats, ints=hop2_ints)
hops[0] = int(hops[0] / neg)
g = EgoGraph(vertices, [neighbor_hop_1, neighbor_hop_2], schema, hops)
return g, total_dim
def test_homogeneous_graph(self):
spec = FeatureSpec(10)
for i in range(3):
spec.append_dense()
total_dim = 3
for i in range(7):
dim = random.randint(8, 10)
spec.append_sparse(20 + 10 * i, dim, False)
total_dim += dim
hops = [4, 5]
# the centric vertices share the same spec with 2-hop neighbors
schema = [("nodes", spec), ("nodes", spec), ("nodes", spec)]
# [f_num, batch_size] = [3, 2]
batch_floats = np.array([[1.0 * i, 2.0 * i] for i in range(3)])
batch_floats = tf.convert_to_tensor(batch_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2]
batch_ints = np.array([[i, 2 * i] for i in range(7)])
batch_ints = tf.convert_to_tensor(batch_ints, dtype=tf.int64)
vertices = Vertex(floats=batch_floats, ints=batch_ints)
# [f_num, batch_size] = [3, 2 * 4]
hop1_floats = np.array([[1.0 * i, 2.0 * i] * hops[0] for i in range(3)])
hop1_floats = tf.convert_to_tensor(hop1_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * 4]
hop1_ints = np.array([[i, 2 * i] * hops[0] for i in range(7)])
hop1_ints = tf.convert_to_tensor(hop1_ints, dtype=tf.int64)
neighbor_hop_1 = Vertex(floats=hop1_floats, ints=hop1_ints)
# [f_num, batch_size] = [3, 2 * 4 * 5]
hop2_floats = np.array([[1.0 * i, 2.0 * i] * hops[0] * hops[1] for i in range(3)])
hop2_floats = tf.convert_to_tensor(hop2_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2 * 4 * 5]
hop2_ints = np.array([[i, 2 * i] * hops[0] * hops[1] for i in range(7)])
hop2_ints = tf.convert_to_tensor(hop2_ints, dtype=tf.int64)
neighbor_hop_2 = Vertex(floats=hop2_floats, ints=hop2_ints)
g1 = EgoGraph(vertices, [neighbor_hop_1, neighbor_hop_2], schema, hops)
g2 = g1.forward()
self.assertListEqual(list(g2.expands), hops)
x_list = [g2.nodes, g2.hop(0), g2.hop(1)]
with tf.Session() as sess:
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
ret = sess.run(x_list)
self.assertListEqual(list(ret[0].shape), [2, total_dim])
self.assertListEqual(list(ret[1].shape), [2 * hops[0], total_dim])
self.assertListEqual(list(ret[2].shape), [2 * hops[0] * hops[1], total_dim])
def forward(self, shared_param=None):
"""feature encoding function, which returns
encoded feature embedding as nodes with shape [batch_size, dim].
"""
def tf_transpose(item):
if item is not None:
item = tf.transpose(item)
return item
if self.node_schema is None:
return self
vertex_handler = FeatureHandler(self.node_schema[0],
self.node_schema[1].feature_spec)
ints, floats, strings =\
tf_transpose(self.nodes.int_attrs),\
tf_transpose(self.nodes.float_attrs),\
tf_transpose(self.nodes.string_attrs)
node = Vertex(self.nodes.ids, ints, floats, strings)
node_tensor = vertex_handler.forward(node)
graph = BatchGraph(self.edge_index,
node_tensor,
self.node_schema,
self.node_graph_id,
additional_keys=self.additional_keys)
for key in self.additional_keys:
graph[key] = self[key]
return graph
def build_node_from_tensors(feature_schema, tensors):
"""Constructs Vertices in Tensor format.
Args:
feature_schema: A (name, Decoder) tuple used to parse the feature.
Returns:
Vertex.
"""
ids = next(tensors)
if feature_schema[1].int_attr_num > 0:
int_attrs = next(tensors)
else:
int_attrs = None
if feature_schema[1].float_attr_num > 0:
float_attrs = next(tensors)
else:
float_attrs = None
if feature_schema[1].string_attr_num > 0:
string_attrs = next(tensors)
else:
string_attrs = None
feature_tensor = Vertex(ids,
ints=int_attrs,
floats=float_attrs,
strings=string_attrs)
return feature_tensor
def _build_feature(cls, graphs, type='node'):
def list_append(list, item):
if item is not None:
list.append(item)
return list
def np_concat(list):
if list:
return np.concatenate(list, axis=0)
return None
#TODO(baole): support other labels and weights.
ids_list = []
int_attrs_list = []
float_attrs_list = []
string_attrs_list = []
for graph in graphs:
if type == 'node':
item = graph.nodes
else:
return None
# flatten format.
ids_list = list_append(ids_list, item.ids)
int_attrs_list = list_append(int_attrs_list, item.int_attrs)
float_attrs_list = list_append(float_attrs_list, item.float_attrs)
string_attrs_list = list_append(string_attrs_list,
item.string_attrs)
ids = np_concat(ids_list)
ints = np_concat(int_attrs_list)
floats = np_concat(float_attrs_list)
strings = np_concat(string_attrs_list)
return Vertex(ids, ints, floats, strings)
def get_sub_graph_data(self):
spec = FeatureSpec(10)
for i in range(3):
spec.append_dense()
total_dim = 3
for i in range(7):
dim = random.randint(8, 10)
spec.append_sparse(20 + 10 * i, dim, False)
total_dim += dim
N = 10
# [f_num, batch_size] = [3, N]
batch_floats = np.random.random([3, N])
batch_floats = tf.convert_to_tensor(batch_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, N]
batch_ints = np.array([[i * j for j in range(N)] for i in range(7)])
batch_ints = tf.convert_to_tensor(batch_ints, dtype=tf.int64)
vertices = Vertex(floats=batch_floats, ints=batch_ints)
adj = np.zeros([N, N], dtype=np.float32)
for i in range(N):
for j in range(N):
adj[i][j] = i
tf_adj = tf.convert_to_tensor(adj, dtype=tf.float32)
g = SubGraph(vertices, tf_adj, schema=("nodes", spec))
return g, N, total_dim
def test_popular_cases_of_tensor(self):
spec = FeatureSpec(10)
for i in range(3):
spec.append_dense()
total_dim = 3
for i in range(7):
dim = random.randint(8, 10)
spec.append_sparse(20 + 10 * i, dim, False)
total_dim += dim
handler = FeatureHandler("popular_tensor", spec)
self.assertEqual(len(handler._float_fg), 3)
self.assertEqual(len(handler._int_fg), 0)
self.assertEqual(len(handler._fused_int_fg) > 0, True)
self.assertEqual(len(handler._string_fg), 0)
# input_shape = [f_num, batch_size]
batch1_floats = np.array([1.0 * i for i in range(3)], dtype=np.float32)
batch2_floats = np.array([[1.0 * i, 2.0 * i] for i in range(3)],
dtype=np.float32)
batch1_ints = np.array([i for i in range(7)])
batch2_ints = np.array([[i, 2 * i] for i in range(7)])
input1 = Vertex(floats=tf.convert_to_tensor(batch1_floats,
dtype=tf.float32),
ints=tf.convert_to_tensor(batch1_ints, dtype=tf.int64))
input2 = Vertex(floats=tf.convert_to_tensor(batch2_floats,
dtype=tf.float32),
ints=tf.convert_to_tensor(batch2_ints, dtype=tf.int64))
# output_shape = [batch_size, f_num]
output1 = handler(input1)
output2 = handler(input2)
with tf.Session() as sess:
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
ret1 = sess.run(output1)
ret2 = sess.run(output2)
self.assertListEqual(list(ret1.shape), [total_dim])
self.assertListEqual(list(ret2.shape), [2, total_dim])
def test_only_ints_with_fusion(self):
spec = FeatureSpec(10)
total_dim = 0
for i in range(8):
dim = random.randint(8, 10)
spec.append_sparse(100 + 10 * i, dim, False)
total_dim += dim
spec.append_sparse(100, 4, True) # two features need hash
spec.append_sparse(100, 4, True)
total_dim += 8
handler = FeatureHandler("ints_with_fusion", spec, fuse_embedding=True)
self.assertEqual(len(handler._float_fg), 0)
self.assertEqual(len(handler._int_fg) < 10, True)
self.assertEqual(len(handler._fused_int_fg) > 0, True)
self.assertEqual(len(handler._string_fg), 0)
# input_shape = [f_num, batch_size]
batch1_ints = np.array([i for i in range(10)])
batch2_ints = np.array([[i, 2 * i] for i in range(10)])
input1 = Vertex(ints=batch1_ints)
input2 = Vertex(ints=batch2_ints)
# output_shape = [batch_size, total_dim]
output1 = handler(input1)
output2 = handler(input2)
with tf.Session() as sess:
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
ret1 = sess.run(output1)
ret2 = sess.run(output2)
self.assertListEqual(list(ret1.shape),
[total_dim]) # 1d array, batch_size = 1
self.assertListEqual(list(ret2.shape),
[2, total_dim]) # 2d array, batch_size = 2
def test_only_strings(self):
spec = FeatureSpec(3)
total_dim = 0
for i in range(3):
dim = random.randint(8, 10)
spec.append_multival(10 + 10 * i, dim, ",")
total_dim += dim
handler = FeatureHandler("only_strings", spec)
self.assertEqual(len(handler._float_fg), 0)
self.assertEqual(len(handler._int_fg), 0)
self.assertEqual(len(handler._fused_int_fg), 0)
self.assertEqual(len(handler._string_fg), 3)
# input_shape = [f_num, batch_size]
batch1_ss = np.array(
["just,a,test", "items,splited,by,comma", "the,third,feature"])
batch2_ss = np.array([["f1,batch1", "f1,batch2,others"],
["f2,batch1,others", "f2,batch2"],
["f3,batch1,haha", "f3,batch2,hh,kk"]])
input1 = Vertex(strings=batch1_ss)
input2 = Vertex(strings=batch2_ss)
# output_shape = [batch_size, total_dim]
output1 = handler(input1)
output2 = handler(input2)
with tf.Session() as sess:
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
ret1 = sess.run(output1)
ret2 = sess.run(output2)
self.assertListEqual(list(ret1.shape),
[total_dim]) # 1d array, batch_size = 1
self.assertListEqual(list(ret2.shape),
[2, total_dim]) # 2d array, batch_size = 2
def test_only_floats(self):
spec = FeatureSpec(10)
for i in range(10):
spec.append_dense()
handler = FeatureHandler("only_floats", spec)
self.assertEqual(len(handler._float_fg), 10)
self.assertEqual(len(handler._int_fg), 0)
self.assertEqual(len(handler._fused_int_fg), 0)
self.assertEqual(len(handler._string_fg), 0)
# input_shape = [f_num, batch_size]
batch1_floats = np.array([1.0 * i for i in range(10)])
batch2_floats = np.array([[1.0 * i, 2.0 * i] for i in range(10)])
input1 = Vertex(floats=batch1_floats)
input2 = Vertex(floats=batch2_floats)
# output_shape = [batch_size, f_num]
output1 = handler(input1)
output2 = handler(input2)
with tf.Session() as sess:
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
ret1 = sess.run(output1)
ret2 = sess.run(output2)
self.assertListEqual(list(ret1.shape),
list(batch1_floats.transpose().shape))
self.assertListEqual(list(ret1), list(batch1_floats.transpose()))
self.assertListEqual(list(ret2.shape),
list(batch2_floats.transpose().shape))
input2_t = batch2_floats.transpose()
for i in range(ret2.shape[0]):
self.assertListEqual(list(ret2[i]), list(input2_t[i]))
def _build_entity(self, value):
entities = {}
cursor = [-1]
def pop(mask):
if mask:
cursor[0] += 1
return value[cursor[0]]
return None
for alias, handler in self._handler_dict.items():
_, _, feat_masks, id_masks, degree_masks = handler
ints, floats, strings, labels, weights, ids, _, out_degrees, in_degrees = \
[pop(msk) for msk in feat_masks + id_masks + degree_masks]
# TODO(wenting.swt): To support edges.
entities[alias] = Vertex(ids, ints, floats, strings, labels,
weights, out_degrees, in_degrees)
return entities
def test_heterogeneous_graph(self):
u_spec = FeatureSpec(10)
for i in range(3):
u_spec.append_dense()
u_total_dim = 3
for i in range(7):
dim = random.randint(8, 10)
u_spec.append_sparse(20 + 10 * i, dim, False)
u_total_dim += dim
i_spec = FeatureSpec(19)
for i in range(6):
i_spec.append_dense()
i_total_dim = 6
for i in range(13):
dim = random.randint(8, 11)
i_spec.append_sparse(30 + 10 * i, dim, False)
i_total_dim += dim
u_out_dim = 16
i_out_dim = 12
out_dim = 9
hops = [4, 5]
# the centric vertices share the same spec with 2-hop neighbors
# metapath: u--i--i
schema = [("u_nodes", u_spec), ("nbr", i_spec), ("nbr", i_spec)]
# [f_num, batch_size] = [3, 2]
batch_floats = np.array([[1.0 * i, 2.0 * i] for i in range(3)])
batch_floats = tf.convert_to_tensor(batch_floats, dtype=tf.float32)
# [i_num, batch_size] = [7, 2]
batch_ints = np.array([[i, 2 * i] for i in range(7)])
batch_ints = tf.convert_to_tensor(batch_ints, dtype=tf.int64)
vertices = Vertex(floats=batch_floats, ints=batch_ints)
# [f_num, batch_size] = [6, 2 * 4]
hop1_floats = np.array([[1.0 * i, 2.0 * i] * hops[0]
for i in range(6)])
hop1_floats = tf.convert_to_tensor(hop1_floats, dtype=tf.float32)
# [i_num, batch_size] = [13, 2 * 4]
hop1_ints = np.array([[i, 2 * i] * hops[0] for i in range(13)])
hop1_ints = tf.convert_to_tensor(hop1_ints, dtype=tf.int64)
neighbor_hop_1 = Vertex(floats=hop1_floats, ints=hop1_ints)
# [f_num, batch_size] = [6, 2 * 4 * 5]
hop2_floats = np.array([[1.0 * i, 2.0 * i] * hops[0] * hops[1]
for i in range(6)])
hop2_floats = tf.convert_to_tensor(hop2_floats, dtype=tf.float32)
# [i_num, batch_size] = [13, 2 * 4 * 5]
hop2_ints = np.array([[i, 2 * i] * hops[0] * hops[1]
for i in range(13)])
hop2_ints = tf.convert_to_tensor(hop2_ints, dtype=tf.int64)
neighbor_hop_2 = Vertex(floats=hop2_floats, ints=hop2_ints)
g = EgoGraph(vertices, [neighbor_hop_1, neighbor_hop_2], schema, hops)
layer_ui = EgoSAGELayer("heter_uv",
input_dim=(u_total_dim, i_total_dim),
output_dim=u_out_dim,
agg_type="mean",
com_type="concat")
layer_ii = EgoSAGELayer("heter_ii",
input_dim=i_total_dim,
output_dim=i_out_dim,
agg_type="mean",
com_type="concat")
layer_uii = EgoSAGELayer("heter_uii",
input_dim=(u_out_dim, i_out_dim),
output_dim=out_dim,
agg_type="sum",
com_type="concat")
layer_group_1 = EgoSAGELayerGroup([layer_ui, layer_ii])
layer_group_2 = EgoSAGELayerGroup([layer_uii])
model = EgoGraphSAGE([layer_group_1, layer_group_2],
bn_fn=None,
active_fn=tf.nn.relu,
droput=0.1)
embeddings = model.forward(g)
with tf.Session() as sess:
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
ret = sess.run(embeddings)
self.assertListEqual([2, 9],
list(ret.shape)) # [batch_size, output_dim]