Exemple #1
0
 def init_core(self):
     return GraphCore(
         edge_block=EdgeBlockCore(
             layers=self.config.core.edge.layers,
             dropout=self.config.core.edge.dropout,
             layer_norm=self.config.core.edge.layer_norm,
         ),
         node_block=NodeBlockCore(
             layers=self.config.core.node.layers,
             dropout=self.config.core.node.dropout,
             aggregator=self.config.edge_block_to_node_aggregators,
             aggregator_activation=self.config.get_activation(
                 self.config.aggregator_activation
             ),
             layer_norm=self.config.core.node.layer_norm,
         ),
         global_block=GlobalBlockCore(
             layers=self.config.core.glob.layers,
             dropout=self.config.core.glob.dropout,
             edge_aggregator=self.config.global_block_to_edge_aggregators,
             node_aggregator=self.config.global_block_to_node_aggregators,
             layer_norm=self.config.core.glob.layer_norm,
             aggregator_activation=self.config.get_activation(
                 self.config.aggregator_activation
             ),
         ),
         pass_global_to_node=self.config.pass_global_to_node,
         pass_global_to_edge=self.config.pass_global_to_edge,
     )
Exemple #2
0
 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,
     )
Exemple #3
0
    def _init_core(self):
        edge_layers = [self.config["sizes"]["latent"]["edge"]
                       ] * self.config["sizes"]["latent"]["edge_depth"]
        node_layers = [self.config["sizes"]["latent"]["node"]
                       ] * self.config["sizes"]["latent"]["node_depth"]
        global_layers = [self.config["sizes"]["latent"]["global"]
                         ] * self.config["sizes"]["latent"]["global_depth"]

        return GraphCore(
            AggregatingEdgeBlock(
                torch.nn.Sequential(
                    Flex(MLP)(Flex.d(),
                              *edge_layers,
                              dropout=self.config["dropout"],
                              layer_norm=True), )),
            AggregatingNodeBlock(
                torch.nn.Sequential(
                    Flex(MLP)(Flex.d(),
                              *node_layers,
                              dropout=self.config["dropout"],
                              layer_norm=True), ),
                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=self.config["dropout"],
                              layer_norm=True), ),
                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"],
        )
Exemple #4
0
    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,
        )
Exemple #5
0
    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)),
        )