def message(self, x_j, edge_index, size2): return x_j # x_j has shape [E, out_channels] row, col = edge_index deg = pyg_utils.degree(row, size2[0], dtype=x_j.dtype) deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] return norm.view(-1, 1) * x_j
def forward(self, x, edge_index): x = self.fc(x) row, col = edge_index deg = degree(row, x.size(0), dtype = x.dtype) + 1 deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] return self.propagate(edge_index, x=x, norm=norm)
def aggregate(self, inputs: Tensor, index: Tensor, dim_size: Optional[int] = None) -> Tensor: outs = [aggr(inputs, index, dim_size) for aggr in self.aggregators] out = torch.cat(outs, dim=-1) deg = degree(index, dim_size, dtype=inputs.dtype).view(-1, 1, 1) outs = [scaler(out, deg, self.avg_deg) for scaler in self.scalers] return torch.cat(outs, dim=-1)
def message(self, x_j, edge_index, size): # x_j has shape [E, out_channels] # Step 3: Normalize node features. src, dst = edge_index # we assume source_to_target message passing deg = degree(src, size[0], dtype=x_j.dtype) deg = deg.pow(-1) norm = deg[dst] return norm.view(-1, 1) * x_j # broadcasting the normalization term to all out_channels === hidden features
def message(self, x_i, x_j, edge_index, size): # Compute messages # x_j has shape [E, out_channels] row, col = edge_index deg = pyg_utils.degree(row, size[0], dtype=x_j.dtype) deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] return x_j
def __call__(self, data): data.x = torch.zeros((data.num_nodes, 1), dtype=torch.float) data = TwoMalkin()(data) data = ConnectedThreeMalkin()(data) data.x = degree(data.edge_index[0], data.num_nodes, dtype=torch.long) data.x = F.one_hot( data.x // self.div, num_classes=(degrees[self.args.dataset]) // self.div + 1).to( torch.float) return data
def message(self, x_j, edge_index, size): # x_j has shape [E, out_channels] # Step 3: Normalize node features. row, col = edge_index deg = degree(row, size[0], dtype=x_j.dtype) deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] return norm.view(-1, 1) * x_j
def __call__(self, data): col, x = data.edge_index[1], data.x deg = degree(col, data.num_nodes) #print(deg) if self.norm: deg = deg / (deg.max() if self.max is None else self.max) deg = deg.view(-1, 1) if x is None: data.x = deg return data
def forward(self, x, edge_index, edge_attr): edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0)) self_loop_attr = torch.zeros((x.size(0), edge_attr.size(1)), dtype=edge_attr.dtype).to(x.device) edge_attr = torch.cat((edge_attr, self_loop_attr), dim=0) row, col = edge_index deg = degree(col, x.size(0), dtype=x.dtype) deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] return self.propagate(edge_index, x=x, edge_attr=edge_attr, norm=norm)
def forward(self, data): degrees = degree(data.edge_index[0], num_nodes=data.num_nodes, dtype=torch.float) density = torch.mm( degrees.view(1, -1), sparse_to_dense(data.edge_index, num_nodes=data.num_nodes)).flatten() return (degrees * 1000000 + density)
def get_max_deg(dataset): max_deg = 0 for data in dataset: row, col = data.edge_index num_nodes = data.num_nodes deg = degree(row, num_nodes) deg = max(deg).item() if deg > max_deg: max_deg = int(deg) return max_deg
def forward(self, x, edge_index, edge_weight=None): """""" edge_index, edge_weight = remove_self_loops(edge_index, edge_weight) row, col = edge_index num_nodes, num_edges, order_filter = x.size(0), row.size(0), self.weight.size(0) if edge_weight is None: edge_weight = x.new_ones((num_edges,)) edge_weight = edge_weight.view(-1) assert edge_weight.size(0) == edge_index.size(1) deg = degree(row, num_nodes, dtype=x.dtype) # Compute normalized and rescaled Laplacian. deg = deg.pow(-0.5) deg[deg == float('inf')] = 0 lap = -deg[row] * edge_weight * deg[col] def weight_mult(x, w): y = torch.einsum('fgrs,ifrs->igrs', w, x) return y def lap_mult(edge_index, lap, x): L = torch.sparse.IntTensor(edge_index, lap, torch.Size([x.shape[0], x.shape[0]])).to_dense() x_tilde = torch.einsum('ij,ifrs->jfrs', L, x) return x_tilde # Perform filter operation recurrently. horizon = x.shape[1] x = x.permute(0, 2, 1) x_hat = torch.rfft(x, 1, normalized=True, onesided=True) Tx_0 = x_hat y_hat = weight_mult(Tx_0, self.weight[0, :]) if order_filter > 1: Tx_1 = lap_mult(edge_index, lap, x_hat) y_hat = y_hat + weight_mult(Tx_1, self.weight[1, :]) for k in range(2, order_filter): Tx_2 = 2 * lap_mult(edge_index, lap, Tx_1) - Tx_0 y_hat = y_hat + weight_mult(Tx_2, self.weight[k, :]) Tx_0, Tx_1 = Tx_1, Tx_2 y = torch.irfft(y_hat, 1, normalized=True, onesided=True, signal_sizes=(horizon,)) y = y.permute(0, 2, 1) if self.bias is not None: y = y + self.bias return y
def train_pyg(model_name, model, train_loader, device, optimizer): global epoch_load_time, epoch_forward_time, epoch_backward_time, epoch_batch_time, epoch_update_time t10 = time.time() model.train() n_data = 0 loss_all = 0 epoch_train_acc = 0 t4 = time.time() t2 = time.time() for data in train_loader: torch.cuda.synchronize() epoch_load_time = epoch_load_time + time.time() - t2 print("data load time:", time.time() - t2) t7 = time.time() data = data.to(device) torch.cuda.synchronize() print("data to gpu:", time.time() - t7) if model_name == 'monet': row, col = data.edge_index deg = degree(col, data.num_nodes) data.edge_attr = torch.stack([1 / torch.sqrt(deg[row]), 1 / torch.sqrt(deg[col])], dim=-1) optimizer.zero_grad() torch.cuda.synchronize() t9 = time.time() #print(data.edge_attr) output = model(data.x, data.edge_index, data.edge_attr, data.batch) torch.cuda.synchronize() epoch_forward_time = epoch_forward_time + time.time() - t9 print("forward:", time.time() - t9) loss = F.nll_loss(output, data.y) torch.cuda.synchronize() t3 = time.time() #loss.register_hook(print) loss.backward() torch.cuda.synchronize() epoch_backward_time = epoch_backward_time + time.time() - t3 print("backward:", time.time() - t3) t5 = time.perf_counter() optimizer.step() torch.cuda.synchronize() epoch_update_time = epoch_update_time + time.perf_counter() - t5 print("update:", time.perf_counter() - t5) loss_all += loss.item() * data.num_graphs n_data += data.y.size(0) epoch_train_acc += accuracy(output, data.y) torch.cuda.synchronize() epoch_batch_time = epoch_batch_time + time.time() -t7 print("a batch:", time.time() - t7) t2 = time.time() epoch_train_acc /= n_data print("train:", time.time() - t10) return loss_all / n_data, epoch_train_acc, optimizer
def __init__(self, net_path, device, max_iter): with open(net_path, "rb") as f: self.edge_list = pkl.load(f) # edge_list with size [3, E], (src_node, tar_node, weight) self.device = device self.src_nodes = T.LongTensor(self.edge_list[0]).to(device) self.tar_nodes = T.LongTensor(self.edge_list[1]).to(device) self.weights = T.FloatTensor(self.edge_list[2]).to(device) self.cave_index = T.LongTensor(self.edge_list[3]).to(device) self.N = max([T.max(self.src_nodes), T.max(self.tar_nodes)]).item()+1 self.E = len(self.src_nodes) self.d = degree(self.tar_nodes, num_nodes=self.N).to(device) self.out_d = degree(self.src_nodes, num_nodes=self.N).to(device) self.out_weight_d = scatter_add(self.weights, self.src_nodes).to(device) self.G = nx.DiGraph() self.G.add_edges_from(self.edge_list[:2].T) self.max_iter = max_iter
def get_dataset(name, sparse=True, cleaned=False): if name == 'node': path = osp.join(os.environ['GNN_TRAINING_DATA_ROOT'], name) print(path) dataset = HitGraphDataset2(path, directed=False, categorical=True) else: path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', name) dataset = TUDataset(path, name, cleaned=cleaned) dataset.data.edge_attr = None if dataset.data.x is None: max_degree = 0 degs = [] for data in dataset: degs += [degree(data.edge_index[0], dtype=torch.long)] max_degree = max(max_degree, degs[-1].max().item()) if max_degree < 1000: dataset.transform = T.OneHotDegree(max_degree) else: deg = torch.cat(degs, dim=0).to(torch.float) mean, std = deg.mean().item(), deg.std().item() dataset.transform = NormalizedDegree(mean, std) if not sparse: num_nodes = max_num_nodes = 0 for data in dataset: num_nodes += data.num_nodes max_num_nodes = max(data.num_nodes, max_num_nodes) # Filter out a few really large graphs in order to apply DiffPool. if name == 'REDDIT-BINARY': num_nodes = min(int(num_nodes / len(dataset) * 1.5), max_num_nodes) else: num_nodes = min(int(num_nodes / len(dataset) * 5), max_num_nodes) indices = [] for i, data in enumerate(dataset): if data.num_nodes <= num_nodes: indices.append(i) dataset = dataset[torch.tensor(indices)] if dataset.transform is None: dataset.transform = T.ToDense(num_nodes) else: dataset.transform = T.Compose( [dataset.transform, T.ToDense(num_nodes)]) return dataset
def forward(self, x, edge_index): edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0)) row, col = edge_index deg = degree(row, x.size(0), dtype=x.dtype) deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] x = self.lin(x) return self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x, norm=norm)
def message(self, x_i, x_j, edge_index, size): if self.cached_norm is None: row, col = edge_index deg = pyg_utils.degree(row, size[0], dtype=x_j.dtype) deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] if self.cache: self.cached_norm = norm else: norm = self.cached_norm return norm.view(-1, 1) * x_j
def mask_features(x, edge_index, edge_weight, sigma): source, target = edge_index h_s, h_t = x[source], x[target] h = (h_t - h_s) / sigma h = edge_weight.view(-1, 1) * h * h mask = torch.zeros(x.size(), device=device) mask.index_add_(0, source, h) deg = degree(edge_index[0]) mask = torch.exp(- mask / deg.view(-1, 1)) x = x * mask return x
def aggregate(self, inputs, index, ptr=None, dim_size=None): if self.aggr in ['add', 'mean', 'max', None]: return super(GenMessagePassing, self).aggregate(inputs, index, ptr, dim_size) elif self.aggr in ['softmax_sg', 'softmax', 'softmax_sum']: if self.learn_t: out = scatter_softmax(inputs*self.t, index, dim=self.node_dim) else: with torch.no_grad(): out = scatter_softmax(inputs*self.t, index, dim=self.node_dim) out = scatter(inputs*out, index, dim=self.node_dim, dim_size=dim_size, reduce='sum') if self.aggr == 'softmax_sum': self.sigmoid_y = torch.sigmoid(self.y) degrees = degree(index, num_nodes=dim_size).unsqueeze(1) out = torch.pow(degrees, self.sigmoid_y) * out return out elif self.aggr in ['power', 'power_sum']: min_value, max_value = 1e-7, 1e1 torch.clamp_(inputs, min_value, max_value) out = scatter(torch.pow(inputs, self.p), index, dim=self.node_dim, dim_size=dim_size, reduce='mean') torch.clamp_(out, min_value, max_value) out = torch.pow(out, 1/self.p) if self.aggr == 'power_sum': self.sigmoid_y = torch.sigmoid(self.y) degrees = degree(index, num_nodes=dim_size).unsqueeze(1) out = torch.pow(degrees, self.sigmoid_y) * out return out else: raise NotImplementedError('To be implemented')
def compute_pr(edge_index, damp: float = 0.85, k: int = 10): num_nodes = edge_index.max().item() + 1 deg_out = degree(edge_index[0]) x = torch.ones((num_nodes, )).to(edge_index.device).to(torch.float32) for i in range(k): edge_msg = x[edge_index[0]] / deg_out[edge_index[0]] agg_msg = scatter(edge_msg, edge_index[1], reduce='sum') x = (1 - damp) * x + damp * agg_msg return x
def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor: """""" if batch is None: out = F.instance_norm( x.t().unsqueeze(0), self.running_mean, self.running_var, self.weight, self.bias, self.training or not self.track_running_stats, self.momentum, self.eps) return out.squeeze(0).t() batch_size = int(batch.max()) + 1 mean = var = unbiased_var = x # Dummies. if self.training or not self.track_running_stats: norm = degree(batch, batch_size, dtype=x.dtype).clamp_(min=1) norm = norm.view(-1, 1) unbiased_norm = (norm - 1).clamp_(min=1) mean = scatter(x, batch, dim=0, dim_size=batch_size, reduce='add') / norm x = x - mean[batch] var = scatter(x * x, batch, dim=0, dim_size=batch_size, reduce='add') unbiased_var = var / unbiased_norm var = var / norm momentum = self.momentum if self.running_mean is not None: self.running_mean = ( 1 - momentum) * self.running_mean + momentum * mean.mean(0) if self.running_var is not None: self.running_var = ( 1 - momentum ) * self.running_var + momentum * unbiased_var.mean(0) else: if self.running_mean is not None: mean = self.running_mean.view(1, -1).expand(batch_size, -1) if self.running_var is not None: var = self.running_var.view(1, -1).expand(batch_size, -1) x = x - mean[batch] out = x / (var + self.eps).sqrt()[batch] if self.weight is not None and self.bias is not None: out = out * self.weight.view(1, -1) + self.bias.view(1, -1) return out
def deg(self): # Compute in-degree histogram over training data. deg = torch.zeros(5, dtype=torch.long) if self.dataset_train is not None: for data in self.dataset_train: d = degree( data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long ) deg += torch.bincount(d, minlength=deg.numel()) else: deg = torch.tensor([[0, 41130, 117278, 70152, 3104]]) return deg.numpy().tolist()
def __call__(self, data): idx, x = data.edge_index[1 if self.in_degree else 0], data.x deg = truncate_degree(degree(idx, data.num_nodes, dtype=torch.long)) deg = F.one_hot(deg, num_classes=self.max_degree + 1).to(torch.float) if x is not None and self.cat: x = x.view(-1, 1) if x.dim() == 1 else x data.x = torch.cat([x, deg.to(x.dtype)], dim=-1) else: data.x = deg return data
def __call__(self, data): row, x = data.edge_index[0], data.x deg = degree(row, data.num_nodes) deg = one_hot(deg, num_classes=self.max_degree + 1) if x is not None and self.cat: x = x.view(-1, 1) if x.dim() == 1 else x data.x = torch.cat([x, deg.to(x.dtype)], dim=-1) else: data.x = deg return data
def batched_negative_sampling(edge_index, batch, num_neg_samples=None, method="sparse", force_undirected=False): r"""Samples random negative edges of multiple graphs given by :attr:`edge_index` and :attr:`batch`. Args: edge_index (LongTensor): The edge indices. batch (LongTensor): Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each node to a specific example. num_neg_samples (int, optional): The number of negative samples to return. If set to :obj:`None`, will try to return a negative edge for every positive edge. (default: :obj:`None`) method (string, optional): The method to use for negative sampling, *i.e.*, :obj:`"sparse"` or :obj:`"dense"`. This is a memory/runtime trade-off. :obj:`"sparse"` will work on any graph of any size, while :obj:`"dense"` can perform faster true-negative checks. (default: :obj:`"sparse"`) force_undirected (bool, optional): If set to :obj:`True`, sampled negative edges will be undirected. (default: :obj:`False`) :rtype: LongTensor """ split = degree(batch[edge_index[0]], dtype=torch.long).tolist() edge_indices = torch.split(edge_index, split, dim=1) num_nodes = degree(batch, dtype=torch.long) cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)[:-1]]) neg_edge_indices = [] for edge_index, N, C in zip(edge_indices, num_nodes.tolist(), cum_nodes.tolist()): neg_edge_index = negative_sampling(edge_index - C, N, num_neg_samples, method, force_undirected) + C neg_edge_indices.append(neg_edge_index) return torch.cat(neg_edge_indices, dim=1)
def __init__(self, data, state_feats=None, gamma=0.99, epsilon=lambda x: 0.5, episode_len=10, num_walks=10, hidden=64, network=None, edata=False, frozen=True): # Is there really not a better way to do this? self.max_actions = max( degree(data.edge_index[0]).max().long().item(), degree(data.edge_index[1]).max().long().item()) if not state_feats: state_feats = data.num_nodes self.state_feats = state_feats self.hidden = hidden if not network: network = Q_Network_No_Action(self.state_feats, self.max_actions, hidden) super().__init__(data, state_feats=state_feats, action_feats=0, gamma=gamma, epsilon=epsilon, episode_len=episode_len, num_walks=num_walks, hidden=hidden, network=network, edata=edata, frozen=True) self.action_map = None
def global_mean_pool_sparse(x, batch): #-------------- global average pooling index = torch.stack( [batch, torch.tensor(list(range(batch.shape[0])), device=x.device)], 0) x_sparse = torch.sparse.FloatTensor( index, x, torch.Size([torch.max(batch) + 1, x.shape[0], x.shape[1]])) graph_sizes = degree(batch).float() graph_sizes[graph_sizes == 0.0] = 1.0 return torch.sparse.sum(x_sparse, 1).to_dense() / graph_sizes.unsqueeze(1)
def forward(self, x: torch.Tensor, idx: torch.Tensor, dim_size: Optional[int] = None, dim: int = 0) -> torch.Tensor: outs = [aggr(x, idx, dim_size) for aggr in self.aggregators] # concatenate the different aggregator results, considering the shape of the hypercomplex components. out = phm_cat(tensors=outs, phm_dim=self.phm_dim, dim=-1) if self.scalers is not None: deg = degree(idx, dim_size, dtype=x.dtype).view(-1, 1) # concatenate the different aggregator results, considering the shape of the hypercomplex components. outs = [scaler(out, deg, self.avg_deg) for scaler in self.scalers] out = phm_cat(tensors=outs, phm_dim=self.phm_dim, dim=-1) out = self.transform(out) return out
def aggregate(self, inputs: torch.Tensor, index: torch.Tensor, dim_size: Optional[int] = None) -> torch.Tensor: # inputs has shape (*, self.phm_dim * self.in_feats) outs = [aggr(inputs, index, dim_size) for aggr in self.aggregators] # concatenate the different aggregator results, considering the shape of the hypercomplex components. out = phm_cat(tensors=outs, phm_dim=self.phm_dim, dim=-1) deg = degree(index, dim_size, dtype=inputs.dtype).view(-1, 1) # concatenate the different aggregator results, considering the shape of the hypercomplex components. outs = [scaler(out, deg, self.avg_deg) for scaler in self.scalers] out = phm_cat(tensors=outs, phm_dim=self.phm_dim, dim=-1) return out
def forward(self, x, edge_index, edge_attr): x = self.fc(x) edge_embedding = self.bond_encoder(edge_attr) row, col = edge_index deg = degree(row, x.size(0), dtype=x.dtype) + 1 deg_inv_sqrt = deg.pow(-0.5) norm = deg_inv_sqrt[row] * deg_inv_sqrt[col] return self.propagate( edge_index, x=x, edge_attr=edge_embedding, norm=norm ) + F.relu(x + self.root_emb.weight) * 1. / deg.view(-1, 1)