Class GraphNetwork

climart/models/GraphNet/graph_network.py:14–100

class GraphNetwork(BaseModel):
    def __init__(self,
                 # input_dim: Dict[str, int],
                 hidden_dims: Sequence[int],
                 input_transform: AbstractGraphTransform,
                 readout_which_output: Optional[str] = NODES,
                 update_mlp_n_layers: int = 1,
                 aggregator_funcs: Union[str, Dict[AggregationTypes, int]] = 'sum',
                 net_normalization: str = 'layer_norm',
                 residual: Union[bool, Dict[str, bool]] = True,
                 activation_function: str = 'Gelu',
                 output_activation_function: Optional[str] = None,
                 output_net_normalization: bool = True,
                 dropout: float = 0.0,
                 *args, **kwargs):
        """
        Args:
             readout_which_output: Which graph part to return (default: edges),
                                    can be {EDGES, NODES, GLOBALS, 'graph', None}
                                   If None or 'graph', the whole graph is returned.
        """
        super().__init__(input_transform=input_transform, *args, **kwargs)
        self.save_hyperparameters(ignore="verbose_mlp")
        assert len(self.hparams.hidden_dims) >= 1
        assert update_mlp_n_layers >= 1
        self.input_transform: AbstractGraphTransform = self.input_transform
        in_dim = self.input_transform.output_dim

        senders, receivers = self.input_transform.get_edge_idxs()
        gn_layers = []
        dims = [in_dim] + list(hidden_dims)
        for i in range(1, len(dims)):
            out_activation_function = output_activation_function if i == len(dims) - 1 else activation_function
            out_net_norm = output_net_normalization if i == len(dims) - 1 else True
            gn_layers += [
                GraphNetBlock(
                    in_dims=in_dim,
                    out_dims=dims[i],
                    senders=senders,
                    receivers=receivers,
                    n_layers=update_mlp_n_layers,
                    residual=residual,
                    net_norm=net_normalization,
                    activation=activation_function,
                    dropout=dropout,
                    output_normalization=out_net_norm,
                    output_activation_function=out_activation_function,
                    aggregator_funcs=aggregator_funcs,
                )]
            in_dim = dims[i]

        self.layers: nn.ModuleList[GraphNetBlock] = nn.ModuleList(gn_layers)
        self.output_type = readout_which_output
        if self.output_type not in [NODES, EDGES, GLOBALS, 'graph', None]:
            raise ValueError("Unsupported argument for GraphNetwork `output_type`", readout_which_output)
        if hasattr(self.input_transform, "n_edges"):
            err_msg = f"GN inferred {self.n_edges} edges, but input_tranform refers to {self.input_transform.n_edges}"
            assert self.n_edges == self.input_transform.n_edges, err_msg