AutoGL/docs/docfile/tutorial/t_model.rst

.. _model:

AutoGL Model
============

In AutoGL, we use ``model`` and ``automodel`` to define the logic of graph nerual networks and make it compatible with hyper parameter optimization. Currently we support the following models for given tasks.

+----------------------+----------------------------+
| Tasks                | Models                     |
+======================+============================+
| Node Classification  | ``gcn``, ``gat``, ``sage`` |
+----------------------+----------------------------+
| Graph Classification | ``gin``, ``topk``          |
+----------------------+----------------------------+
| Link Prediction      | ``gcn``, ``gat``, ``sage`` |
+----------------------+----------------------------+

Lazy Initialization
-------------------

In current AutoGL pipeline, some important hyper-parameters related with model cannot be set outside before the pipeline (e.g. input dimensions, which can only be caluclated during running after feature engineered). Therefore, in ``automodel``, we use lazy initialization to initialize the core ``model``. When the ``automodel`` initialization method ``__init__()`` is called with argument ``init`` be ``False``, only (part of) the hyper-parameters will be set. The ``automodel`` will have its core ``model`` only after ``initialize()`` is explicitly called, which will be done automatically in ``solver`` and ``from_hyper_parameter()``, after all the hyper-parameters are set properly.

Define your own model and automodel
-----------------------------------

We highly recommend you to define both ``model`` and ``automodel``, although you only need your ``automodel`` to communicate with ``solver`` and ``trainer``. The ``model`` will be responsible for the parameters initialization and forward logic declaration, while the ``automodel`` will be responsible for the hyper-parameter definiton and organization.

General customization
^^^^^^^^^^^^^^^^^^^^^

Let's say you want to implement a simple MLP for node classification and want to let AutoGL find the best hyper-parameters for you. You can first define the logics assuming all the hyper-parameters are given.

.. code-block:: python

    import torch

    # define mlp model, need to inherit from torch.nn.Module
    class MyMLP(torch.nn.Module):
        # assume you already get all the hyper-parameters
        def __init__(self, in_channels, num_classes, layer_num, dim):
            super().__init__()
            if layer_num == 1:
                ops = [torch.nn.Linear(in_channels, num_classes)]
            else:
                ops = [torch.nn.Linear(in_channels, dim)]
                for i in range(layer_num - 2):
                    ops.append(torch.nn.Linear(dim, dim))
                ops.append(torch.nn.Linear(dim, num_classes))
        
            self.core = torch.nn.Sequential(*ops)
        
        # this method is required
        def forward(self, data):
            # data: torch_geometric.data.Data
            assert hasattr(data, 'x'), 'MLP only support graph data with features'
            x = data.x
            return torch.nn.functional.log_softmax(self.core(x))


After you define the logic of ``model``, you can now define your ``automodel`` to manage the hyper-parameters.

.. code-block:: python

    from autogl.module.model import BaseModel
    
    # define your automodel, need to inherit from BaseModel
    class MyAutoMLP(BaseModel):
        def __init__(self):
            # (required) make sure you call __init__ of super with init argument properly set.
            # if you do not want to initialize inside __init__, please pass False.
            super().__init__(init=False)

            # (required) define the search space
            self.space = [
                {'parameterName': 'layer_num', 'type': 'INTEGER', 'minValue': 1, 'maxValue': 5, 'scalingType': 'LINEAR'},
                {'parameterName': 'dim', 'type': 'INTEGER', 'minValue': 64, 'maxValue': 128, 'scalingType': 'LINEAR'}
            ]

            # set default hyper-parameters
            self.layer_num = 2
            self.dim = 72

            # for the hyper-parameters that are related with dataset, you can just set them to None
            self.num_classes = None
            self.num_features = None

            # (required) since we don't know the num_classes and num_features until we see the dataset,
            # we cannot initialize the models when instantiated. the initialized will be set to False.
            self.initialized = False

            # (required) set the device of current auto model
            self.device = torch.device('cuda')

        # (required) get current hyper-parameters of this automodel
        # need to return a dictionary whose keys are the same with self.space
        def get_hyper_parameter(self):
            return {
                'layer_num': self.layer_num,
                'dim': self.dim
            }
        
        # (required) override to interact with num_classes
        def get_num_classes(self):
            return self.num_classes
        
        # (required) override to interact with num_classes
        def set_num_classes(self, n_classes):
            self.num_classes = n_classes
        
        # (required) override to interact with num_features
        def get_num_features(self):
            return self.num_features
        
        # (required) override to interact with num_features
        def set_num_features(self, n_features):
            self.num_features = n_features

        # (required) instantiate the core MLP model using corresponding hyper-parameters
        def initialize(self):
            # (required) you need to make sure the core model is named as `self.model`
            self.model = MyMLP(
                in_channels = self.num_features,
                num_classes = self.num_classes,
                layer_num = self.layer_num,
                dim = self.dim
            ).to(self.device)

            self.initialized = True
        
        # (required) override to create a copy of model using provided hyper-parameters
        def from_hyper_parameter(self, hp):
            # hp is a dictionary that contains keys and values corrsponding to your self.space
            # in this case, it will be in form {'layer_num': XX, 'dim': XX}
            
            # create a new instance
            ret = self.__class__()

            # set the hyper-parameters related to dataset and device
            ret.num_classes = self.num_classes
            ret.num_features = self.num_features
            ret.device = self.device

            # set the hyper-parameters according to hp
            ret.layer_num = hp['layer_num']
            ret.dim = hp['dim']

            # initialize it before returning
            ret.initialize()

            return ret
        

Then, you can use this node classification model as part of AutoNodeClassifier ``solver``.

.. code-block :: python

    from autogl.solver import AutoNodeClassifier

    solver = AutoNodeClassifier(graph_models=(MyAutoMLP(),))


The model for graph classification is generally the same, except that you can now also receive the ``num_graph_features`` (the dimension of the graph-level feature) through overriding ``set_num_graph_features(self, n_graph_features)`` of ``BaseModel``. Also, please remember to return graph-level logits instead of node-level one in ``forward()`` of ``model``.

Model for link prediction
^^^^^^^^^^^^^^^^^^^^^^^^^

For link prediction, the definition of model is a bit different with the common forward definition. You need to implement the ``lp_encode(self, data)`` and ``lp_decode(self, x, pos_edge_index, neg_edge_index)`` to interact with ``LinkPredictionTrainer`` and ``AutoLinkPredictor``. Taking the class ``MyMLP`` defined above for example, if you want to perform link prediction:

.. code-block:: python

    class MyMLPForLP(torch.nn.Module):
        # num_classes is removed since it is invalid for link prediction
        def __init__(self, in_channels, layer_num, dim):
            super().__init__()
            ops = [torch.nn.Linear(in_channels, dim)]
            for i in range(layer_num - 1):
                ops.append(torch.nn.Linear(dim, dim))
        
            self.core = torch.nn.Sequential(*ops)

        # (required) for interaction with link prediction trainer and solver
        def lp_encode(self, data):
            return self.core(data.x)

        # (required) for interaction with link prediction trainer and solver
        def lp_decode(self, x, pos_edge_index, neg_edge_index):
            # first, get all the edge_index need calculated
            edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1)
            # then, use dot-products to calculate logits, you can use whatever decode method you want
            logits = (x[edge_index[0]] * x[edge_index[1]]).sum(dim=-1)
            return logits

    class MyAutoMLPForLP(MyAutoMLP):
        def initialize(self):
            # init MyMLPForLP instead of MyMLP
            self.model = MyMLPForLP(
                in_channels = self.num_features,
                layer_num = self.layer_num,
                dim = self.dim
            ).to(self.device)

            self.initialized = True


Model with sampling support
^^^^^^^^^^^^^^^^^^^^^^^^^^^

Towards efficient representation learning on large-scale graph, AutoGL currently support node classification using sampling techniques including node-wise sampling, layer-wise sampling, and graph-wise sampling. See more about sampling in :ref:`trainer`.

In order to conduct node classification using sampling technique with your custom model, further adaptation and modification are generally required.
According to the Message Passing mechanism of Graph Neural Network (GNN), numerous nodes in the multi-hop neighborhood of evaluation set or test set are potentially involved to evaluate the GNN model on large-scale graph dataset.
As the representations for those numerous nodes are likely to occupy large amount of computational resource, the common forwarding process is generally infeasible for model evaluation on large-scale graph.
An iterative representation learning mechanism is a practical and feasible way to evaluate **Sequential Model**,
which only consists of multiple sequential layers, with each layer taking a ``Data`` aggregate as input. The input ``Data`` has the same functionality with ``torch_geometric.data.Data``, which conventionally provides properties ``x``, ``edge_index``, and optional ``edge_weight``.
If your custom model is composed of concatenated layers, you would better make your model inherit ``ClassificationSupportedSequentialModel`` to utilize the layer-wise representation learning mechanism to efficiently conduct representation learning for your custom sequential model.

.. code-block:: python

    import autogl
    from autogl.module.model.base import ClassificationSupportedSequentialModel

    # override Linear so that it can take graph data as input
    class Linear(torch.nn.Linear):
        def forward(self, data):
            return super().forward(data.x)

    class MyMLPSampling(ClassificationSupportedSequentialModel):
        def __init__(self, in_channels, num_classes, layer_num, dim):
            super().__init__()
            if layer_num == 1:
                ops = [Linear(in_channels, num_classes)]
            else:
                ops = [Linear(in_channels, dim)]
                for i in range(layer_num - 2):
                    ops.append(Linear(dim, dim))
                ops.append(Linear(dim, num_classes))

            self.core = torch.nn.ModuleList(ops)

        # (required) override sequential_encoding_layers property to interact with sampling
        @property
        def sequential_encoding_layers(self) -> torch.nn.ModuleList:
            return self.core
        
        # (required) define the encode logic of classification for sampling
        def cls_encode(self, data):
            # if you use sampling, the data will be passed in two possible ways,
            # you can judge it use following rules
            if hasattr(data, 'edge_indexes'):
                # the edge_indexes are a list of edge_index, one for each layer
                edge_indexes = data.edge_indexes
                edge_weights = [None] * len(self.core) if getattr(data, 'edge_weights', None) is None else data.edge_weights
            else:
                # the edge_index and edge_weight will stay the same as default
                edge_indexes = [data.edge_index] * len(self.core)
                edge_weights = [getattr(data, 'edge_weight', None)] * len(self.core)

            x = data.x
            for i in range(len(self.core)):
                data = autogl.data.Data(x=x, edge_index=edge_indexes[i])
                data.edge_weight = edge_weights[i]
                x = self.sequential_encoding_layers[i](data)
            return x

        # (required) define the decode logic of classification for sampling
        def cls_decode(self, x):
            return torch.nn.functional.log_softmax(x)