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
