def __init__(
self,
hidden_size=64,
num_layer=2,
readout="avg",
layernorm: bool = False,
set2set_lstm_layer: int = 3,
set2set_iter: int = 6,
):
super(UnsupervisedGCN, self).__init__()
self.layers = nn.ModuleList([
GCNLayer(
in_feats=hidden_size,
out_feats=hidden_size,
activation=F.relu if i + 1 < num_layer else None,
residual=False,
batchnorm=False,
dropout=0.0,
) for i in range(num_layer)
])
if readout == "avg":
self.readout = AvgPooling()
elif readout == "set2set":
self.readout = Set2Set(hidden_size,
n_iters=set2set_iter,
n_layers=set2set_lstm_layer)
self.linear = nn.Linear(2 * hidden_size, hidden_size)
elif readout == "root":
# HACK: process outside the model part
self.readout = lambda _, x: x
else:
raise NotImplementedError
self.layernorm = layernorm
if layernorm:
self.ln = nn.LayerNorm(hidden_size, elementwise_affine=False)
def __init__(
self,
node_input_dim=15,
edge_input_dim=5,
output_dim=12,
node_hidden_dim=64,
edge_hidden_dim=128,
num_step_message_passing=6,
num_step_set2set=6,
num_layer_set2set=3,
):
super(MPNNModel, self).__init__()
self.num_step_message_passing = num_step_message_passing
self.lin0 = nn.Linear(node_input_dim, node_hidden_dim)
edge_network = nn.Sequential(
nn.Linear(edge_input_dim, edge_hidden_dim),
nn.ReLU(),
nn.Linear(edge_hidden_dim, node_hidden_dim * node_hidden_dim),
)
self.conv = NNConv(
in_feats=node_hidden_dim,
out_feats=node_hidden_dim,
edge_func=edge_network,
aggregator_type="sum",
)
self.gru = nn.GRU(node_hidden_dim, node_hidden_dim)
self.set2set = Set2Set(node_hidden_dim, num_step_set2set,
num_layer_set2set)
self.lin1 = nn.Linear(2 * node_hidden_dim, node_hidden_dim)
self.lin2 = nn.Linear(node_hidden_dim, output_dim)
def __init__(self,
in_node_feats,
in_edge_feats,
node_hidden_dim,
edge_hidden_dim,
num_step_message_passing,
num_step_set2set,
num_layer_set2set,
n_tasks,
regressor_hidden_feats=128,
dropout=0.):
super(MPNNRegressorBypass,
self).__init__(readout_feats=4 * node_hidden_dim,
n_tasks=n_tasks,
regressor_hidden_feats=regressor_hidden_feats,
dropout=dropout)
self.shared_gnn = MPNNGNN(in_node_feats, in_edge_feats,
node_hidden_dim, edge_hidden_dim,
num_step_message_passing)
for _ in range(n_tasks):
self.task_gnns.append(
MPNNGNN(in_node_feats, in_edge_feats, node_hidden_dim,
edge_hidden_dim, num_step_message_passing))
self.readouts.append(
Set2Set(2 * node_hidden_dim, num_step_set2set,
num_layer_set2set))
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,
node_input_dim=42,
edge_input_dim=10,
node_hidden_dim=42,
edge_hidden_dim=42,
num_step_message_passing=6,
interaction='dot',
num_step_set2_set=2,
num_layer_set2set=1,
):
super(CIGINModel, self).__init__()
self.node_input_dim = node_input_dim
self.node_hidden_dim = node_hidden_dim
self.edge_input_dim = edge_input_dim
self.edge_hidden_dim = edge_hidden_dim
self.num_step_message_passing = num_step_message_passing
self.interaction = interaction
self.solute_gather = GatherModel(self.node_input_dim, self.edge_input_dim,
self.node_hidden_dim, self.edge_input_dim,
self.num_step_message_passing,
)
self.solvent_gather = GatherModel(self.node_input_dim, self.edge_input_dim,
self.node_hidden_dim, self.edge_input_dim,
self.num_step_message_passing,
)
self.fc1 = nn.Linear(8 * self.node_hidden_dim, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, 1)
self.imap = nn.Linear(80, 1)
self.num_step_set2set = num_step_set2_set
self.num_layer_set2set = num_layer_set2set
self.set2set_solute = Set2Set(2 * node_hidden_dim, self.num_step_set2set, self.num_layer_set2set)
self.set2set_solvent = Set2Set(2 * node_hidden_dim, self.num_step_set2set, self.num_layer_set2set)
def __init__(self,
node_in_feats,
edge_in_feats,
global_feats,
node_out_feats=64,
edge_hidden_feats=128,
global_hidden_feats=32,
n_tasks=1,
num_step_message_passing=6,
num_step_set2set=6,
num_layer_set2set=3,
output_f=None):
super(MPNNPredictor, self).__init__()
self.gnn = MPNNGNN(
node_in_feats=node_in_feats,
node_out_feats=node_out_feats,
edge_in_feats=edge_in_feats,
edge_hidden_feats=edge_hidden_feats,
num_step_message_passing=num_step_message_passing,
)
self.readout = Set2Set(
input_dim=node_out_feats,
n_iters=num_step_set2set,
n_layers=num_layer_set2set,
)
self.global_subnet = nn.Sequential(
nn.Linear(global_feats, global_hidden_feats),
nn.ReLU(),
nn.Linear(global_hidden_feats, global_hidden_feats),
nn.ReLU(),
)
self.predict = nn.Sequential(
nn.Linear(2 * node_out_feats + global_hidden_feats,
node_out_feats),
nn.ReLU(),
nn.Linear(node_out_feats, n_tasks),
)
self.output_f = output_f
def __init__(self,
in_node_feats,
in_edge_feats,
node_hidden_dim,
edge_hidden_dim,
num_step_message_passing,
num_step_set2set,
num_layer_set2set,
n_tasks,
regressor_hidden_feats=128,
dropout=0.):
super(MPNNRegressor,
self).__init__(readout_feats=2 * node_hidden_dim,
n_tasks=n_tasks,
regressor_hidden_feats=regressor_hidden_feats,
dropout=dropout)
self.gnn = MPNNGNN(in_node_feats, in_edge_feats, node_hidden_dim,
edge_hidden_dim, num_step_message_passing)
self.readout = Set2Set(node_hidden_dim, num_step_set2set,
num_layer_set2set)
def __init__(self,
node_input_dim=42,
edge_input_dim=10,
node_hidden_dim=42,
edge_hidden_dim=42,
num_step_message_passing=6,
):
super(GatherModel, self).__init__()
self.num_step_message_passing = num_step_message_passing
self.lin0 = nn.Linear(node_input_dim, node_hidden_dim)
self.set2set = Set2Set(node_hidden_dim, 2, 1)
self.message_layer = nn.Linear(2 * node_hidden_dim, node_hidden_dim)
edge_network = nn.Sequential(
nn.Linear(edge_input_dim, edge_hidden_dim), nn.ReLU(),
nn.Linear(edge_hidden_dim, node_hidden_dim * node_hidden_dim))
self.conv = NNConv(in_feats=node_hidden_dim,
out_feats=node_hidden_dim,
edge_func=edge_network,
aggregator_type='sum',
residual=True
)
def __init__(self, in_feats, out_feats, n_hidden, n_iter_readout):
super(GGNN, self).__init__()
self.in_feats = in_feats
self.out_feats = out_feats
self.hidden = n_hidden
self.conv1 = GatedGraphConv(in_feats=in_feats,
out_feats=out_feats,
n_etypes=6,
n_steps=5).cuda()
self.conv2 = GraphConv(in_feats=in_feats, out_feats=out_feats)
self.edge_net = nn.Linear(6, n_hidden)
self.feat_net = nn.Linear(in_feats, n_hidden)
self.lin_1 = nn.Linear(out_feats, n_hidden)
self.lstm = nn.LSTM(input_size=n_hidden,
hidden_size=n_hidden,
num_layers=3,
batch_first=True)
self.lin_2 = nn.Linear(in_features=n_hidden, out_features=3)
self.predict = nn.Linear(2 * out_feats, 3)
self.dropout = nn.Dropout(p=0.1)
self.set2set = Set2Set(input_dim=out_feats,
n_iters=n_iter_readout,
n_layers=3)
def __init__(self,
node_in_feats,
edge_in_feats,
node_out_feats=64,
edge_hidden_feats=128,
num_step_message_passing=6,
num_step_set2set=6,
num_layer_set2set=3):
super(MPNN_encoder, self).__init__()
self.gnn = MPNNGNN(node_in_feats=node_in_feats,
node_out_feats=node_out_feats,
edge_in_feats=edge_in_feats,
edge_hidden_feats=edge_hidden_feats,
num_step_message_passing=num_step_message_passing)
self.readout = Set2Set(input_dim=node_out_feats,
n_iters=num_step_set2set,
n_layers=num_layer_set2set)
self.process = nn.Sequential(
# nn.Linear(2 * node_out_feats + 114, node_out_feats),
nn.Linear(2 * node_out_feats, node_out_feats),
nn.ReLU(),
nn.Dropout(p=0.2)
)
def __init__(
self,
positional_embedding_size=32,
max_node_freq=8,
max_edge_freq=8,
max_degree=128,
freq_embedding_size=32,
degree_embedding_size=32,
output_dim=32,
node_hidden_dim=32,
edge_hidden_dim=32,
num_layers=6,
num_heads=4,
num_step_set2set=6,
num_layer_set2set=3,
norm=False,
gnn_model="mpnn",
degree_input=False,
lstm_as_gate=False,
):
super(GraphEncoder, self).__init__()
if degree_input:
node_input_dim = positional_embedding_size + degree_embedding_size + 1
else:
node_input_dim = positional_embedding_size + 1
edge_input_dim = freq_embedding_size + 1
if gnn_model == "mpnn":
self.gnn = UnsupervisedMPNN(
output_dim=output_dim,
node_input_dim=node_input_dim,
node_hidden_dim=node_hidden_dim,
edge_input_dim=edge_input_dim,
edge_hidden_dim=edge_hidden_dim,
num_step_message_passing=num_layers,
lstm_as_gate=lstm_as_gate,
)
elif gnn_model == "gat":
self.gnn = UnsupervisedGAT(
node_input_dim=node_input_dim,
node_hidden_dim=node_hidden_dim,
edge_input_dim=edge_input_dim,
num_layers=num_layers,
num_heads=num_heads,
)
elif gnn_model == "gin":
self.gnn = UnsupervisedGIN(
num_layers=num_layers,
num_mlp_layers=2,
input_dim=node_input_dim,
hidden_dim=node_hidden_dim,
output_dim=output_dim,
final_dropout=0.5,
learn_eps=False,
graph_pooling_type="sum",
neighbor_pooling_type="sum",
use_selayer=False,
)
self.gnn_model = gnn_model
self.max_node_freq = max_node_freq
self.max_edge_freq = max_edge_freq
self.max_degree = max_degree
self.degree_input = degree_input
if degree_input:
self.degree_embedding = nn.Embedding(num_embeddings=max_degree + 1, embedding_dim=degree_embedding_size)
self.set2set = Set2Set(node_hidden_dim, num_step_set2set, num_layer_set2set)
self.lin_readout = nn.Sequential(
nn.Linear(2 * node_hidden_dim, node_hidden_dim),
nn.ReLU(),
nn.Linear(node_hidden_dim, output_dim),
)
self.norm = norm
