def __call__(self, edge_index, x=None, edge_weight=None): if self.method_type == "emb": if isinstance(edge_index, np.ndarray): edge_index = torch.from_numpy(edge_index) edge_index = (edge_index[:, 0], edge_index[:, 1]) data = Graph(edge_index=edge_index, edge_weight=edge_weight) self.model = build_model(self.args) embeddings = self.model(data) elif self.method_type == "gnn": num_nodes = edge_index.max().item() + 1 if x is None: print("No input node features, use random features instead.") x = np.random.randn(num_nodes, self.num_features) if isinstance(x, np.ndarray): x = torch.from_numpy(x).float() if isinstance(edge_index, np.ndarray): edge_index = torch.from_numpy(edge_index) edge_index = (edge_index[:, 0], edge_index[:, 1]) data = Graph(x=x, edge_index=edge_index, edge_weight=edge_weight) torch.save(data, self.data_path) dataset = NodeDataset(path=self.data_path, scale_feat=False, metric="accuracy") self.args.dataset = dataset model = train(self.args) embeddings = model.embed(data.to(model.device)) embeddings = embeddings.detach().cpu().numpy() return embeddings
def build_g_feat(self, A): edge2type = {} edges = [] weights = [] for k, mat in enumerate(A): edges.append(mat[0].cpu().numpy()) weights.append(mat[1].cpu().numpy()) for u, v in zip(*edges[-1]): edge2type[(u, v)] = k edges = np.concatenate(edges, axis=1) weights = np.concatenate(weights) edges = torch.tensor(edges).to(self.device) weights = torch.tensor(weights).to(self.device) g = Graph(edge_index=edges, edge_weight=weights) g = g.to(self.device) e_feat = [] for u, v in zip(*g.edge_index): u = u.cpu().item() v = v.cpu().item() e_feat.append(edge2type[(u, v)]) e_feat = torch.tensor(e_feat, dtype=torch.long).to(self.device) g.edge_type = e_feat self.g = g