Ejemplo n.º 1
0
    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]
Ejemplo n.º 2
0
    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
Ejemplo n.º 3
0
    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
Ejemplo n.º 4
0
  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])
Ejemplo n.º 5
0
    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
Ejemplo n.º 6
0
   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
Ejemplo n.º 7
0
    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)
Ejemplo n.º 8
0
    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
Ejemplo n.º 9
0
    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])
Ejemplo n.º 10
0
    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
Ejemplo n.º 11
0
    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
Ejemplo n.º 12
0
    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]))
Ejemplo n.º 13
0
    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
Ejemplo n.º 14
0
    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]