def _preprocessing(self, graph, x, drop_edge=False): device = x.device graph.to("cpu") x = x.to("cpu") graph.eval() op_embedding = [x] edge_index = graph.edge_index if self.undirected: edge_index = to_undirected(edge_index) if drop_edge: edge_index, _ = dropout_adj(edge_index, drop_rate=self.dropedge_rate) graph = get_adj(graph, remove_diag=self.remove_diag) for norm in self.adj_norm: with graph.local_graph(): graph.edge_index = edge_index graph.normalize(norm) if self.diffusion == "ppr": results = multi_hop_ppr_diffusion(graph, graph.x, self.num_propagations) else: results = multi_hop_sgc(graph, graph.x, self.num_propagations) op_embedding.extend(results) graph.to(device) return torch.cat(op_embedding, dim=1).to(device)
def _preprocessing(self, graph, x): op_embedding = [] op_embedding.append(x) edge_index = graph.edge_index # Convert to numpy arrays on cpu edge_index, _ = dropout_adj(edge_index, drop_rate=self.dropedge_rate) # if self.undirected: # edge_index = to_undirected(edge_index, num_nodes) graph = get_adj(graph, asymm_norm=self.asymm_norm, set_diag=self.set_diag, remove_diag=self.remove_diag) with graph.local_graph(): graph.edge_index = edge_index for _ in range(self.num_propagations): x = spmm(graph, x) op_embedding.append(x) for _ in range(self.num_propagations): nx = spmm(graph, x) op_embedding.append(nx) return torch.cat(op_embedding, dim=1)
def forward(self, graph: Graph) -> torch.Tensor: x = graph.x if self.improved and not hasattr(graph, "unet_improved"): row, col = graph.edge_index row = torch.cat( [row, torch.arange(0, x.shape[0], device=x.device)], dim=0) col = torch.cat( [col, torch.arange(0, x.shape[0], device=x.device)], dim=0) graph.edge_index = (row, col) graph["unet_improved"] = True graph.row_norm() with graph.local_graph(): if self.training and self.adj_dropout > 0: graph.edge_index, graph.edge_weight = dropout_adj( graph.edge_index, graph.edge_weight, self.adj_dropout) x = F.dropout(x, p=self.n_dropout, training=self.training) h = self.in_gcn(graph, x) h = self.act(h) h_list = self.unet(graph, h) h = h_list[-1] h = F.dropout(h, p=self.n_dropout, training=self.training) return self.out_gcn(graph, h)
def forward(self, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: Optional[torch.Tensor] = None) -> torch.Tensor: if self.cache_edge_attr is None: edge_index, _ = add_remaining_self_loops(edge_index) if self.improved: self_loop = torch.stack([torch.arange(0, x.shape[0])] * 2, dim=0).to(x.device) edge_index = torch.cat([edge_index, self_loop], dim=1) edge_attr = row_normalization(x.shape[0], edge_index) self.cache_edge_attr = edge_attr self.cache_edge_index = edge_index else: edge_index = self.cache_edge_index edge_attr = self.cache_edge_attr if self.training and self.adj_dropout > 0: edge_index, edge_attr = dropout_adj(edge_index, edge_attr, self.adj_dropout) x = F.dropout(x, p=self.n_dropout, training=self.training) h = self.in_gcn(x, edge_index, edge_attr) h = self.act(h) h_list = self.unet(h, edge_index, edge_attr) h = h_list[-1] h = F.dropout(h, p=self.n_dropout, training=self.training) return self.out_gcn(h, edge_index, edge_attr)
def prop( self, graph: Graph, x: torch.Tensor, drop_feature_rate: float = 0.0, drop_edge_rate: float = 0.0, ): x = dropout_features(x, drop_feature_rate) with graph.local_graph(): graph.edge_index, graph.edge_weight = dropout_adj( graph.edge_index, graph.edge_weight, drop_edge_rate) return self.model.forward(graph, x)
def forward(self, graph: Graph) -> torch.Tensor: x = graph.x if self.improved and not hasattr(graph, "unet_improved"): self_loop = torch.stack([torch.arange(0, x.shape[0])] * 2, dim=0).to(x.device) graph.edge_index = torch.cat([graph.edge_index, self_loop], dim=1) graph["unet_improved"] = True graph.row_norm() with graph.local_graph(): if self.training and self.adj_dropout > 0: graph.edge_index, graph.edge_weight = dropout_adj(graph.edge_index, graph.edge_weight, self.adj_dropout) x = F.dropout(x, p=self.n_dropout, training=self.training) h = self.in_gcn(graph, x) h = self.act(h) h_list = self.unet(graph, h) h = h_list[-1] h = F.dropout(h, p=self.n_dropout, training=self.training) return self.out_gcn(graph, h)
def _preprocessing(self, x, edge_index): num_nodes = x.shape[0] op_embedding = [] op_embedding.append(x) # Convert to numpy arrays on cpu edge_index, _ = dropout_adj(edge_index, drop_rate=self.dropedge_rate) row, col = edge_index if self.undirected: edge_index = to_undirected(edge_index, num_nodes) row, col = edge_index # adj matrix edge_index, edge_attr = get_adj(row, col, asymm_norm=self.asymm_norm, set_diag=self.set_diag, remove_diag=self.remove_diag) nx = x for _ in range(self.num_propagations): nx = spmm(edge_index, edge_attr, nx) op_embedding.append(nx) # transpose adj matrix edge_index, edge_attr = get_adj(col, row, asymm_norm=self.asymm_norm, set_diag=self.set_diag, remove_diag=self.remove_diag) nx = x for _ in range(self.num_propagations): nx = spmm(edge_index, edge_attr, nx) op_embedding.append(nx) return torch.cat(op_embedding, dim=1)
def forward(self, graph): graph.requires_grad = False edge_index, edge_weight = dropout_adj( graph.edge_index, drop_rate=self.drop_edge_rate, renorm=None, training=self.training ) h = graph.x h = F.dropout(h, self.dropout, training=self.training) with graph.local_graph(): graph.edge_index = edge_index graph.sym_norm() assert (graph.degrees() > 0).all() h = self.layers[0](graph, h) mask = shared_dropout(h, self.dropout) for i in range(1, len(self.layers) - 1): h = self.layers[i](graph, h, mask) h = self.norm(h) h = self.act(h) h = F.dropout(h, p=self.dropout, training=self.training) h = self.layers[-1](graph, h) return h
def transform_data(self): self.graph.edge_index, _ = dropout_adj(edge_index=self.edge_index, drop_rate=self.dropedge_rate) return self.graph.to(self.device)