def __init__(self, num_tasks = 1, num_layers = 5, emb_dim = 300, gnn_type = 'gin', virtual_node = True, residual = False, drop_ratio = 0, JK = "last", graph_pooling = "sum"): ''' num_tasks (int): number of labels to be predicted virtual_node (bool): whether to add virtual node or not ''' super(GNN, self).__init__() self.num_layers = num_layers self.drop_ratio = drop_ratio self.JK = JK self.emb_dim = emb_dim self.num_tasks = num_tasks self.graph_pooling = graph_pooling if self.num_layers < 2: raise ValueError("Number of GNN layers must be greater than 1.") ### GNN to generate node embeddings if virtual_node: self.gnn_node = GNN_node_Virtualnode(num_layers, emb_dim, JK = JK, drop_ratio = drop_ratio, residual = residual, gnn_type = gnn_type) else: self.gnn_node = GNN_node(num_layers, emb_dim, JK = JK, drop_ratio = drop_ratio, residual = residual, gnn_type = gnn_type) ### Pooling function to generate whole-graph embeddings if self.graph_pooling == "sum": self.pool = SumPooling() elif self.graph_pooling == "mean": self.pool = AvgPooling() elif self.graph_pooling == "max": self.pool = MaxPooling elif self.graph_pooling == "attention": self.pool = GlobalAttentionPooling( gate_nn = nn.Sequential(nn.Linear(emb_dim, 2*emb_dim), nn.BatchNorm1d(2*emb_dim), nn.ReLU(), nn.Linear(2*emb_dim, 1))) elif self.graph_pooling == "set2set": self.pool = Set2Set(emb_dim, n_iters = 2, n_layers = 2) else: raise ValueError("Invalid graph pooling type.") if graph_pooling == "set2set": self.graph_pred_linear = nn.Linear(2*self.emb_dim, self.num_tasks) else: self.graph_pred_linear = nn.Linear(self.emb_dim, self.num_tasks)
def __init__(self, num_tasks=1, num_layers=5, emb_dim=300, gnn_type='gin', virtual_node=True, residual=False, drop_ratio=0, JK="last", graph_pooling="sum"): ''' num_tasks (int): number of labels to be predicted virtual_node (bool): whether to add virtual node or not ''' super(GNN, self).__init__() self.num_layers = num_layers self.drop_ratio = drop_ratio self.JK = JK self.emb_dim = emb_dim self.num_tasks = num_tasks self.graph_pooling = graph_pooling if self.num_layers < 2: raise ValueError("Number of GNN layers must be greater than 1.") ### GNN to generate node embeddings if virtual_node: self.gnn_node = GNN_node_Virtualnode(num_layers, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) else: self.gnn_node = GNN_node(num_layers, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) ### Pooling function to generate whole-graph embeddings if self.graph_pooling == "sum": self.pool = gnn.GraphPool(pool_type="sum") elif self.graph_pooling == "mean": self.pool = gnn.GraphPool(pool_type="mean") elif self.graph_pooling == "max": self.pool = gnn.GraphPool(pool_type="max") else: raise ValueError("Invalid graph pooling type.") self.graph_pred_linear = nn.Linear(self.emb_dim, self.num_tasks)
def __init__(self, num_vocab, max_seq_len, node_encoder, num_layer = 5, emb_dim = 300, gnn_type = 'gin', virtual_node = True, residual = False, drop_ratio = 0.5, JK = "last", graph_pooling = "mean"): ''' num_tasks (int): number of labels to be predicted virtual_node (bool): whether to add virtual node or not ''' super(GNN, self).__init__() self.num_layer = num_layer self.drop_ratio = drop_ratio self.JK = JK self.emb_dim = emb_dim self.num_vocab = num_vocab self.max_seq_len = max_seq_len self.graph_pooling = graph_pooling if self.num_layer < 2: raise ValueError("Number of GNN layers must be greater than 1.") ### GNN to generate node embeddings if virtual_node: self.gnn_node = GNN_node_Virtualnode(num_layer, emb_dim, node_encoder, JK = JK, drop_ratio = drop_ratio, residual = residual, gnn_type = gnn_type) else: self.gnn_node = GNN_node(num_layer, emb_dim, node_encoder, JK = JK, drop_ratio = drop_ratio, residual = residual, gnn_type = gnn_type) ### Pooling function to generate whole-graph embeddings if self.graph_pooling == "sum": self.pool = global_add_pool elif self.graph_pooling == "mean": self.pool = global_mean_pool elif self.graph_pooling == "max": self.pool = global_max_pool elif self.graph_pooling == "attention": self.pool = GlobalAttention(gate_nn = torch.nn.Sequential(torch.nn.Linear(emb_dim, 2*emb_dim), torch.nn.BatchNorm1d(2*emb_dim), torch.nn.ReLU(), torch.nn.Linear(2*emb_dim, 1))) elif self.graph_pooling == "set2set": self.pool = Set2Set(emb_dim, processing_steps = 2) else: raise ValueError("Invalid graph pooling type.") self.graph_pred_linear_list = torch.nn.ModuleList() if graph_pooling == "set2set": for i in range(max_seq_len): self.graph_pred_linear_list.append(torch.nn.Linear(2*emb_dim, self.num_vocab)) else: for i in range(max_seq_len): self.graph_pred_linear_list.append(torch.nn.Linear(emb_dim, self.num_vocab))
def __init__(self, num_tasks, num_layer=5, emb_dim=300, gnn_type='gin', virtual_node=True, residual=False, drop_ratio=0.5, drop_path_p=0.01, JK="last", graph_pooling="mean", net_linear=False, net_seed=47, edge_p=0.6): ''' num_tasks (int): number of labels to be predicted virtual_node (bool): whether to add virtual node or not ''' super(GNN, self).__init__() self.num_layer = num_layer self.drop_ratio = drop_ratio self.JK = JK self.emb_dim = emb_dim self.num_tasks = num_tasks self.graph_pooling = graph_pooling if self.num_layer < 2: raise ValueError("Number of GNN layers must be greater than 1.") ### GNN to generate node embeddings if gnn_type == 'randomgin': if virtual_node: self.gnn_node = RandomGNN_node_Virtualnode( num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type, drop_path_p=drop_path_p, net_linear=net_linear, net_seed=net_seed, edge_p=edge_p) else: self.gnn_node = RandomGNN_node(num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type, drop_path_p=drop_path_p, net_linear=net_linear, net_seed=net_seed, edge_p=edge_p) else: if virtual_node: self.gnn_node = GNN_node_Virtualnode(num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) else: self.gnn_node = GNN_node(num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) ### Pooling function to generate whole-graph embeddings if self.graph_pooling == "sum": self.pool = global_add_pool elif self.graph_pooling == "mean": self.pool = global_mean_pool elif self.graph_pooling == "max": self.pool = global_max_pool elif self.graph_pooling == "attention": self.pool = GlobalAttention(gate_nn=torch.nn.Sequential( torch.nn.Linear(emb_dim, 2 * emb_dim), torch.nn.BatchNorm1d(2 * emb_dim), torch.nn.ReLU(), torch.nn.Linear(2 * emb_dim, 1))) elif self.graph_pooling == "set2set": self.pool = Set2Set(emb_dim, processing_steps=2) else: raise ValueError("Invalid graph pooling type.") if graph_pooling == "set2set": self.graph_pred_linear = torch.nn.Linear(2 * self.emb_dim, self.num_tasks) else: self.graph_pred_linear = torch.nn.Linear(self.emb_dim, self.num_tasks)
def __init__(self, num_tasks, num_layer=5, emb_dim=300, gnn_type='gin', virtual_node=True, residual=False, drop_ratio=0.5, JK="last", graph_pooling="mean", transformers=False, controller=False): ''' num_tasks (int): number of labels to be predicted virtual_node (bool): whether to add virtual node or not ''' super(GNN, self).__init__() self.num_layer = num_layer self.drop_ratio = drop_ratio self.JK = JK self.emb_dim = emb_dim self.num_tasks = num_tasks self.graph_pooling = graph_pooling if transformers: self.transformer_layers = torch.nn.ModuleList([ torch.nn.TransformerEncoderLayer(d_model=emb_dim, dim_feedforward=emb_dim * 4, nhead=4) for _ in range(2) ]) if controller: self.controller_layers = torch.nn.ModuleList([ ControllerTransformer(depth=1, expansion_ratio=4, n_heads=4, s2g_sharing=True, in_features=300, out_features=1, set_fn_feats=[128, 128, 128, 128, 5], method='lin2', hidden_mlp=[128], predict_diagonal=False, attention=True) for _ in range(1) ]) self.transformers = transformers self.controller = controller if self.num_layer < 2: raise ValueError("Number of GNN layers must be greater than 1.") ### GNN to generate node embeddings if virtual_node: self.gnn_node = GNN_node_Virtualnode(num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) else: self.gnn_node = GNN_node(num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) ### Pooling function to generate whole-graph embeddings if self.graph_pooling == "sum": self.pool = global_add_pool elif self.graph_pooling == "mean": self.pool = global_mean_pool elif self.graph_pooling == "max": self.pool = global_max_pool elif self.graph_pooling == "attention": self.pool = GlobalAttention(gate_nn=torch.nn.Sequential( torch.nn.Linear(emb_dim, 2 * emb_dim), torch.nn.BatchNorm1d(2 * emb_dim), torch.nn.ReLU(), torch.nn.Linear(2 * emb_dim, 1))) elif self.graph_pooling == "set2set": self.pool = Set2Set(emb_dim, processing_steps=2) else: raise ValueError("Invalid graph pooling type.") if graph_pooling == "set2set": self.graph_pred_linear = torch.nn.Linear(2 * self.emb_dim, self.num_tasks) else: self.graph_pred_linear = torch.nn.Linear(self.emb_dim, self.num_tasks)
def __init__(self, num_tasks=1, num_layers=5, emb_dim=300, gnn_type='gin', virtual_node=True, residual=False, drop_ratio=0, JK="last", graph_pooling="sum", last_layer_only=True): super(BayesianGNN, self).__init__(num_tasks, num_layers, emb_dim, gnn_type, virtual_node, residual, drop_ratio, JK, graph_pooling) ### GNN to generate node embeddings if virtual_node & (not last_layer_only): self.gnn_node = Bayesian_GNN_node_Virtualnode( num_layers, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) # only last layer is replaced with Bayesian neural network elif virtual_node & last_layer_only: self.gnn_node = GNN_node_Virtualnode(num_layers, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type) else: raise Exception("not implemented") # KL-divergence loss for Bayesian Neural Network self.kl_loss = bnn.BKLLoss(reduction='mean', last_layer_only=False) self.kl_weight = 0.01 # change graph_pred_linear if graph_pooling == "set2set": embedding_dim = 2 * self.emb_dim else: embedding_dim = self.emb_dim self.graph_pred_linear = torch.nn.Sequential( bnn.BayesLinear(prior_mu=0, prior_sigma=0.1, in_features=embedding_dim, out_features=100), torch.nn.ReLU(), bnn.BayesLinear(prior_mu=0, prior_sigma=0.1, in_features=100, out_features=100), torch.nn.ReLU(), bnn.BayesLinear(prior_mu=0, prior_sigma=0.1, in_features=100, out_features=100), torch.nn.ReLU(), bnn.BayesLinear(prior_mu=0, prior_sigma=0.1, in_features=100, out_features=self.num_tasks), )
def __init__(self, num_tasks, num_layer=5, emb_dim=300, gnn_type='gin', virtual_node=True, residual=False, drop_ratio=0.5, JK="last", graph_pooling="mean", laf_fun='mean', laf_layers='false', device='cuda', lafgrad=True): ''' num_tasks (int): number of labels to be predicted virtual_node (bool): whether to add virtual node or not ''' super(GNN, self).__init__() self.num_layer = num_layer self.drop_ratio = drop_ratio self.JK = JK self.emb_dim = emb_dim self.num_tasks = num_tasks self.graph_pooling = graph_pooling if self.num_layer < 2: raise ValueError("Number of GNN layers must be greater than 1.") ### GNN to generate node embeddings if virtual_node: self.gnn_node = GNN_node_Virtualnode(num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type, laf_fun=laf_fun, laf_layers=laf_layers, device=device) else: self.gnn_node = GNN_node(num_layer, emb_dim, JK=JK, drop_ratio=drop_ratio, residual=residual, gnn_type=gnn_type, laf_fun=laf_fun, laf_layers=laf_layers, device=device) ### Pooling function to generate whole-graph embeddings if self.graph_pooling == "sum": self.pool = global_add_pool elif self.graph_pooling == "mean": self.pool = global_mean_pool elif self.graph_pooling == "max": self.pool = global_max_pool elif self.graph_pooling == "attention": self.pool = GlobalAttention(gate_nn=torch.nn.Sequential( torch.nn.Linear(emb_dim, 2 * emb_dim), torch.nn.BatchNorm1d(2 * emb_dim), torch.nn.ReLU(), torch.nn.Linear(2 * emb_dim, 1))) elif self.graph_pooling == "set2set": self.pool = Set2Set(emb_dim, processing_steps=2) elif self.graph_pooling == "laf": if laf_fun == 'exp': self.pool = ScatterExponentialLAF(device=device) else: self.pool = ScatterAggregationLayer(function=laf_fun, grad=lafgrad, device=device) else: raise ValueError("Invalid graph pooling type.") if graph_pooling == "set2set": self.graph_pred_linear = torch.nn.Linear(2 * self.emb_dim, self.num_tasks) else: self.graph_pred_linear = torch.nn.Linear(self.emb_dim, self.num_tasks)