示例#1
0
文件: network.py 项目: jvrana/caldera
 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"],
             )),
     )
示例#2
0
 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),
             )
         ),
     )
示例#3
0
文件: network.py 项目: jvrana/caldera
 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"](),
             )),
     )
示例#4
0
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)
示例#5
0
文件: network.py 项目: jvrana/caldera
    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)),
        )
示例#6
0
    def __init__(
        self,
        latent_sizes=(128, 128, 1),
        output_sizes=(1, 1, 1),
        depths=(1, 1, 1),
        layer_norm: bool = True,
        dropout: float = None,
        pass_global_to_edge: bool = True,
        pass_global_to_node: bool = True,
    ):
        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],
            },
            "output_size": {
                "edge": output_sizes[0],
                "node": output_sizes[1],
                "global": output_sizes[2],
            },
            "node_block_aggregator": "add",
            "global_block_to_node_aggregator": "add",
            "global_block_to_edge_aggregator": "add",
            "pass_global_to_edge": pass_global_to_edge,
            "pass_global_to_node": pass_global_to_node,
        }

        def mlp(*layer_sizes):
            return Flex(MLP)(Flex.d(),
                             *layer_sizes,
                             layer_norm=layer_norm,
                             dropout=dropout)

        self.encoder = GraphEncoder(
            EdgeBlock(mlp(latent_sizes[0])),
            NodeBlock(mlp(latent_sizes[1])),
            GlobalBlock(mlp(latent_sizes[2])),
        )

        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(mlp(*edge_layers)),
            AggregatingNodeBlock(
                mlp(*node_layers),
                Aggregator(self.config["node_block_aggregator"])),
            AggregatingGlobalBlock(
                mlp(*global_layers),
                edge_aggregator=Aggregator(
                    self.config["global_block_to_edge_aggregator"]),
                node_aggregator=Aggregator(
                    self.config["global_block_to_node_aggregator"]),
            ),
            pass_global_to_edge=self.config["pass_global_to_edge"],
            pass_global_to_node=self.config["pass_global_to_node"],
        )

        self.decoder = GraphEncoder(
            EdgeBlock(mlp(latent_sizes[0])),
            NodeBlock(mlp(latent_sizes[1])),
            GlobalBlock(mlp(latent_sizes[2])),
        )

        self.output_transform = GraphEncoder(
            EdgeBlock(Flex(torch.nn.Linear)(Flex.d(), output_sizes[0])),
            NodeBlock(Flex(torch.nn.Linear)(Flex.d(), output_sizes[1])),
            GlobalBlock(Flex(torch.nn.Linear)(Flex.d(), output_sizes[2])),
        )
示例#7
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)),
        )