def init_encoder(self) -> GraphEncoder:
return GraphEncoder(
edge_block=EdgeBlockEncoder(
self.config.encode.edge.size,
self.config.encode.edge.dropout,
self.config.get_activation(self.config.encode.edge.activation),
),
node_block=NodeBlockEncoder(
self.config.encode.node.size,
self.config.encode.node.dropout,
self.config.get_activation(self.config.encode.node.activation),
),
global_block=GlobalBlockEncoder(
self.config.encode.glob.size,
self.config.encode.glob.dropout,
self.config.get_activation(self.config.encode.glob.activation),
),
)
def init_out_transform(self):
return GraphEncoder(
edge_block=EdgeBlock(
LinearTransformation(
self.config.out.edge.size,
self.config.get_activation(self.config.out.edge.activation),
),
),
node_block=NodeBlock(
LinearTransformation(
self.config.out.node.size,
self.config.get_activation(self.config.out.node.activation),
)
),
global_block=GlobalBlock(
LinearTransformation(
self.config.out.glob.size,
self.config.get_activation(self.config.out.glob.activation),
)
),
)
def _init_encoder(self):
return GraphEncoder(
EdgeBlock(
Flex(MLP)(
Flex.d(),
self.config["sizes"]["latent"]["edge"],
dropout=self.config["dropout"],
)),
NodeBlock(
Flex(MLP)(
Flex.d(),
self.config["sizes"]["latent"]["node"],
dropout=self.config["dropout"],
)),
GlobalBlock(
Flex(MLP)(
Flex.d(),
self.config["sizes"]["latent"]["global"],
dropout=self.config["dropout"],
)),
)
def _init_out_transform(self):
return GraphEncoder(
EdgeBlock(
torch.nn.Sequential(
Flex(torch.nn.Linear)(Flex.d(),
self.config["sizes"]["out"]["edge"]),
self.config["sizes"]["out"]["activation"](),
)),
NodeBlock(
torch.nn.Sequential(
Flex(torch.nn.Linear)(Flex.d(),
self.config["sizes"]["out"]["node"]),
self.config["sizes"]["out"]["activation"](),
)),
GlobalBlock(
torch.nn.Sequential(
Flex(torch.nn.Linear)(
Flex.d(), self.config["sizes"]["out"]["global"]),
self.config["sizes"]["out"]["activation"](),
)),
)
class Networks:
"""Networks that will be used in the tests."""
n = NamedNetwork
linear_block = n(
"linear",
lambda: torch.nn.Sequential(torch.nn.Linear(5, 16), torch.nn.ReLU(),
torch.nn.Linear(16, 1)),
)
mlp_block = n(
"mlp",
lambda: torch.nn.Sequential(
Flex(MLP)(Flex.d(), 16),
Flex(torch.nn.Linear)(Flex.d(), 1)),
)
node_block = n(
"node_block",
lambda: torch.nn.Sequential(
NodeBlock(Flex(MLP)(Flex.d(), 25, 25, layer_norm=False)),
Flex(torch.nn.Linear)(Flex.d(), 1),
),
)
edge_block = n(
"edge_block",
lambda: torch.nn.Sequential(
EdgeBlock(Flex(MLP)(Flex.d(), 25, 25, layer_norm=False)),
Flex(torch.nn.Linear)(Flex.d(), 1),
),
)
global_block = n(
"global_block",
lambda: torch.nn.Sequential(
GlobalBlock(Flex(MLP)(Flex.d(), 25, 25, layer_norm=False)),
Flex(torch.nn.Linear)(Flex.d(), 1),
),
)
graph_encoder = n(
"graph_encoder",
lambda: GraphEncoder(
EdgeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(), 5, 5, layer_norm=False),
Flex(torch.nn.Linear)(Flex.d(), 1),
)),
NodeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(), 5, 5, layer_norm=False),
Flex(torch.nn.Linear)(Flex.d(), 1),
)),
GlobalBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(), 5, 5, layer_norm=False),
Flex(torch.nn.Linear)(Flex.d(), 1),
)),
),
)
def create_graph_core(pass_global_to_edge: bool,
pass_global_to_node: bool):
return GraphCore(
AggregatingEdgeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(), 5, 5, layer_norm=False),
Flex(torch.nn.Linear)(Flex.d(), 1),
)),
AggregatingNodeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(), 5, 5, layer_norm=False),
Flex(torch.nn.Linear)(Flex.d(), 1),
),
edge_aggregator=Aggregator("add"),
),
AggregatingGlobalBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(), 5, 5, layer_norm=False),
Flex(torch.nn.Linear)(Flex.d(), 1),
),
edge_aggregator=Aggregator("add"),
node_aggregator=Aggregator("add"),
),
pass_global_to_edge=pass_global_to_edge,
pass_global_to_node=pass_global_to_node,
)
graph_core = n("graph_core", create_graph_core)
def create_graph_core_multi_agg(pass_global_to_edge: bool,
pass_global_to_node: bool):
agg = lambda: Flex(MultiAggregator)(Flex.d(),
["add", "mean", "max", "min"])
return GraphCore(
AggregatingEdgeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(),
5,
5,
layer_norm=True,
activation=torch.nn.LeakyReLU),
Flex(torch.nn.Linear)(Flex.d(), 1),
)),
AggregatingNodeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(),
5,
5,
layer_norm=True,
activation=torch.nn.LeakyReLU),
Flex(torch.nn.Linear)(Flex.d(), 1),
),
edge_aggregator=agg(),
),
AggregatingGlobalBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(),
5,
5,
layer_norm=True,
activation=torch.nn.LeakyReLU),
Flex(torch.nn.Linear)(Flex.d(), 1),
),
edge_aggregator=agg(),
node_aggregator=agg(),
),
pass_global_to_edge=pass_global_to_edge,
pass_global_to_node=pass_global_to_node,
)
graph_core_multi_agg = n("graph_core(multiagg)",
create_graph_core_multi_agg)
@staticmethod
def reset(net: torch.nn.Module):
def weight_reset(model):
for layer in model.children():
if hasattr(layer, "reset_parameters"):
layer.reset_parameters()
net.apply(weight_reset)
def __init__(
self,
latent_sizes=(16, 16, 1),
out_sizes=(1, 1, 1),
latent_depths=(1, 1, 1),
dropout: float = None,
pass_global_to_edge: bool = True,
pass_global_to_node: bool = True,
activation=defaults.activation,
out_activation=defaults.activation,
edge_to_node_aggregators=tuple(["add", "max", "mean", "min"]),
edge_to_global_aggregators=tuple(["add", "max", "mean", "min"]),
node_to_global_aggregators=tuple(["add", "max", "mean", "min"]),
aggregator_activation=defaults.activation,
):
super().__init__()
self.config = {
"sizes": {
"latent": {
"edge": latent_sizes[0],
"node": latent_sizes[1],
"global": latent_sizes[2],
"edge_depth": latent_depths[0],
"node_depth": latent_depths[1],
"global_depth": latent_depths[2],
},
"out": {
"edge": out_sizes[0],
"node": out_sizes[1],
"global": out_sizes[2],
"activation": out_activation,
},
},
"activation": activation,
"dropout": dropout,
"node_block_aggregator": edge_to_node_aggregators,
"global_block_to_node_aggregator": node_to_global_aggregators,
"global_block_to_edge_aggregator": edge_to_global_aggregators,
"aggregator_activation": aggregator_activation,
"pass_global_to_edge": pass_global_to_edge,
"pass_global_to_node": pass_global_to_node,
}
###########################
# encoder
###########################
self.encoder = self._init_encoder()
self.core = self._init_core()
self.decoder = self._init_encoder()
self.output_transform = self._init_out_transform()
self.output_transform = GraphEncoder(
EdgeBlock(
torch.nn.Sequential(
Flex(torch.nn.Linear)(Flex.d(), 1), torch.nn.Sigmoid())),
NodeBlock(
torch.nn.Sequential(
Flex(torch.nn.Linear)(Flex.d(), 1), torch.nn.Sigmoid())),
GlobalBlock(Flex(torch.nn.Linear)(Flex.d(), 1)),
)
def __init__(
self,
latent_sizes=(128, 128, 1),
depths=(1, 1, 1),
dropout: float = None,
pass_global_to_edge: bool = True,
pass_global_to_node: bool = True,
edge_to_node_aggregators=tuple(["add", "max", "mean", "min"]),
edge_to_global_aggregators=tuple(["add", "max", "mean", "min"]),
node_to_global_aggregators=tuple(["add", "max", "mean", "min"]),
aggregator_activation=defaults.activation,
):
super().__init__()
self.config = {
"latent_size": {
"node": latent_sizes[1],
"edge": latent_sizes[0],
"global": latent_sizes[2],
"core_node_block_depth": depths[0],
"core_edge_block_depth": depths[1],
"core_global_block_depth": depths[2],
},
"node_block_aggregator": edge_to_node_aggregators,
"global_block_to_node_aggregator": node_to_global_aggregators,
"global_block_to_edge_aggregator": edge_to_global_aggregators,
"aggregator_activation": aggregator_activation,
"pass_global_to_edge": pass_global_to_edge,
"pass_global_to_node": pass_global_to_node,
}
self.encoder = GraphEncoder(
EdgeBlock(Flex(MLP)(Flex.d(), latent_sizes[0], dropout=dropout)),
NodeBlock(Flex(MLP)(Flex.d(), latent_sizes[1], dropout=dropout)),
GlobalBlock(Flex(MLP)(Flex.d(), latent_sizes[2], dropout=dropout)),
)
edge_layers = [self.config["latent_size"]["edge"]
] * self.config["latent_size"]["core_edge_block_depth"]
node_layers = [self.config["latent_size"]["node"]
] * self.config["latent_size"]["core_node_block_depth"]
global_layers = [
self.config["latent_size"]["global"]
] * self.config["latent_size"]["core_global_block_depth"]
self.core = GraphCore(
AggregatingEdgeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(),
*edge_layers,
dropout=dropout,
layer_norm=True),
# Flex(torch.nn.Linear)(Flex.d(), edge_layers[-1])
)),
AggregatingNodeBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(),
*node_layers,
dropout=dropout,
layer_norm=True),
# Flex(torch.nn.Linear)(Flex.d(), node_layers[-1])
),
Flex(MultiAggregator)(
Flex.d(),
self.config["node_block_aggregator"],
activation=self.config["aggregator_activation"],
),
),
AggregatingGlobalBlock(
torch.nn.Sequential(
Flex(MLP)(Flex.d(),
*global_layers,
dropout=dropout,
layer_norm=True),
# Flex(torch.nn.Linear)(Flex.d(), global_layers[-1])
),
edge_aggregator=Flex(MultiAggregator)(
Flex.d(),
self.config["global_block_to_edge_aggregator"],
activation=self.config["aggregator_activation"],
),
node_aggregator=Flex(MultiAggregator)(
Flex.d(),
self.config["global_block_to_node_aggregator"],
activation=self.config["aggregator_activation"],
),
),
pass_global_to_edge=self.config["pass_global_to_edge"],
pass_global_to_node=self.config["pass_global_to_node"],
)
self.decoder = GraphEncoder(
EdgeBlock(
Flex(MLP)(Flex.d(),
latent_sizes[0],
latent_sizes[0],
dropout=dropout)),
NodeBlock(
Flex(MLP)(Flex.d(),
latent_sizes[1],
latent_sizes[1],
dropout=dropout)),
GlobalBlock(Flex(MLP)(Flex.d(), latent_sizes[2])),
)
self.output_transform = GraphEncoder(
EdgeBlock(Flex(torch.nn.Linear)(Flex.d(), 1)),
NodeBlock(Flex(torch.nn.Linear)(Flex.d(), 1)),
GlobalBlock(Flex(torch.nn.Linear)(Flex.d(), 1)),
)