def __init__(self, mlp_layers, N, f_dict):
"""
Simpler version of the alternating model. In this model there is no
encoder network, we only have 1 layer of GNN on each processing step.
We condition on the output global embedding from the processing on the
previous graph, and we only condition the node computations since there
are less nodes than edges (this is a choice that can be discussed).
We aggregate nodes with attention in the global model.
We use the same gnn for processing both inputs.
In this model, since we may want to chain the passes, we let the number
of input features unchanged.
"""
super(AlternatingSimple, self).__init__(f_dict)
model_fn = gn.mlp_fn(mlp_layers)
self.N = N
self.component = 'MPGNN'
# f_e, f_x, f_u, f_out = self.get_features(f_dict)
self.gnn = gn.GNN(
gn.EdgeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fe),
gn.NodeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fx),
gn.GlobalModel_NodeOnly(self.fx, 2 * self.fu, model_fn, self.fu))
self.mlp = model_fn(2 * self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(DeepSetPlus_A, self).__init__(f_dict)
self.N = N # we allow multiple rounds
mlp_fn = gn.mlp_fn(mlp_layers)
self.deepset = gn.DeepSetPlus(
gn.DS_NodeModel(self.fx, self.fu, mlp_fn, self.fx),
gn.DS_GlobalModel_A(self.fx, self.fu, self.h, mlp_fn, self.fu))
self.mlp = mlp_fn(self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(RecurrentGraphEmbeddingDS, self).__init__(f_dict)
self.N = N
model_fn = gn.mlp_fn(mlp_layers)
self.component = 'DS'
self.ds1 = gn.DeepSet(model_fn, self.h, self.h, self.h)
self.ds2 = gn.DeepSet(model_fn, 2 * self.h, self.h, self.h)
self.mlp = model_fn(self.h, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(GNN_NEAgg, self).__init__(f_dict)
self.N = N
mlp_fn = gn.mlp_fn(mlp_layers)
self.gnn = gn.GNN(
gn.EdgeModel(self.fe, self.fx, self.fu, mlp_fn, self.fe),
gn.NodeModel(self.fe, self.fx, self.fu, mlp_fn, self.fx),
gn.GlobalModel(self.fe, self.fx, self.fu, mlp_fn, self.fx))
self.mlp = mlp_fn(self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(AlternatingDoubleDS, self).__init__(f_dict)
model_fn = gn.mlp_fn(mlp_layers)
self.N = N
self.component = 'DS'
# f_e, f_x, f_u, f_out = self.get_features(f_dict)
self.ds1 = gn.DeepSet(model_fn, self.fu + self.fx, self.h, self.fx)
self.ds2 = gn.DeepSet(model_fn, self.fu + self.fx, self.h, self.fx)
self.mlp = model_fn(2 * self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(N_GNN_A, self).__init__(f_dict)
self.N = N
mlp_fn = gn.mlp_fn(mlp_layers)
self.gnn = gn.N_GNN(
gn.EdgeModel_NoMem(self.fx, self.fu, mlp_fn, self.fe),
gn.NodeModel_A(self.fe, self.fx, self.fu, self.h, mlp_fn, self.fx),
gn.GlobalModel_NodeOnly_A(self.fx, self.fu, self.h, mlp_fn,
self.fx))
self.mlp = mlp_fn(self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(AlternatingSimplev2, self).__init__(f_dict)
model_fn = gn.mlp_fn(mlp_layers)
self.N = N
self.component = 'MPGNN'
self.proj = torch.nn.Linear(self.fx, self.h)
self.gnn = gn.GNN(
gn.EdgeModel(self.h, self.h, 2 * self.h, model_fn, self.h),
gn.NodeModel(self.h, self.h, 2 * self.h, model_fn, self.h),
gn.GlobalModel_NodeOnly(self.h, 2 * self.h, model_fn, self.h))
self.mlp = model_fn(2 * self.h, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(AlternatingSimple, self).__init__(f_dict)
model_fn = gn.mlp_fn(mlp_layers)
self.N = N
self.component = 'RDS'
# f_e, f_x, f_u, f_out = self.get_features(f_dict)
self.gnn = gn.DeepSetPlus(
gn.DS_NodeModel(self.fx, 2 * self.fu, model_fn, self.fx),
gn.DS_GlobalModel(self.fx, 2 * self.fu, model_fn, self.fu))
self.mlp = model_fn(2 * self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(RecurrentGraphEmbeddingRDS, self).__init__(f_dict)
self.N = N
model_fn = gn.mlp_fn(mlp_layers)
self.component = 'RDS'
self.gnn1 = gn.DeepSetPlus(
gn.DS_NodeModel(self.fx, self.fu, model_fn, self.fx),
gn.DS_GlobalModel(self.fx, self.fu, model_fn, self.fu))
self.gnn2 = gn.DeepSetPlus(
gn.DS_NodeModel(self.fx, 2 * self.fu, model_fn, self.fx),
gn.DS_GlobalModel(self.fx, 2 * self.fu, model_fn, self.fu))
self.mlp = model_fn(self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(RecurrentGraphEmbeddingRDSv2, self).__init__(f_dict)
model_fn = gn.mlp_fn(mlp_layers)
self.N = N
self.component = 'RDS'
self.proj = torch.nn.Linear(self.fx, self.h)
self.gnn1 = gn.DeepSetPlus(
gn.DS_NodeModel(self.h, self.h, model_fn, self.h),
gn.DS_GlobalModel(self.h, self.h, model_fn, self.h))
self.gnn2 = gn.DeepSetPlus(
gn.DS_NodeModel(self.h, 2 * self.h, model_fn, self.h),
gn.DS_GlobalModel(self.h, 2 * self.h, model_fn, self.h))
self.mlp = model_fn(self.h, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(RecurrentGraphEmbedding, self).__init__(f_dict)
self.N = N
model_fn = gn.mlp_fn(mlp_layers)
self.component = 'MPGNN'
self.gnn1 = gn.GNN(
gn.EdgeModel(self.fe, self.fx, self.fu, model_fn, self.fe),
gn.NodeModel(self.fe, self.fx, self.fu, model_fn, self.fx),
gn.GlobalModel_NodeOnly(self.fx, self.fu, model_fn, self.fx))
self.gnn2 = gn.GNN(
gn.EdgeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fe),
gn.NodeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fx),
gn.GlobalModel_NodeOnly(self.fx, 2 * self.fu, model_fn, self.fx))
self.mlp = model_fn(self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(AlternatingDoubleRDSv2, self).__init__(f_dict)
model_fn = gn.mlp_fn(mlp_layers)
self.N = N
self.component = 'RDS'
# f_e, f_x, f_u, f_out = self.get_features(f_dict)
self.proj = torch.nn.Linear(self.fx, self.h)
self.gnn1 = gn.DeepSetPlus(
gn.DS_NodeModel(self.h, 2 * self.h, model_fn, self.h),
gn.DS_GlobalModel(self.h, 2 * self.h, model_fn, self.h))
self.gnn2 = gn.DeepSetPlus(
gn.DS_NodeModel(self.h, 2 * self.h, model_fn, self.h),
gn.DS_GlobalModel(self.h, 2 * self.h, model_fn, self.h))
self.mlp = model_fn(2 * self.h, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(RecurrentGraphEmbeddingv2, self).__init__(f_dict)
self.N = N
self.component = 'MPGNN'
model_fn = gn.mlp_fn(mlp_layers)
self.proj = torch.nn.Linear(self.fx, self.h)
self.gnn1 = gn.GNN(
gn.EdgeModel(self.h, self.h, self.h, model_fn, self.h),
gn.NodeModel(self.h, self.h, self.h, model_fn, self.h),
gn.GlobalModel_NodeOnly(self.h, self.h, model_fn, self.h))
self.gnn2 = gn.GNN(
gn.EdgeModel(self.h, self.h, 2 * self.h, model_fn, self.h),
gn.NodeModel(self.h, self.h, 2 * self.h, model_fn, self.h),
gn.GlobalModel_NodeOnly(self.h, 2 * self.h, model_fn, self.h))
self.mlp = model_fn(self.h, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(AlternatingDouble, self).__init__(f_dict)
model_fn = gn.mlp_fn(mlp_layers)
self.N = N
self.component = 'MPGNN'
# f_e, f_x, f_u, f_out = self.get_features(f_dict)
self.gnn1 = gn.GNN(
gn.EdgeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fe),
gn.NodeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fx),
gn.GlobalModel_NodeOnly(self.fx, 2 * self.fu, model_fn, self.fu))
self.gnn2 = gn.GNN(
gn.EdgeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fe),
gn.NodeModel(self.fe, self.fx, 2 * self.fu, model_fn, self.fx),
gn.GlobalModel_NodeOnly(self.fx, 2 * self.fu, model_fn, self.fu))
self.mlp = model_fn(2 * self.fu, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(DeepSet, self).__init__(f_dict)
mlp_fn = gn.mlp_fn(mlp_layers)
self.deepset = gn.DeepSet(mlp_fn, self.fx, self.h, self.fout)
def __init__(self, mlp_layers, N, f_dict):
super(TGNN, self).__init__(f_dict)
self.N = N
mlp_fn = gn.mlp_fn(mlp_layers)
self.tgnn = gn.MultiHeadAttention(self.fx, 8, self.h)
self.agg = gn.SumAggreg()
