def drop_adj(self, graph: Graph, drop_rate: float = 0.5): if drop_rate < 0.0 or drop_rate > 1.0: raise ValueError("Dropout probability has to be between 0 and 1, " "but got {}".format(drop_rate)) if not self.training: return graph num_edges = graph.num_edges mask = torch.full((num_edges,), 1 - drop_rate, dtype=torch.float) mask = torch.bernoulli(mask).to(torch.bool) edge_index = graph.edge_index edge_weight = graph.edge_weight graph.edge_index = edge_index[:, mask] graph.edge_weight = edge_weight[mask] return graph
def _construct_propagation_matrix(self, sample_adj, sample_id, num_neighbors): row, col = sample_adj.edge_index value = sample_adj.edge_weight """add self connection""" num_row = row.max() + 1 row = torch.cat([torch.arange(0, num_row).long(), row], dim=0) col = torch.cat([torch.arange(0, num_row).long(), col], dim=0) value = torch.cat([self.diag[sample_id[:num_row]], value], dim=0) value = value * self.degree[sample_id[row]] / num_neighbors new_graph = Graph() new_graph.edge_index = torch.stack([row, col]) new_graph.edge_weight = value return new_graph