class GIN(torch.nn.Module): """ 多層化対応モデル(AutoGraphで使われていたモデル) グラフ構造が重要なので最初から畳み込み層に入力する """ def __init__(self, num_node_features=100, num_class=18, hidden=16, dropout_rate=0.5, num_layers=2, eps=0, train_eps=True): super(GIN, self).__init__() self.first_conv = GINConv( Sequential(Linear(num_node_features, hidden), ReLU(), Linear(hidden, hidden)), eps, train_eps) self.first_bn = BatchNorm1d(hidden) self.nns = torch.nn.ModuleList() self.convs = torch.nn.ModuleList() self.bns = torch.nn.ModuleList() for i in range(num_layers): self.nns.append( Sequential(Linear(hidden, hidden), ReLU(), Linear(hidden, hidden))) self.bns.append(BatchNorm1d(hidden)) self.convs.append(GINConv(self.nns[i], eps, train_eps)) self.lin1 = Linear(hidden, hidden) self.lin2 = Linear(hidden, num_class) self.dropout_rate = dropout_rate def reset_parameters(self): self.first_conv.reset_parameters() self.first_bn.reset_parameters() for nn, conv, bn in zip(self.nns, self.convs, self.bns): nn.reset_parameters() conv.reset_parameters() bn.reset_parameters() self.lin1.reset_parameters() self.lin2.reset_parameters() def forward(self, data): x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight #x = F.relu(self.first_lin(x)) x = F.relu(self.first_conv(x, edge_index)) x = self.first_bn(x) for conv, bn in zip(self.convs, self.bns): x = F.relu(conv(x, edge_index)) x = bn(x) #x = F.dropout(x, self.dropout_rate, training=self.training) x = F.relu(self.lin1(x)) x = F.dropout(x, self.dropout_rate, training=self.training) x = self.lin2(x) return F.log_softmax(x, dim=-1) def __repr__(self): return self.__class__.__name__
class GIN0WithJK(torch.nn.Module): def __init__(self, dataset, num_layers, hidden, mode='cat'): super(GIN0WithJK, self).__init__() self.conv1 = GINConv(Sequential( Linear(dataset.num_features, hidden), ReLU(), Linear(hidden, hidden), ReLU(), BN(hidden), ), train_eps=False) self.convs = torch.nn.ModuleList() for i in range(num_layers - 1): self.convs.append( GINConv(Sequential( Linear(hidden, hidden), ReLU(), Linear(hidden, hidden), ReLU(), BN(hidden), ), train_eps=False)) self.jump = JumpingKnowledge(mode) if mode == 'cat': self.lin1 = Linear(num_layers * hidden, hidden) else: self.lin1 = Linear(hidden, hidden) self.lin2 = Linear(hidden, dataset.num_classes) def reset_parameters(self): self.conv1.reset_parameters() for conv in self.convs: conv.reset_parameters() self.jump.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 x = self.conv1(x, edge_index) xs = [x] for conv in self.convs: x = conv(x, edge_index) xs += [x] x = self.jump(xs) x = global_mean_pool(x, batch) 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) def __repr__(self): return self.__class__.__name__
class GIN0(torch.nn.Module): def __init__(self, dataset, num_layers, hidden, add_pool=False): super(GIN0, self).__init__() self.conv1 = GINConv(Sequential( Linear(dataset.num_features, hidden), ReLU(), Linear(hidden, hidden), ReLU(), BN(hidden), ), train_eps=False) self.convs = torch.nn.ModuleList() for i in range(num_layers - 1): self.convs.append( GINConv(Sequential( Linear(hidden, hidden), ReLU(), Linear(hidden, hidden), ReLU(), BN(hidden), ), train_eps=False)) self.lin1 = Linear(hidden, hidden) self.lin2 = Linear(hidden, dataset.num_classes) self.add_pool = add_pool def reset_parameters(self): self.conv1.reset_parameters() for conv in self.convs: conv.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 x = self.conv1(x, edge_index) for conv in self.convs: x = conv(x, edge_index) if self.add_pool: x = global_add_pool(x, batch) else: x = global_mean_pool(x, batch) 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) def __repr__(self): return self.__class__.__name__
class GINLayer(nn.Module): def __init__(self, in_channels, out_channels): super(GINLayer, self).__init__() seq = nn.Sequential( nn.Linear(in_channels, out_channels), nn.ReLU(), nn.Linear(out_channels, out_channels), ) self.conv = GINConv(seq) self.bn = nn.BatchNorm1d(out_channels) def reset_parameters(self): self.conv.reset_parameters() # TODO: reset bn? def forward(self, x, edge_index): x = F.relu(self.conv(x, edge_index)) x = self.bn(x) return x
class GIN(torch.nn.Module): def __init__(self, num_layers=3, hidden=32, features_num=32, num_class=2): super(GIN, self).__init__() self.lin1 = Linear(hidden, hidden) self.lin3 = Linear(hidden, hidden) self.conv1 = GINConv(self.lin1) self.conv2 = GINConv(self.lin3) self.lin2 = Linear(hidden, num_class) self.fuse_weight = torch.nn.Parameter(torch.FloatTensor(num_layers),requires_grad=True) self.fuse_weight.data.fill_(float(1) / (3)) self.first_lin = Linear(features_num, hidden) def reset_parameters(self): self.first_lin.reset_parameters() self.conv1.reset_parameters() for conv in self.convs: conv.reset_parameters() self.lin2.reset_parameters() def forward(self, data): x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight x = F.relu(self.first_lin(x)) x = F.dropout(x, p=0.5, training=self.training) x_first = x x = F.relu(self.conv1(x, edge_index)) x = F.dropout(x, p=0.2, training=self.training) x = x + self.fuse_weight[0] * x_first x_first = x x = F.relu(self.conv2(x, edge_index)) #, edge_attr=edge_weight)) x = F.dropout(x, p=0.2, training=self.training) x = x + self.fuse_weight[1] * x_first #x=x+first_x x = self.lin2(x) return F.log_softmax(x, dim=-1)
class GIN(torch.nn.Module): def __init__(self, args): super(GIN, self).__init__() self.args = set_default(args, { 'num_layers': 2, 'hidden': 64, 'hidden2': 32, 'dropout': 0.5, 'lr': 0.005, 'epoches': 300, 'weight_decay': 5e-4, 'act': 'leaky_relu', 'withbn': True, }) self.timer = self.args['timer'] self.dropout = self.args['dropout'] self.agg = self.args['agg'] self.conv1 = GINConv( Sequential( Linear(self.args['features_num'], self.args['hidden']), ReLU(), BatchNorm1d(self.args['hidden']), ), train_eps=True) self.convs = torch.nn.ModuleList() hd = [self.args['hidden']] for i in range(self.args['num_layers'] - 1): hd.append(self.args['hidden2']) self.convs.append( GINConv( Sequential( Linear(self.args['hidden'], self.args['hidden2']), ReLU(), BatchNorm1d(self.args['hidden2']), ), train_eps=True)) if self.args['agg'] == 'concat': outdim = sum(hd) elif self.args['agg'] == 'self': outdim = hd[-1] if self.args['act'] == 'leaky_relu': self.act = F.leaky_relu elif self.args['act'] == 'tanh': self.act = torch.tanh else: self.act = lambda x: x self.lin2 = Linear(outdim, self.args['num_class']) def reset_parameters(self): self.conv1.reset_parameters() for conv in self.convs: conv.reset_parameters() self.lin2.reset_parameters() def forward(self, data): x, edge_index, edge_weight = data.x, data.edge_index, data.edge_weight x = self.act(self.conv1(x, edge_index)) xs = [x] for conv in self.convs: x = self.act(conv(x, edge_index)) xs.append(x) #x = F.dropout(x, p=self.dropout, training=self.training) if self.agg == 'concat': x = torch.cat(xs, dim=1) elif self.agg == 'self': x = xs[-1] x = self.lin2(x) return F.log_softmax(x, dim=-1) def train_predict(self, data, train_mask=None, val_mask=None, return_out=True): if train_mask is None: train_mask = data.train_mask optimizer = torch.optim.Adam(self.parameters(), lr=self.args['lr'], weight_decay=self.args['weight_decay']) flag_end = False st = time.time() for epoch in range(1, self.args['epoches']): self.train() optimizer.zero_grad() res = self.forward(data) loss = F.nll_loss(res[train_mask], data.y[train_mask]) loss.backward() optimizer.step() if epoch%50 == 0: cost = (time.time()-st)/epoch*50 if max(cost*10, 5) > self.timer.remain_time(): flag_end = True break test_mask = data.test_mask self.eval() with torch.no_grad(): res = self.forward(data) if return_out: pred = res else: pred = res[test_mask] if val_mask is not None: return pred, res[val_mask], flag_end return pred, flag_end def __repr__(self): return self.__class__.__name__
class JK(nn.Module): def __init__(self, nfeat, nhid, nclass, dropout=0.5, lr=0.01, weight_decay=5e-4, n_edge=1, with_relu=True, drop=False, with_bias=True, device=None): super(JK, self).__init__() assert device is not None, "Please specify 'device'!" self.device = device self.nfeat = nfeat self.hidden_sizes = [nhid] self.nclass = int(nclass) self.dropout = dropout self.lr = lr self.drop = drop if not with_relu: self.weight_decay = 0 else: self.weight_decay = weight_decay self.with_relu = with_relu self.with_bias = with_bias self.n_edge = n_edge self.output = None self.best_model = None self.best_output = None self.adj_norm = None self.features = None self.gate = Parameter(torch.rand(1)) # creat a generator between [0,1] # self.beta = Parameter(torch.Tensor(self.n_edge)) nclass = int(nclass) """JK from torch-geometric""" num_features = nfeat dim = nhid nn1 = Sequential( Linear(num_features, dim), ReLU(), ) self.gc1 = GINConv(nn1) self.bn1 = torch.nn.BatchNorm1d(dim) nn2 = Sequential( Linear(dim, dim), ReLU(), ) self.gc2 = GINConv(nn2) nn3 = Sequential( Linear(dim, dim), ReLU(), ) self.gc3 = GINConv(nn3) self.jump = JumpingKnowledge(mode='cat') # 'cat', 'lstm', 'max' self.bn2 = torch.nn.BatchNorm1d(dim) # self.fc1 = Linear(dim*3, dim) self.fc2 = Linear(dim * 2, int(nclass)) def forward(self, x, adj): """we don't change the edge_index, just update the edge_weight; some edge_weight are regarded as removed if it equals to zero""" x = x.to_dense() edge_index = adj._indices() """GJK-Nets""" if self.attention: adj = self.att_coef(x, adj, i=0) x1 = F.relu( self.gc1(x, edge_index=edge_index, edge_weight=adj._values())) if self.attention: # if attention=True, use attention mechanism adj_2 = self.att_coef(x1, adj, i=1) adj_values = self.gate * adj._values() + ( 1 - self.gate) * adj_2._values() else: adj_values = adj._values() x1 = F.dropout(x1, self.dropout, training=self.training) x2 = F.relu(self.gc2(x1, edge_index=edge_index, edge_weight=adj_values)) # x2 = self.bn1(x2) # if self.attention: # if attention=True, use attention mechanism # adj_3 = self.att_coef(x2, adj, i=1) # adj_values = self.gate * adj_2._values() + (1 - self.gate) * adj_3._values() # else: # adj_values = adj._values() x2 = F.dropout(x2, self.dropout, training=self.training) # x3 = F.relu(self.gc2(x2, edge_index=edge_index, edge_weight=adj_values)) # x3 = F.dropout(x3, self.dropout, training=self.training) x_last = self.jump([x1, x2]) x_last = F.dropout(x_last, self.dropout, training=self.training) x_last = self.fc2(x_last) return F.log_softmax(x_last, dim=1) def initialize(self): self.gc1.reset_parameters() self.gc2.reset_parameters() self.fc2.reset_parameters() try: self.jump.reset_parameters() self.fc1.reset_parameters() self.gc3.reset_parameters() except: pass def att_coef(self, fea, edge_index, is_lil=False, i=0): if is_lil == False: edge_index = edge_index._indices() else: edge_index = edge_index.tocoo() n_node = fea.shape[0] row, col = edge_index[0].cpu().data.numpy()[:], edge_index[1].cpu( ).data.numpy()[:] # row, col = edge_index[0], edge_index[1] fea_copy = fea.cpu().data.numpy() sim_matrix = cosine_similarity(X=fea_copy, Y=fea_copy) # try cosine similarity sim = sim_matrix[row, col] sim[sim < 0.1] = 0 # print('dropped {} edges'.format(1-sim.nonzero()[0].shape[0]/len(sim))) # """use jaccard for binary features and cosine for numeric features""" # fea_start, fea_end = fea[edge_index[0]], fea[edge_index[1]] # isbinray = np.array_equal(fea_copy, fea_copy.astype(bool)) # check is the fea are binary # np.seterr(divide='ignore', invalid='ignore') # if isbinray: # fea_start, fea_end = fea_start.T, fea_end.T # sim = jaccard_score(fea_start, fea_end, average=None) # similarity scores of each edge # else: # fea_copy[np.isinf(fea_copy)] = 0 # fea_copy[np.isnan(fea_copy)] = 0 # sim_matrix = cosine_similarity(X=fea_copy, Y=fea_copy) # try cosine similarity # sim = sim_matrix[edge_index[0], edge_index[1]] # sim[sim < 0.01] = 0 """build a attention matrix""" att_dense = lil_matrix((n_node, n_node), dtype=np.float32) att_dense[row, col] = sim if att_dense[0, 0] == 1: att_dense = att_dense - sp.diags( att_dense.diagonal(), offsets=0, format="lil") # normalization, make the sum of each row is 1 att_dense_norm = normalize(att_dense, axis=1, norm='l1') """add learnable dropout, make character vector""" if self.drop: character = np.vstack( (att_dense_norm[row, col].A1, att_dense_norm[col, row].A1)) character = torch.from_numpy(character.T) drop_score = self.drop_learn_1(character) drop_score = torch.sigmoid( drop_score ) # do not use softmax since we only have one element mm = torch.nn.Threshold(0.5, 0) drop_score = mm(drop_score) mm_2 = torch.nn.Threshold(-0.49, 1) drop_score = mm_2(-drop_score) drop_decision = drop_score.clone().requires_grad_() # print('rate of left edges', drop_decision.sum().data/drop_decision.shape[0]) drop_matrix = lil_matrix((n_node, n_node), dtype=np.float32) drop_matrix[row, col] = drop_decision.cpu().data.numpy().squeeze(-1) att_dense_norm = att_dense_norm.multiply( drop_matrix.tocsr()) # update, remove the 0 edges if att_dense_norm[ 0, 0] == 0: # add the weights of self-loop only add self-loop at the first layer degree = (att_dense_norm != 0).sum(1).A1 # degree = degree.squeeze(-1).squeeze(-1) lam = 1 / (degree + 1) # degree +1 is to add itself self_weight = sp.diags(np.array(lam), offsets=0, format="lil") att = att_dense_norm + self_weight # add the self loop else: att = att_dense_norm att_adj = edge_index att_edge_weight = att[row, col] att_edge_weight = np.exp(att_edge_weight) # exponent, kind of softmax att_edge_weight = torch.tensor(np.array(att_edge_weight)[0], dtype=torch.float32).cuda() shape = (n_node, n_node) new_adj = torch.sparse.FloatTensor(att_adj, att_edge_weight, shape) return new_adj def add_loop_sparse(self, adj, fill_value=1): # make identify sparse tensor row = torch.range(0, int(adj.shape[0] - 1), dtype=torch.int64) i = torch.stack((row, row), dim=0) v = torch.ones(adj.shape[0], dtype=torch.float32) shape = adj.shape I_n = torch.sparse.FloatTensor(i, v, shape) return adj + I_n.to(self.device) def fit( self, features, adj, labels, idx_train, idx_val=None, idx_test=None, train_iters=81, att_0=None, attention=False, model_name=None, initialize=True, verbose=False, normalize=False, patience=500, ): ''' train the gcn model, when idx_val is not None, pick the best model according to the validation loss ''' self.sim = None self.attention = attention self.idx_test = idx_test # self.device = self.gc1.weight.device if initialize: self.initialize() if type(adj) is not torch.Tensor: features, adj, labels = utils.to_tensor(features, adj, labels, device=self.device) else: features = features.to(self.device) adj = adj.to(self.device) labels = labels.to(self.device) # normalize = False # we don't need normalize here, the norm is conducted in the GCN (self.gcn1) model # if normalize: # if utils.is_sparse_tensor(adj): # adj_norm = utils.normalize_adj_tensor(adj, sparse=True) # else: # adj_norm = utils.normalize_adj_tensor(adj) # else: # adj_norm = adj adj = self.add_loop_sparse(adj) """Make the coefficient D^{-1/2}(A+I)D^{-1/2}""" self.adj_norm = adj self.features = features self.labels = labels if idx_val is None: self._train_without_val(labels, idx_train, train_iters, verbose) else: if patience < train_iters: self._train_with_early_stopping(labels, idx_train, idx_val, train_iters, patience, verbose) else: self._train_with_val(labels, idx_train, idx_val, train_iters, verbose) def _train_without_val(self, labels, idx_train, train_iters, verbose): self.train() optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay) for i in range(train_iters): optimizer.zero_grad() output = self.forward(self.features, self.adj_norm) loss_train = F.nll_loss( output[idx_train], labels[idx_train], weight=None ) # this weight is the weight of each training nodes 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.features, self.adj_norm) self.output = output def _train_with_val(self, labels, idx_train, idx_val, train_iters, verbose): if verbose: print('=== training gcn model ===') optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay) best_loss_val = 100 best_acc_val = 0 for i in range(train_iters): self.train() optimizer.zero_grad() output = self.forward(self.features, self.adj_norm) loss_train = F.nll_loss(output[idx_train], labels[idx_train]) loss_train.backward() optimizer.step() # pred = output[self.idx_test].max(1)[1] # acc_test =accuracy(output[self.idx_test], labels[self.idx_test]) # acc_test = pred.eq(labels[self.idx_test]).sum().item() / self.idx_test.shape[0] self.eval() output = self.forward(self.features, self.adj_norm) loss_val = F.nll_loss(output[idx_val], labels[idx_val]) acc_val = utils.accuracy(output[idx_val], labels[idx_val]) if verbose and i % 20 == 0: print('Epoch {}, training loss: {}, test acc: {}'.format( i, loss_train.item(), acc_val)) if best_loss_val > loss_val: best_loss_val = loss_val self.output = output weights = deepcopy(self.state_dict()) if acc_val > best_acc_val: best_acc_val = acc_val self.output = output weights = deepcopy(self.state_dict()) if verbose: print( '=== picking the best model according to the performance on validation ===' ) self.load_state_dict(weights) def _train_with_early_stopping(self, labels, idx_train, idx_val, train_iters, patience, verbose): if verbose: print('=== training gcn model ===') optimizer = optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay) early_stopping = patience best_loss_val = 100 for i in range(train_iters): self.train() optimizer.zero_grad() output = self.forward(self.features, self.adj_norm) loss_train = F.nll_loss(output[idx_train], labels[idx_train]) loss_train.backward() optimizer.step() self.eval() output = self.forward(self.features, self.adj_norm) if verbose and i % 10 == 0: print('Epoch {}, training loss: {}'.format( i, loss_train.item())) loss_val = F.nll_loss(output[idx_val], labels[idx_val]) 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, idx_test, model_name=None): # self.model_name = model_name self.eval() output = self.predict() # output = self.output loss_test = F.nll_loss(output[idx_test], self.labels[idx_test]) acc_test = utils.accuracy(output[idx_test], self.labels[idx_test]) print("Test set results:", "loss= {:.4f}".format(loss_test.item()), "accuracy= {:.4f}".format(acc_test.item())) return acc_test, output def _set_parameters(self): # TODO pass def predict(self, features=None, adj=None): '''By default, inputs are unnormalized data''' # self.eval() if features is None and adj is None: return self.forward(self.features, self.adj_norm) else: if type(adj) is not torch.Tensor: features, adj = utils.to_tensor(features, adj, device=self.device) self.features = features if utils.is_sparse_tensor(adj): self.adj_norm = utils.normalize_adj_tensor(adj, sparse=True) else: self.adj_norm = utils.normalize_adj_tensor(adj) return self.forward(self.features, self.adj_norm)
class cut_MPNN(torch.nn.Module): def __init__(self, dataset, num_layers, hidden1, hidden2, deltas, elasticity=0.01, num_iterations = 30): super(cut_MPNN, self).__init__() self.hidden1 = hidden1 self.hidden2 = hidden2 self.conv1 = GINConv(Sequential( Linear(1, self.hidden1), ReLU(), Linear(self.hidden1, self.hidden1), ReLU(), BN( self.hidden1), ),train_eps=False) self.num_iterations = num_iterations self.convs = torch.nn.ModuleList() self.deltas = deltas self.numlayers = num_layers self.elasticity = elasticity self.bns = torch.nn.ModuleList() for i in range(num_layers-1): self.bns.append(BN( self.hidden1)) self.convs = torch.nn.ModuleList() for i in range(num_layers - 1): self.convs.append(GINConv(Sequential( Linear( self.hidden1, self.hidden1), ReLU(), Linear( self.hidden1, self.hidden1), ReLU(), BN(self.hidden1), ),train_eps=False)) self.conv2 = GATAConv( self.hidden1, self.hidden2 ,heads=8) self.lin1 = Linear(8*self.hidden2, self.hidden1) self.bn2 = BN(self.hidden1) self.lin2 = Linear(self.hidden1, 1) def reset_parameters(self): self.conv1.reset_parameters() self.conv2.reset_parameters() for conv in self.convs: conv.reset_parameters() for bn in self.bns: bn.reset_parameters() self.lin1.reset_parameters() self.bn2.reset_parameters() self.lin2.reset_parameters() def forward(self, data, tvol = None): x = data.x edge_index = data.edge_index batch = data.batch xinit= x.clone() row, col = edge_index mask = get_mask(x,edge_index,1).to(x.dtype).unsqueeze(-1) x = self.conv1(x, edge_index) xpostconv1 = x.detach() x = x*mask for conv, bn in zip(self.convs, self.bns): if(x.dim()>1): x = x + conv(x, edge_index) mask = get_mask(mask,edge_index,1).to(x.dtype) x = x*mask x = bn(x) x = self.conv2(x, edge_index) mask = get_mask(mask,edge_index,1).to(x.dtype) x = x*mask xpostconvs = x.detach() # x = F.leaky_relu(self.lin1(x)) x = x*mask x = self.bn2(x) xpostlin1 = x.detach() x = F.dropout(x, p=0.5, training=self.training) x = F.leaky_relu(self.lin2(x)) x = x*mask xprethresh = x.detach() N_size = x.shape[0] batch_max = scatter_max(x, batch, 0, dim_size= N_size)[0] batch_max = torch.index_select(batch_max, 0, batch) batch_min = scatter_min(x, batch, 0, dim_size= N_size)[0] batch_min = torch.index_select(batch_min, 0, batch) #min-max normalize x = (x-batch_min)/(batch_max+1e-6-batch_min) x = x*mask + mask*1e-6 #add dirac in the set x = x + xinit.unsqueeze(-1) #calculate x2 = x.detach() r, c = edge_index tv = total_var(x, edge_index, batch) deg = degree(r).unsqueeze(-1) conduct_1 = (tv) totalvol = scatter_add(deg.detach()*torch.ones_like(x, device=device), batch, 0)+1e-6 totalcard = scatter_add(torch.ones_like(x, device=device), batch, 0)+1e-6 #receptive field recvol_hard = scatter_add(deg*mask.float(), batch, 0, dim_size = batch.max().item()+1)+1e-6 reccard_hard = scatter_add(mask.float(), batch, 0, dim_size = batch.max().item()+1)+1e-6 assert recvol_hard.mean()/totalvol.mean() <=1, "Something went wrong! Receptive field is larger than total volume." target = torch.zeros_like(totalvol) #generate target vol if tvol is None: feasible_vols = data.recfield_vol/data.total_vol-0.0 target = torch.rand_like(feasible_vols, device=device)*feasible_vols*0.85 + 0.1 target = target.squeeze(-1)*totalvol.squeeze(-1) else: target = tvol*totalvol.squeeze(-1) a = torch.ones((batch.max().item()+1,1), device = device) xfilt = x ############################################################################### #iterative rescaling counter_no2 = 0 for iteration in range(self.num_iterations): counter_no2 += 1 keep = (((a[batch]*xfilt)<1).to(x.dtype)) x_k, d_k, d_nk = xfilt*keep*mask, deg*keep*mask, deg*(1-keep)*mask diff = target.unsqueeze(-1) - scatter_add(d_nk, batch, 0) dot = scatter_add(x_k*d_k, batch, 0) a = diff/(dot+1e-5) volcur = (scatter_add(torch.clamp(a[batch]*xfilt,max = 1., min = 0.)*deg,batch,0)) volcheck = (torch.abs(target - volcur.squeeze(-1))>0.1) checki = torch.abs(target.squeeze(-1)-volcur.squeeze(-1))>0.01 targetcheck = torch.abs(volcur.squeeze(-1) - target) check = (targetcheck<= self.elasticity*target).to(x.dtype) if (tvol is not None): pass if(check.sum()>=batch.max().item()+1): break; probs = torch.clamp(a[batch]*x*mask, max = 1., min = 0.) ############################################################################### #collect useful numbers x2 = ((probs - torch.rand_like(x, device = device))>0).float() vol_1 = scatter_add(probs*deg, batch, 0)+1e-6 card_1 = scatter_add(probs, batch,0) rec_field = scatter_add(mask, batch, 0)+1e-6 cut_size = scatter_add(x2, batch, 0) tv_hard = total_var(x2, edge_index, batch) vol_hard = scatter_add(deg*x2, batch, 0, dim_size = batch.max().item()+1)+1e-6 conduct_hard = tv_hard/vol_hard rec_field_ratio = cut_size/rec_field rec_field_volratio = vol_hard/recvol_hard total_vol_ratio = vol_hard/totalvol #calculate loss expected_cut = scatter_add(probs*deg, batch, 0) - scatter_add((probs[row]*probs[col]), batch[row], 0) loss = expected_cut #return dict retdict = {} retdict["output"] = [probs.squeeze(-1),"hist"] #output #retdict["|Expected_vol - Target|"]= [targetcheck, "sequence"] #absolute distance from targetvol retdict["Expected_volume"] = [vol_1.mean(),"sequence"] #volume retdict["Expected_cardinality"] = [card_1.mean(),"sequence"] retdict["volume_hard"] = [vol_hard.mean(),"sequence"] #volume2 #retdict["cut1"] = [tv.mean(),"sequence"] #cut1 retdict["cut_hard"] = [tv_hard.mean(),"sequence"] #cut1 retdict["Average cardinality ratio of receptive field "] = [rec_field_ratio.mean(),"sequence"] retdict["Recfield volume/Total volume"] = [recvol_hard.mean()/totalvol.mean(), "sequence"] retdict["Average ratio of receptive field volume"]= [rec_field_volratio.mean(),'sequence'] retdict["Average ratio of total volume"]= [total_vol_ratio.mean(),'sequence'] retdict["mask"] = [mask, "aux"] #mask retdict["xinit"] = [xinit,"hist"] #layer input diracs retdict["xpostlin1"] = [xpostlin1.mean(1),"hist"] #after first linear layer retdict["xprethresh"] = [xprethresh.mean(1),"hist"] #pre thresholding activations 195 x 1 retdict["lossvol"] = [lossvol.mean(),"sequence"] #volume constraint retdict["losscard"] = [losscard.mean(),"sequence"] #cardinality constraint retdict["loss"] = [loss.mean().squeeze(),"sequence"] #final loss return retdict def __repr__(self): return self.__class__.__name__
class clique_MPNN(torch.nn.Module): def __init__(self, dataset, num_layers, hidden1, hidden2, deltas, elasticity=0.01, num_iterations = 30): super(cliqueMPNN_hindsight_earlyGAT, self).__init__() self.hidden1 = hidden1 self.hidden2 = hidden2 self.momentum = 0.1 self.num_iterations = num_iterations self.convs = torch.nn.ModuleList() self.deltas = deltas self.numlayers = num_layers self.elasticity = elasticity self.heads = 8 self.concat = True self.bns = torch.nn.ModuleList() for i in range(num_layers-1): self.bns.append(BN(self.heads*self.hidden1, momentum=self.momentum)) self.convs = torch.nn.ModuleList() for i in range(num_layers - 1): self.convs.append(GINConv(Sequential( Linear( self.heads*self.hidden1, self.heads*self.hidden1), ReLU(), Linear( self.heads*self.hidden1, self.heads*self.hidden1), ReLU(), BN(self.heads*self.hidden1, momentum=self.momentum), ),train_eps=True)) self.bn1 = BN(self.heads*self.hidden1) self.conv1 = GINConv(Sequential(Linear(self.hidden2, self.heads*self.hidden1), ReLU(), Linear( self.heads*self.hidden1, self.heads*self.hidden1), ReLU(), BN(self.heads*self.hidden1, momentum=self.momentum), ),train_eps=True) if self.concat: self.lin1 = Linear(self.heads*self.hidden1, self.hidden1) else: self.lin1 = Linear(self.hidden1, self.hidden1) self.lin2 = Linear(self.hidden1, 1) self.gnorm = GraphSizeNorm() def reset_parameters(self): self.conv1.reset_parameters() for conv in self.convs: conv.reset_parameters() for bn in self.bns: bn.reset_parameters() self.bn1.reset_parameters() self.lin1.reset_parameters() self.lin2.reset_parameters() def forward(self, data, edge_dropout = None, penalty_coefficient = 0.25): x = data.x edge_index = data.edge_index batch = data.batch num_graphs = batch.max().item() + 1 row, col = edge_index total_num_edges = edge_index.shape[1] N_size = x.shape[0] if edge_dropout is not None: edge_index = dropout_adj(edge_index, edge_attr = (torch.ones(edge_index.shape[1], device=device)).long(), p = edge_dropout, force_undirected=True)[0] edge_index = add_remaining_self_loops(edge_index, num_nodes = batch.shape[0])[0] reduced_num_edges = edge_index.shape[1] current_edge_percentage = (reduced_num_edges/total_num_edges) no_loop_index,_ = remove_self_loops(edge_index) no_loop_row, no_loop_col = no_loop_index xinit= x.clone() x = x.unsqueeze(-1) mask = get_mask(x,edge_index,1).to(x.dtype) x = F.leaky_relu(self.conv1(x, edge_index))# +x x = x*mask x = self.gnorm(x) x = self.bn1(x) for conv, bn in zip(self.convs, self.bns): if(x.dim()>1): x = x+F.leaky_relu(conv(x, edge_index)) mask = get_mask(mask,edge_index,1).to(x.dtype) x = x*mask x = self.gnorm(x) x = bn(x) xpostconvs = x.detach() # x = F.leaky_relu(self.lin1(x)) x = x*mask xpostlin1 = x.detach() x = F.leaky_relu(self.lin2(x)) x = x*mask #calculate min and max batch_max = scatter_max(x, batch, 0, dim_size= N_size)[0] batch_max = torch.index_select(batch_max, 0, batch) batch_min = scatter_min(x, batch, 0, dim_size= N_size)[0] batch_min = torch.index_select(batch_min, 0, batch) #min-max normalize x = (x-batch_min)/(batch_max+1e-6-batch_min) probs=x x2 = x.detach() deg = degree(row).unsqueeze(-1) totalvol = scatter_add(deg.detach()*torch.ones_like(x, device=device), batch, 0)+1e-6 totalcard = scatter_add(torch.ones_like(x, device=device), batch, 0)+1e-6 x2 = ((x2 - torch.rand_like(x, device = device))>0).float() vol_1 = scatter_add(probs*deg, batch, 0)+1e-6 card_1 = scatter_add(probs, batch,0) set_size = scatter_add(x2, batch, 0) vol_hard = scatter_add(deg*x2, batch, 0, dim_size = batch.max().item()+1)+1e-6 total_vol_ratio = vol_hard/totalvol #calculating the terms for the expected distance between clique and graph pairwise_prodsums = torch.zeros(num_graphs, device = device) for graph in range(num_graphs): batch_graph = (batch==graph) pairwise_prodsums[graph] = (torch.conv1d(probs[batch_graph].unsqueeze(-1), probs[batch_graph].unsqueeze(-1))).sum()/2 ###calculate loss terms self_sums = scatter_add((probs*probs), batch, 0, dim_size = num_graphs) expected_weight_G = scatter_add(probs[no_loop_row]*probs[no_loop_col], batch[no_loop_row], 0, dim_size = num_graphs)/2. expected_clique_weight = (pairwise_prodsums.unsqueeze(-1) - self_sums)/1. expected_distance = (expected_clique_weight - expected_weight_G) ###useful numbers max_set_weight = (scatter_add(torch.ones_like(x)[no_loop_row], batch[no_loop_row], 0, dim_size = num_graphs)/2).squeeze(-1) set_weight = (scatter_add(x2[no_loop_row]*x2[no_loop_col], batch[no_loop_row], 0, dim_size = num_graphs)/2)+1e-6 clique_edges_hard = (set_size*(set_size-1)/2) +1e-6 clique_dist_hard = set_weight/clique_edges_hard clique_check = ((clique_edges_hard != clique_edges_hard)) setedge_check = ((set_weight != set_weight)) assert ((clique_dist_hard>=1.1).sum())<=1e-6, "Invalid set vol/clique vol ratio." ###calculate loss expected_loss = (penalty_coefficient)*expected_distance*0.5 - 0.5*expected_weight_G loss = expected_loss retdict = {} retdict["output"] = [probs.squeeze(-1),"hist"] #output retdict["Expected_cardinality"] = [card_1.mean(),"sequence"] retdict["Expected_cardinality_hist"] = [card_1,"hist"] retdict["losses histogram"] = [loss.squeeze(-1),"hist"] retdict["Set sizes"] = [set_size.squeeze(-1),"hist"] retdict["volume_hard"] = [vol_hard.mean(),"aux"] #volume2 retdict["cardinality_hard"] = [set_size[0],"sequence"] #volumeq retdict["Expected weight(G)"]= [expected_weight_G.mean(), "sequence"] retdict["Expected maximum weight"] = [expected_clique_weight.mean(),"sequence"] retdict["Expected distance"]= [expected_distance.mean(), "sequence"] retdict["Currvol/Cliquevol"] = [clique_dist_hard.mean(),'sequence'] retdict["Currvol/Cliquevol all graphs in batch"] = [clique_dist_hard.squeeze(-1),'hist'] retdict["Average ratio of total volume"]= [total_vol_ratio.mean(),'sequence'] retdict["cardinalities"] = [cardinalities.squeeze(-1),"hist"] retdict["Current edge percentage"] = [torch.tensor(current_edge_percentage),'sequence'] retdict["loss"] = [loss.mean().squeeze(),"sequence"] #final loss return retdict def __repr__(self): return self.__class__.__name__
class clique_MPNN(torch.nn.Module): def __init__(self, dataset, num_layers, hidden1, hidden2, deltas): super(clique_MPNN, self).__init__() self.hidden1 = hidden1 self.hidden2 = hidden2 self.momentum = 0.1 self.convs = torch.nn.ModuleList() self.deltas = deltas self.numlayers = num_layers self.heads = 8 self.concat = True self.bns = torch.nn.ModuleList() for i in range(num_layers - 1): self.bns.append( BN(self.heads * self.hidden1, momentum=self.momentum)) self.convs = torch.nn.ModuleList() for i in range(num_layers - 1): self.convs.append( GINConv(Sequential( Linear(self.heads * self.hidden1, self.heads * self.hidden1), ReLU(), Linear(self.heads * self.hidden1, self.heads * self.hidden1), ReLU(), BN(self.heads * self.hidden1, momentum=self.momentum), ), train_eps=True)) self.bn1 = BN(self.heads * self.hidden1) self.conv1 = GINConv(Sequential( Linear(self.hidden2, self.heads * self.hidden1), ReLU(), Linear(self.heads * self.hidden1, self.heads * self.hidden1), ReLU(), BN(self.heads * self.hidden1, momentum=self.momentum), ), train_eps=True) if self.concat: self.lin1 = Linear(self.heads * self.hidden1, self.hidden1) else: self.lin1 = Linear(self.hidden1, self.hidden1) self.lin2 = Linear(self.hidden1, 1) self.gnorm = GraphSizeNorm() def reset_parameters(self): self.conv1.reset_parameters() for conv in self.convs: conv.reset_parameters() for bn in self.bns: bn.reset_parameters() self.bn1.reset_parameters() self.lin1.reset_parameters() self.lin2.reset_parameters() def forward(self, data, edge_dropout=None, penalty_coefficient=0.25): x = data.x edge_index = data.edge_index batch = data.batch num_graphs = batch.max().item() + 1 row, col = edge_index total_num_edges = edge_index.shape[1] N_size = x.shape[0] if edge_dropout is not None: edge_index = dropout_adj( edge_index, edge_attr=(torch.ones(edge_index.shape[1], device=device)).long(), p=edge_dropout, force_undirected=True)[0] edge_index = add_remaining_self_loops(edge_index, num_nodes=batch.shape[0])[0] reduced_num_edges = edge_index.shape[1] current_edge_percentage = (reduced_num_edges / total_num_edges) no_loop_index, _ = remove_self_loops(edge_index) no_loop_row, no_loop_col = no_loop_index xinit = x.clone() x = x.unsqueeze(-1) mask = get_mask(x, edge_index, 1).to(x.dtype) x = F.leaky_relu(self.conv1(x, edge_index)) # +x x = x * mask x = self.gnorm(x) x = self.bn1(x) for conv, bn in zip(self.convs, self.bns): if (x.dim() > 1): x = x + F.leaky_relu(conv(x, edge_index)) mask = get_mask(mask, edge_index, 1).to(x.dtype) x = x * mask x = self.gnorm(x) x = bn(x) xpostconvs = x.detach() # x = F.leaky_relu(self.lin1(x)) x = x * mask xpostlin1 = x.detach() x = F.leaky_relu(self.lin2(x)) x = x * mask #calculate min and max batch_max = scatter_max(x, batch, 0, dim_size=N_size)[0] batch_max = torch.index_select(batch_max, 0, batch) batch_min = scatter_min(x, batch, 0, dim_size=N_size)[0] batch_min = torch.index_select(batch_min, 0, batch) #min-max normalize x = (x - batch_min) / (batch_max + 1e-6 - batch_min) probs = x #calculating the terms for the expected distance between clique and graph pairwise_prodsums = torch.zeros(num_graphs, device=device) for graph in range(num_graphs): batch_graph = (batch == graph) pairwise_prodsums[graph] = (torch.conv1d( probs[batch_graph].unsqueeze(-1), probs[batch_graph].unsqueeze(-1))).sum() / 2 ###calculate loss terms self_sums = scatter_add((probs * probs), batch, 0, dim_size=num_graphs) expected_weight_G = scatter_add( probs[no_loop_row] * probs[no_loop_col], batch[no_loop_row], 0, dim_size=num_graphs) / 2. expected_clique_weight = (pairwise_prodsums.unsqueeze(-1) - self_sums) / 1. expected_distance = (expected_clique_weight - expected_weight_G) ###calculate loss expected_loss = (penalty_coefficient ) * expected_distance * 0.5 - 0.5 * expected_weight_G loss = expected_loss retdict = {} retdict["output"] = [probs.squeeze(-1), "hist"] #output retdict["losses histogram"] = [loss.squeeze(-1), "hist"] retdict["Expected weight(G)"] = [expected_weight_G.mean(), "sequence"] retdict["Expected maximum weight"] = [ expected_clique_weight.mean(), "sequence" ] retdict["Expected distance"] = [expected_distance.mean(), "sequence"] retdict["loss"] = [loss.mean().squeeze(), "sequence"] #final loss return retdict def __repr__(self): return self.__class__.__name__
class GIN(torch.nn.Module): def __init__(self, num_layers, num_input_features, hidden): super(GIN, self).__init__() self.conv1 = GINConv( torch.nn.Sequential(Linear(num_input_features, hidden), torch.nn.ReLU(), Linear(hidden, hidden))) self.convs = torch.nn.ModuleList() for i in range(num_layers - 1): self.convs.append( GINConv( torch.nn.Sequential(Linear(hidden, hidden), torch.nn.ReLU(), Linear(hidden, hidden)))) self.lin1 = torch.nn.Linear(3 * hidden, hidden) self.lin2 = torch.nn.Linear(hidden, 2) def reset_parameters(self): # reset all conv and linear layers self.conv1.reset_parameters() for conv in self.convs: conv.reset_parameters( ) # .reset_parameters() is method of the torch_geometric.nn.GINConv class self.lin1.reset_parameters( ) # .reset_parameters() is method of the torch.nn.Linear class self.lin2.reset_parameters() def forward(self, data): # data: Batch(batch=[num_nodes_in_batch], # edge_attr=[2*num_nodes_in_batch,num_edge_features_per_edge], # edge_index=[2,2*num_nodes_in_batch], # pos=[num_nodes_in_batch,2], # x=[num_nodes_in_batch, num_input_features_per_node], # y=[num_graphs_in_batch, num_classes] # example: Batch(batch=[2490], edge_attr=[4980,1], edge_index=[2,4980], pos=[2490,2], x=[2490,33], y=[32,2] x, edge_index, batch = data.x, data.edge_index, data.batch # x.shape: torch.Size([num_nodes_in_batch, num_input_features_per_node]) # edge_index.shape: torch.Size([2, 2*num_nodes_in_batch]) # batch.shape: torch.Size([num_nodes_in_batch]) # example: x.shape = troch.Size([2490,33]) # edge_index.shape = torch.Size([2,4980]) # batch.shape = torch.Size([2490]) # graph convolutions and relu activation x = F.relu(self.conv1(x, edge_index)) # x.shape: torch.Size([num_nodes_in_batch, hidden]) # example: x.shape = torch.Size([2490, 66]) for conv in self.convs: x = F.relu(conv(x, edge_index)) # x.shape: torch.Size([num_nodes_in_batch, hidden]) # example: x.shape = torch.Size([2490, 66]) x = torch.cat([ global_add_pool(x, batch), global_mean_pool(x, batch), global_max_pool(x, batch) ], dim=1) # x.shape: torch.Size([num_graphs_in_batch, 3*hidden) # example: x.shape = torch.Size([32, 3*66]) # linear layers, activation function, dropout x = F.relu(self.lin1(x)) # x.shape: torch.Size([num_graphs_in_batch, hidden) # example: x.shape = torch.Size([32, 66]) x = F.dropout(x, p=0.5, training=self.training) x = self.lin2(x) # x.shape: torch.Size([num_graphs_in_batch, num_classes) # example: x.shape = torch.Size([32, 2]) output = F.log_softmax(x, dim=-1) return output def __repr__(self): #for getting a printable representation of an object return self.__class__.__name__