def __init__(self,in_channels, out_channels1, out_channels2,out_channels3,out_channels4, out_channels5, dropout,batch_size):
super(GraphPooling, self).__init__()
self.batch_size=batch_size
self.out_channels2=out_channels2
"""
Encoding
"""
### Encoding
self.sage1=tnn.SAGEConv(in_channels,out_channels1,normalize=False)
self.sage2=tnn.SAGEConv(out_channels1,out_channels2,normalize=False)
self.sage3=tnn.SAGEConv(out_channels2,out_channels3,normalize=False)
self.sage4=tnn.SAGEConv(out_channels3,out_channels4,normalize=False)
self.tr1=nn.Linear(out_channels4,out_channels5)
self.tr2=nn.Linear(out_channels5,1)
self.drop=torch.nn.Dropout(p=0.2)
## Batch Normalization
self.bano1 = nn.BatchNorm1d(num_features=out_channels1)
self.bano2 = nn.BatchNorm1d(num_features=out_channels2)
self.bano3 = nn.BatchNorm1d(num_features=out_channels3)
self.bano4 = nn.BatchNorm1d(num_features=out_channels4)
self.bano5 = nn.BatchNorm1d(num_features=out_channels5)
self.edge1=EdgePooling(out_channels1, edge_score_method=None, dropout=0., add_to_edge_score=0.5)
self.edge2=EdgePooling(out_channels2, edge_score_method=None, dropout=0., add_to_edge_score=0.5)
self.edge3=EdgePooling(out_channels3, edge_score_method=None, dropout=0., add_to_edge_score=0.5)
def __init__(self, feat_in, hidden_features=10, output_dim=1, dropout=0.1):
"""
A class for graph nodes classification. Contains graph convolution layers
:param feat_in: Number of features of one node
:param hidden_features: Number of features between layers
:param dropout: dropout value (after each layer)
"""
super(TreeSupport, self).__init__()
self.dropout = dropout
self.conv1 = gnn.GINConv(
nn.Sequential(nn.Linear(feat_in, hidden_features), nn.LeakyReLU(),
nn.Linear(hidden_features, hidden_features)),
train_eps=True
)
self.conv2 = gnn.GINConv(
nn.Sequential(nn.Linear(hidden_features, max(hidden_features // 2, 1)), nn.LeakyReLU(),
nn.Linear(max(hidden_features // 2, 1), max(hidden_features // 2, 1))),
train_eps=True
)
self.conv3 = gnn.GINConv(
nn.Sequential(nn.Linear(max(hidden_features // 2, 1), max(hidden_features // 2, 1)), nn.LeakyReLU(),
nn.Linear(max(hidden_features // 2, 1), max(hidden_features // 2, 1))),
train_eps=True
)
self.conv4 = gnn.SAGEConv(max(hidden_features // 2, 1), max(hidden_features // 2, 1))
self.conv5 = gnn.SAGEConv(max(hidden_features // 2, 1), output_dim)
def __init__(self, ppi_adj, g2v_embedding, args):
super(GEX_PPI_SAGE_cat4_MLP, self).__init__()
print('num genes : '+str(args.num_genes))
self.bn1 = nn.BatchNorm1d(args.ecfp_nBits)
self.bn2 = nn.BatchNorm1d(args.num_genes)
self.bn3 = nn.BatchNorm1d(1)
self.bn4 = nn.BatchNorm1d(1)
self.drug_mlp1 = nn.Linear(args.ecfp_nBits,
int(args.ecfp_nBits*2/3 + args.drug_embed_dim/3), bias = True)
self.drug_mlp2 = nn.Linear(int(args.ecfp_nBits*2/3 + args.drug_embed_dim/3),
int(args.ecfp_nBits/3+args.drug_embed_dim*2/3), bias = True)
self.drug_mlp3 = nn.Linear(int(args.ecfp_nBits/3+args.drug_embed_dim*2/3),
args.drug_embed_dim, bias = True)
self.gcns = {}
for i in range(args.num_gcn_hops):
if i==0:
self.gcns[i] = geo_nn.SAGEConv(args.gene2vec_dim, args.gcn_hidden_dim1*4)
else:
self.gcns[i] = geo_nn.SAGEConv(args.gcn_hidden_dim1*4, args.gcn_hidden_dim1*4)
for i in range(args.num_gcn_hops):
self.add_module('gcn_{}'.format(i), self.gcns[i])
# self.pred_emb_dim = args.drug_embed_dim + (args.num_gcn_hops*args.gcn_hidden_dim1*args.gat_num_heads) + 2
self.pred_emb_dim = args.drug_embed_dim + (args.num_gcn_hops*args.gcn_hidden_dim1*args.gat_num_heads) + 2
self.pred_mlp1 = nn.Linear(self.pred_emb_dim, int(self.pred_emb_dim*2/3), bias = True)
self.pred_mlp2 = nn.Linear(int(self.pred_emb_dim*2/3), int(self.pred_emb_dim/3), bias = True)
self.pred_mlp3 = nn.Linear(int(self.pred_emb_dim/3), args.num_classes, bias = True)
self.activ = nn.ReLU()
self.activ2 = nn.Softmax()
#
self.g2v_embeddings = nn.Embedding(args.num_genes, args.gene2vec_dim)
if args.g2v_pretrained == True:
g2v_embedding = F.normalize(torch.from_numpy(g2v_embedding), p = 2)
self.g2v_embeddings.weight.data.copy_(g2v_embedding)
self.ppi_adj = ppi_adj # 2 x num edges
# Read out MLP
self.readout_mlp1 = nn.Linear(args.num_genes,
int(args.num_genes*2/3), bias = True)
self.readout_mlp2 = nn.Linear(int(args.num_genes*2/3),
int(args.num_genes/3), bias = True)
self.readout_mlp3 = nn.Linear(int(args.num_genes/3), 1, bias = True)
def _setup_layers(self):
dims = []
# GNN layers
graph_layers = []
for layer in self.config['graph_layers']:
if type(layer) == int:
dim = dims[-1] if len(dims) != 0 else self.config['n_features']
graph_layers.append(geo_nn.SAGEConv(dim, layer))
dims.append(layer)
elif type(layer) == str and layer in self.activations:
graph_layers.append(self.activations[layer]())
# Fully connected layers
fc_layers = []
for i, layer in enumerate(self.config['fc_layers']):
if type(layer) == int:
if type(self.config['fc_layers'][i - 1]) == int and i != 0:
fc_layers.append(NN.ReLU())
continue
dim = dims[-1] if len(dims) != 0 else self.config['n_features']
fc_layers.append(NN.Linear(dim, layer))
dims.append(layer)
elif type(layer) == str and layer in self.activations:
fc_layers.append(self.activations[layer]())
elif layer == 'drop':
fc_layers.append(NN.Dropout(.3))
# Output layer
fc_layers.append(NN.Linear(dims[-1], config['n_labels']))
return NN.ModuleList(graph_layers), NN.ModuleList(fc_layers)
def __init__(self, \
n_features, \
n_classes, \
n_hidden_GNN=[10], \
n_hidden_FC=[], \
dropout_GNN=0, \
dropout_FC=0):
super(GraphSAGE, self).__init__(\
n_features, n_classes, n_hidden_GNN,\
n_hidden_FC, dropout_FC, dropout_GNN)
self.layers_GNN.append(pyg_nn.SAGEConv(1, n_hidden_GNN[0]))
if self.n_layers_GNN > 1:
for i in range(self.n_layers_GNN - 1):
self.layers_GNN.append(
pyg_nn.SAGEConv(n_hidden_GNN[i], n_hidden_GNN[(i + 1)]))
def __init__(self, feat_in, n_classes=3, max_nodes=45):
"""
Initializes class for alignment classification
:param feat_in: the second dimension size of node feature matrix
:param n_classes: the number of classes in dataset
:param max_nodes: maximum number of nodes in a graph
"""
super(AlignSAGE, self).__init__()
first_hidden = max(feat_in // 2, 1)
second_hidden = max(feat_in // 4, 1)
self.N = max_nodes
self.sage1 = gnn.SAGEConv(feat_in, first_hidden)
self.pooling1 = gnn.SAGPooling(first_hidden)
self.sage2 = gnn.SAGEConv(first_hidden, second_hidden)
self.pooling2 = gnn.SAGPooling(second_hidden)
self.lin1 = nn.Linear(10, 100)
self.lin2 = nn.Linear(100, n_classes)
def build_conv_model(self, model_type, node_in_dim, node_out_dim, edge_dim,
edge_mode, normalize_emb, activation, aggr):
if model_type == 'GCN':
return pyg_nn.GCNConv(node_in_dim, node_out_dim)
elif model_type == 'GraphSage':
return pyg_nn.SAGEConv(node_in_dim, node_out_dim)
elif model_type == 'GAT':
return pyg_nn.GATConv(node_in_dim, node_out_dim)
elif model_type == 'EGCN':
return EGCNConv(node_in_dim, node_out_dim, edge_dim, edge_mode)
elif model_type == 'EGSAGE':
return EGraphSage(node_in_dim, node_out_dim, edge_dim, activation,
edge_mode, normalize_emb, aggr)
def build_conv_model(conv_type, node_in_dim, node_out_dim):
if conv_type == 'GSage':
return pyg_nn.SAGEConv(node_in_dim, node_out_dim)
elif conv_type == 'GCN':
return pyg_nn.GCNConv(node_in_dim, node_out_dim)
elif conv_type == 'GAT':
return pyg_nn.GATConv(node_in_dim, node_out_dim)
elif conv_type == 'GSage2':
from conv_layers import GraphSage2
return GraphSage2(node_in_dim, node_out_dim)
elif conv_type == 'GSageW':
from conv_layers import GraphSageW
return GraphSageW(node_in_dim, node_out_dim)
else:
raise NotImplementedError
def build_gnn_conv_model(conv_type, node_in_dim, node_out_dim):
if conv_type == 'GSage':
return pyg_nn.SAGEConv(node_in_dim,
node_out_dim,
normalize=False,
bias=True)
elif conv_type == 'GCN':
return pyg_nn.GCNConv(node_in_dim,
node_out_dim,
add_self_loops=False,
normalize=True,
bias=True)
elif conv_type == 'GAT':
return pyg_nn.GATConv(node_in_dim,
node_out_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=False,
bias=True)
elif conv_type == 'Tran':
gnn = pyg_nn.TransformerConv(node_in_dim,
node_out_dim,
heads=1,
concat=True,
dropout=0.,
bias=True)
gnn.lin_edge = None # remove redundant parameters
return gnn
elif conv_type == 'GSage2':
from rlkit.torch.networks.conv_layers import GraphSage2
return GraphSage2(node_in_dim, node_out_dim)
elif conv_type == 'GSageW':
from rlkit.torch.networks.conv_layers import GraphSageW
return GraphSageW(node_in_dim, node_out_dim)
else:
raise NotImplementedError
def __init__(self,
input_dim,
hidden_dim,
output_dim,
seq_len,
head_num,
qs_graph_dir,
device,
dropout,
n_hop,
gcn_type,
gcn_layer_num,
gcn_on,
pretrained_embedding=None,
freeze=True):
"""[summary]
Args:
input_dim ([type]): [description]
hidden_dim ([type]): [description]
output_dim ([type]): output dim of GCN
seq_len ([type]): [description]
head_num ([type]): [description]
qs_graph_dir ([type]): [description]
"""
super(Model, self).__init__()
self.device = device
self.gcn_on = gcn_on
if pretrained_embedding is not None:
self.pretrained_embedding = pretrained_embedding
else:
self.pretrained_embedding = None
if gcn_type == 'sgconv' and gcn_on:
self.gcn_layer_num = gcn_layer_num
self.convs = nn.ModuleList()
if gcn_layer_num == 1:
self.convs.append(pyg_nn.SGConv(input_dim, output_dim,
K=n_hop))
elif gcn_layer_num > 1:
self.convs.append(pyg_nn.SGConv(input_dim, hidden_dim,
K=n_hop))
else:
raise ValueError("Unsupported gcn_layer_num {}")
for i in range(self.gcn_layer_num - 1):
if i == self.gcn_layer_num - 2:
self.convs.append(
pyg_nn.SGConv(hidden_dim, output_dim, K=n_hop))
else:
self.convs.append(
pyg_nn.SGConv(hidden_dim, hidden_dim, K=n_hop))
elif gcn_type == 'gconv' and gcn_on:
self.gcn_layer_num = gcn_layer_num
self.convs = nn.ModuleList()
if gcn_layer_num == 1:
self.convs.append(pyg_nn.GCNConv(input_dim, output_dim))
elif gcn_layer_num > 1:
self.convs.append(pyg_nn.GCNConv(input_dim, hidden_dim))
else:
raise ValueError("Unsupported gcn_layer_num {}")
for i in range(self.gcn_layer_num - 1):
if i == self.gcn_layer_num - 2:
self.convs.append(pyg_nn.GCNConv(hidden_dim, output_dim))
else:
self.convs.append(pyg_nn.GCNConv(hidden_dim, hidden_dim))
elif gcn_type == 'gat' and gcn_on:
self.gcn_layer_num = gcn_layer_num
self.convs = nn.ModuleList()
if gcn_layer_num == 1:
self.convs.append(
pyg_nn.GATConv(input_dim,
output_dim,
heads=5,
concat=False))
elif gcn_layer_num > 1:
self.convs.append(
pyg_nn.GATConv(input_dim,
hidden_dim,
heads=5,
concat=False))
else:
raise ValueError("Unsupported gcn_layer_num {}")
for i in range(self.gcn_layer_num - 1):
if i == self.gcn_layer_num - 2:
self.convs.append(pyg_nn.GATConv(hidden_dim, output_dim))
else:
self.convs.append(pyg_nn.GATConv(hidden_dim, hidden_dim))
elif gcn_type == 'sage' and gcn_on:
self.gcn_layer_num = gcn_layer_num
self.convs = nn.ModuleList()
if gcn_layer_num == 1:
self.convs.append(pyg_nn.SAGEConv(input_dim, output_dim))
elif gcn_layer_num > 1:
self.convs.append(pyg_nn.SAGEConv(input_dim, hidden_dim))
else:
raise ValueError("Unsupported gcn_layer_num {}")
for i in range(self.gcn_layer_num - 1):
if i == self.gcn_layer_num - 2:
self.convs.append(pyg_nn.SAGEConv(hidden_dim, output_dim))
else:
self.convs.append(pyg_nn.SAGEConv(hidden_dim, hidden_dim))
self.dropout = dropout
with open(qs_graph_dir, "r") as src:
self.qs_graph = json.load(src)
# ! input shape should be (seq_len - 1)
self.seq_len = seq_len
self.correctness_embedding_layer = nn.Embedding(2, input_dim)
if self.pretrained_embedding is None:
self.node_embedding_layer = nn.Embedding(len(self.qs_graph),
input_dim)
self.linears = nn.ModuleList()
self.linears.append(nn.Linear(output_dim * 2, output_dim, bias=True))
for _ in range(3):
self.linears.append(nn.Linear(output_dim, output_dim, bias=False))
self.linears.append(nn.Linear(output_dim, 1))
self.MHA = nn.MultiheadAttention(embed_dim=output_dim,
num_heads=head_num,
dropout=self.dropout[1])
self.FFN = nn.ModuleList()
for _ in range(2):
self.FFN.append(nn.Linear(output_dim, output_dim))
self.lns = nn.ModuleList()
self.lns.append(nn.LayerNorm(hidden_dim))
self.lns.append(nn.LayerNorm(hidden_dim))
self.lns.append(nn.LayerNorm(output_dim))
self.lns.append(nn.LayerNorm(output_dim))
self.pos_embedding = nn.Embedding(seq_len - 1, output_dim)
self.dropout_layers = nn.ModuleList()
self.dropout_layers.append(nn.Dropout(p=self.dropout[0]))
self.dropout_layers.append(nn.Dropout(p=self.dropout[2]))
def build_conv_model(self, node_in_dim, node_out_dim):
return pyg_nn.SAGEConv(node_in_dim,node_out_dim)
def __init__(
self,
pre_graph_builder,
obs_dim,
action_dim,
node_dim,
conv_type,
use_attention=True,
normalize_emb=False,
num_layer=1,
output_activation=None,
):
super(GraphContextNet, self).__init__()
self.obs_dim = obs_dim
self.action_dim = action_dim
self.node_dim = node_dim
self.conv_type = conv_type
self.num_layer = num_layer
# graph builder
self.pre_graph_builder = pre_graph_builder
# convs
if self.conv_type == 'GContext':
self.conv = ContextConv(
obs_dim,
action_dim,
node_dim,
use_attention=use_attention,
num_layer=num_layer,
normalize_emb=normalize_emb,
negative_slope=0.2,
)
elif self.conv_type == 'GAT':
self.pre_lin_l = Mlp(
input_size=obs_dim + action_dim,
output_size=node_dim,
hidden_sizes=[node_dim] * num_layer,
hidden_activation=nn.LeakyReLU(0.2),
output_activation=nn.LeakyReLU(0.2),
)
self.pre_lin_r = Mlp(
input_size=obs_dim,
output_size=node_dim,
hidden_sizes=[node_dim] * num_layer,
hidden_activation=nn.LeakyReLU(0.2),
output_activation=nn.LeakyReLU(0.2),
)
self.conv = pyg_nn.GATConv(
in_channels=[node_dim, node_dim],
out_channels=node_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=False,
bias=True,
)
self.aft_lin = nn.Linear(node_dim + node_dim, node_dim)
elif self.conv_type in ['GAT2', 'GSage2']:
self.convs = nn.ModuleList()
self.mlps = nn.ModuleList()
if self.conv_type == 'GAT2':
self.convs.append(
pyg_nn.GATConv(
in_channels=[obs_dim + action_dim, obs_dim],
out_channels=node_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=False,
bias=True,
))
elif self.conv_type == 'GSage2':
self.convs.append(
pyg_nn.SAGEConv(
in_channels=[obs_dim + action_dim, obs_dim],
out_channels=node_dim,
normalize=normalize_emb,
bias=True,
))
self.mlps.append(nn.Linear(obs_dim + action_dim, node_dim))
for l in range(1, num_layer):
if self.conv_type == 'GAT2':
self.convs.append(
pyg_nn.GATConv(
in_channels=[node_dim, node_dim],
out_channels=node_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=False,
bias=True,
))
elif self.conv_type == 'GSage2':
self.convs.append(
pyg_nn.SAGEConv(
in_channels=[node_dim, node_dim],
out_channels=node_dim,
normalize=normalize_emb,
bias=True,
))
if l < num_layer - 1:
self.mlps.append(nn.Linear(node_dim, node_dim))
elif self.conv_type in ['GSage3', 'GAT3']:
self.obs_convs = nn.ModuleList()
self.action_convs = nn.ModuleList()
self.mlps = nn.ModuleList()
if self.conv_type == 'GSage3':
self.obs_convs.append(
pyg_nn.SAGEConv(
in_channels=obs_dim,
out_channels=node_dim,
normalize=normalize_emb,
bias=True,
))
self.action_convs.append(
pyg_nn.SAGEConv(
in_channels=[node_dim + action_dim, obs_dim],
out_channels=node_dim,
normalize=normalize_emb,
bias=True,
))
for l in range(1, num_layer):
self.obs_convs.append(
pyg_nn.SAGEConv(
in_channels=node_dim,
out_channels=node_dim,
normalize=normalize_emb,
bias=True,
))
self.action_convs.append(
pyg_nn.SAGEConv(
in_channels=[node_dim + action_dim, node_dim],
out_channels=node_dim,
normalize=normalize_emb,
bias=True,
))
elif self.conv_type == 'GAT3':
self.obs_convs.append(
pyg_nn.GATConv(
in_channels=obs_dim,
out_channels=node_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=True,
bias=True,
))
self.action_convs.append(
pyg_nn.GATConv(
in_channels=[node_dim + action_dim, obs_dim],
out_channels=node_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=False,
bias=True,
))
self.mlps.append(nn.Linear(obs_dim + node_dim, node_dim))
for l in range(1, num_layer):
self.obs_convs.append(
pyg_nn.GATConv(
in_channels=node_dim,
out_channels=node_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=True,
bias=True,
))
self.action_convs.append(
pyg_nn.GATConv(
in_channels=[node_dim + action_dim, node_dim],
out_channels=node_dim,
heads=1,
concat=True,
negative_slope=0.2,
dropout=0.,
add_self_loops=False,
bias=True,
))
self.mlps.append(nn.Linear(node_dim + node_dim, node_dim))
self.output_activation = get_activation(output_activation)
def __init__(self, in_c, out_c, hid_c=64):
super(graphSAGE, self).__init__()
self.conv1 = pyg_nn.SAGEConv(in_c, hid_c)
self.conv2 = pyg_nn.SAGEConv(hid_c, out_c)
