class ASAP(torch.nn.Module): def __init__(self, num_features, num_classes, num_layers, hidden, ratio=0.8, dropout=0): super(ASAP, self).__init__() self.conv1 = ChebConv(num_features, hidden // 2, K=4, aggr='mean') self.conv2 = ChebConv(hidden // 2, hidden, K=4, aggr='mean') self.convs = torch.nn.ModuleList() self.pools = torch.nn.ModuleList() self.convs.extend([ ChebConv(hidden, hidden, K=2, aggr='mean') for i in range(num_layers - 1) ]) self.pools.extend([ ASAPooling(hidden, ratio, dropout=dropout) for i in range((num_layers) // 2) ]) self.jump = JumpingKnowledge(mode='cat') self.lin1 = Linear(num_layers * hidden, hidden) self.lin2 = Linear(hidden, num_classes) def reset_parameters(self): self.conv1.reset_parameters() for conv in self.convs: conv.reset_parameters() for pool in self.pools: pool.reset_parameters() self.lin1.reset_parameters() self.lin2.reset_parameters() def forward(self, data): x, edge_index, batch = data.x, data.edge_index, data.batch edge_weight = data.edge_attr.float() x = F.relu(self.conv1(x, edge_index, edge_weight)) x = F.relu(self.conv2(x, edge_index, edge_weight)) xs = [global_mean_pool(x, batch)] for i, conv in enumerate(self.convs): x = conv(x=x, edge_index=edge_index, edge_weight=edge_weight) x = F.relu(x) xs += [global_mean_pool(x, batch)] if i % 2 == 0 and i < len(self.convs) - 1: pool = self.pools[i // 2] x, edge_index, edge_weight, batch, _ = pool( x=x, edge_index=edge_index, edge_weight=edge_weight, batch=batch) x = self.jump(xs) x = F.relu(self.lin1(x)) x = F.dropout(x, p=0.5, training=self.training) x = self.lin2(x) return F.log_softmax(x, dim=-1)
class Net(torch.nn.Module): def __init__(self, dataset): super(Net, self).__init__() self.conv1 = ChebConv(dataset.num_features, args.hidden, args.num_hops) self.conv2 = ChebConv(args.hidden, dataset.num_classes, args.num_hops) def reset_parameters(self): self.conv1.reset_parameters() self.conv2.reset_parameters() def forward(self, data): x, edge_index = data.x, data.edge_index x = F.relu(self.conv1(x, edge_index)) x = F.dropout(x, p=args.dropout, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1)
class ChebNet(nn.Module): def __init__(self, input_dim, hidden_dim, num_classes, num_hops, dropout): super(ChebNet, self).__init__() self.dropout = dropout self.layer_1 = ChebConv(input_dim, hidden_dim, num_hops) self.layer_2 = ChebConv(hidden_dim, num_classes, num_hops) def reset_parameters(self): self.layer_1.reset_parameters() self.layer_2.reset_parameters() def forward(self, data): x, edge_index = data.x, data.edge_index x = F.dropout(x, p=self.dropout, training=self.training) x = self.layer_1(x, edge_index) x = F.relu(x) x = F.dropout(x, p=self.dropout, training=self.training) x = self.layer_2(x, edge_index) return F.log_softmax(x, dim=1)
class Cheb_Net(torch.nn.Module): def __init__(self, features_num, num_class, hidden, num_hops, dropout): super(Cheb_Net, self).__init__() self.dropout = dropout self.conv1 = ChebConv(features_num, hidden, num_hops) self.conv2 = ChebConv(hidden, num_class, num_hops) def reset_parameters(self): self.conv1.reset_parameters() self.conv2.reset_parameters() def forward(self, data): x, edge_index = data.x, data.edge_index x = F.relu(self.conv1(x, edge_index)) x = F.dropout(x, p=self.dropout, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=-1) def __repr__(self): return self.__class__.__name__
class ChebNet(nn.Module): """ 2 Layer ChebNet based on pytorch geometric. Parameters ---------- nfeat : int size of input feature dimension nhid : int number of hidden units nclass : int size of output dimension num_hops: int number of hops in ChebConv dropout : float dropout rate for ChebNet lr : float learning rate for ChebNet weight_decay : float weight decay coefficient (l2 normalization) for GCN. When `with_relu` is True, `weight_decay` will be set to 0. with_bias: bool whether to include bias term in ChebNet weights. device: str 'cpu' or 'cuda'. Examples -------- We can first load dataset and then train ChebNet. >>> from deeprobust.graph.data import Dataset >>> from deeprobust.graph.defense import ChebNet >>> data = Dataset(root='/tmp/', name='cora') >>> adj, features, labels = data.adj, data.features, data.labels >>> idx_train, idx_val, idx_test = data.idx_train, data.idx_val, data.idx_test >>> cheby = ChebNet(nfeat=features.shape[1], nhid=16, num_hops=3, nclass=labels.max().item() + 1, dropout=0.5, device='cpu') >>> cheby = cheby.to('cpu') >>> pyg_data = Dpr2Pyg(data) # convert deeprobust dataset to pyg dataset >>> cheby.fit(pyg_data, patience=10, verbose=True) # train with earlystopping """ def __init__(self, nfeat, nhid, nclass, num_hops=3, dropout=0.5, lr=0.01, weight_decay=5e-4, with_bias=True, device=None): super(ChebNet, self).__init__() assert device is not None, "Please specify 'device'!" self.device = device self.conv1 = ChebConv(nfeat, nhid, K=num_hops, bias=with_bias) self.conv2 = ChebConv(nhid, nclass, K=num_hops, bias=with_bias) self.dropout = dropout self.weight_decay = weight_decay self.lr = lr self.output = None self.best_model = None self.best_output = None def forward(self, data): x, edge_index = data.x, data.edge_index x = F.relu(self.conv1(x, edge_index)) x = F.dropout(x, p=self.dropout, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1) def initialize(self): """Initialize parameters of ChebNet. """ self.conv1.reset_parameters() self.conv2.reset_parameters() def fit(self, pyg_data, train_iters=200, initialize=True, verbose=False, patience=500, **kwargs): """Train the ChebNet model, when idx_val is not None, pick the best model according to the validation loss. Parameters ---------- pyg_data : pytorch geometric dataset object train_iters : int number of training epochs initialize : bool whether to initialize parameters before training verbose : bool whether to show verbose logs patience : int patience for early stopping, only valid when `idx_val` is given """ self.device = self.conv1.weight.device if initialize: self.initialize() self.data = pyg_data[0].to(self.device) # By default, it is trained with early stopping on validation self.train_with_early_stopping(train_iters, patience, verbose) def train_with_early_stopping(self, train_iters, patience, verbose): """early stopping based on the validation loss """ if verbose: print('=== training ChebNet model ===') optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay) labels = self.data.y train_mask, val_mask = self.data.train_mask, self.data.val_mask early_stopping = patience best_loss_val = 100 for i in range(train_iters): self.train() optimizer.zero_grad() output = self.forward(self.data) loss_train = F.nll_loss(output[train_mask], labels[train_mask]) loss_train.backward() optimizer.step() if verbose and i % 10 == 0: print('Epoch {}, training loss: {}'.format( i, loss_train.item())) self.eval() output = self.forward(self.data) loss_val = F.nll_loss(output[val_mask], labels[val_mask]) if best_loss_val > loss_val: best_loss_val = loss_val self.output = output weights = deepcopy(self.state_dict()) patience = early_stopping else: patience -= 1 if i > early_stopping and patience <= 0: break if verbose: print('=== early stopping at {0}, loss_val = {1} ==='.format( i, best_loss_val)) self.load_state_dict(weights) def test(self): """Evaluate ChebNet performance on test set. Parameters ---------- idx_test : node testing indices """ self.eval() test_mask = self.data.test_mask labels = self.data.y output = self.forward(self.data) # output = self.output loss_test = F.nll_loss(output[test_mask], labels[test_mask]) acc_test = utils.accuracy(output[test_mask], labels[test_mask]) print("Test set results:", "loss= {:.4f}".format(loss_test.item()), "accuracy= {:.4f}".format(acc_test.item())) return acc_test.item()