(self, dims, metapath, label_dim,
feature_idx, feature_max_id,
processing_steps=4,
lstm_layers=2)
| 53 | class Set2SetModel(GraphModel): |
| 54 | |
| 55 | def __init__(self, dims, metapath, label_dim, |
| 56 | feature_idx, feature_max_id, |
| 57 | processing_steps=4, |
| 58 | lstm_layers=2): |
| 59 | super(Set2SetModel, self).__init__(label_dim) |
| 60 | self.gnn = GNN('sage', 'full', dims, None, metapath, |
| 61 | feature_idx, feature_max_id) |
| 62 | self.pool = tf_euler.graph_pool.Set2SetPool( |
| 63 | dims[-1], |
| 64 | processing_steps=processing_steps, |
| 65 | num_layers=lstm_layers) |
| 66 | |
| 67 | def embed(self, n_id, graph_index): |
| 68 | node_emb = self.gnn(n_id) |