model = HeteroGNN(...)

output = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)

from torch_geometric.datasets import OGB_MAG

dataset = OGB_MAG(root='./data', preprocess='metapath2vec')
data = dataset[0]

HeteroData(
  paper={
    x=[736389, 128],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389]
  },
  author={ x=[1134649, 128] },
  institution={ x=[8740, 128] },
  field_of_study={ x=[59965, 128] },
  (author, affiliated_with, institution)={ edge_index=[2, 1043998] },
  (author, writes, paper)={ edge_index=[2, 7145660] },
  (paper, cites, paper)={ edge_index=[2, 5416271] },
  (paper, has_topic, field_of_study)={ edge_index=[2, 7505078] }
)

The :class:`torch_geometric.data.HeteroData` class provides a number of useful utility functions to modify and analyze the given graph.

For example, single node or edge stores can be individually indexed:

.. code-block:: python

    paper_node_data = data['paper']
    cites_edge_data = data['paper', 'cites', 'paper']

In case the edge type can be uniquely identified by only the pair of source and destination node types or the edge type, the following operations work as well:

.. code-block:: python

    cites_edge_data = data['paper', 'paper']
    cites_edge_data = data['cites']

We can add new node types or tensors and remove them:

.. code-block:: python

    data['paper'].year = ...    # Setting a new paper attribute
    del data['field_of_study']  # Deleting 'field_of_study' node type
    del data['has_topic']       # Deleting 'has_topic' edge type

We can access the meta-data of the :obj:`data` object, holding information of all present node and edge types:

.. code-block:: python

    node_types, edge_types = data.metadata()
    print(node_types)
    ['paper', 'author', 'institution']
    print(edge_types)
    [('paper', 'cites', 'paper'),
    ('author', 'writes', 'paper'),
    ('author', 'affiliated_with', 'institution')]

The :obj:`data` object can be transferred between devices as usual:

.. code-block:: python

    data = data.to('cuda:0')
    data = data.cpu()

We further have access to additional helper functions to analyze the given graph

.. code-block:: python

    data.has_isolated_nodes()
    data.has_self_loops()
    data.is_undirected()

and can convert it to a homogeneous "typed" graph via :meth:`~torch_geometric.data.HeteroData.to_homogeneous` which is able to maintain features in case their dimensionalities match across different types:

.. code-block:: python

    homogeneous_data = data.to_homogeneous()
    print(homogeneous_data)
    Data(x=[1879778, 128], edge_index=[2, 13605929], edge_type=[13605929])

Here, :obj:`homogeneous_data.edge_type` represents an edge-level vector that holds the edge type of each edge as an integer.

Heterogeneous Graph Transformations
-----------------------------------

Most `transformations <https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html>`_ for preprocessing regular graphs work as well on the heterogeneous graph :obj:`data` object.

.. code-block:: python

    import torch_geometric.transforms as T

    data = T.ToUndirected()(data)
    data = T.AddSelfLoops()(data)
    data = T.NormalizeFeatures()(data)

Here, :meth:`~torch_geometric.transforms.ToUndirected` transforms a directed graph into (the :pyg:`PyG` representation of) an undirected graph, by adding reverse edges for all edges in the graph.
Thus, future message passing is performed in both direction of all edges.
The function may add reverse edge types to the heterogeneous graph, if necessary.

For all nodes of type :obj:`'node_type'` and all existing edge types of the form :obj:`('node_type', 'edge_type', 'node_type')`, the function :meth:`~torch_geometric.transforms.AddSelfLoops` will add self-loop edges.
As a result, each node might receive one or more (one per appropriate edge type) messages from itself during message passing.

The transform :meth:`~torch_geometric.transforms.NormalizeFeatures` works like in the homogeneous case, and normalizes all specified features (of all types) to sum up to one.

Creating Heterogeneous GNNs
---------------------------

Standard Message Passing GNNs (MP-GNNs) can not trivially be applied to heterogeneous graph data, as node and edge features from different types can not be processed by the same functions due to differences in feature type.
A natural way to circumvent this is to implement message and update functions individually for each edge type.
During runtime, the MP-GNN algorithm would need to iterate over edge type dictionaries during message computation and over node type dictionaries during node updates.

To avoid unnecessary runtime overheads and to make the creation of heterogeneous MP-GNNs as simple as possible, Pytorch Geometric provides three ways for the user to create models on heterogeneous graph data:

#. Automatically convert a homogeneous GNN model to a heterogeneous GNN model by making use of :meth:`torch_geometric.nn.to_hetero` or :meth:`torch_geometric.nn.to_hetero_with_bases`
#. Define individual functions for different types using :pyg:`PyG's` wrapper :class:`torch_geometric.nn.conv.HeteroConv` for heterogeneous convolution
#. Deploy existing (or write your own) heterogeneous GNN operators

In the following, each option is introduced in detail.

Automatically Converting GNN Models

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv, to_hetero


dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x


model = GNN(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

    with torch.no_grad():  # Initialize lazy modules.
        out = model(data.x_dict, data.edge_index_dict)

from torch_geometric.nn import GATConv, Linear, to_hetero

class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x


model = GAT(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x_dict, data.edge_index_dict)
    mask = data['paper'].train_mask
    loss = F.cross_entropy(out['paper'][mask], data['paper'].y[mask])
    loss.backward()
    optimizer.step()
    return float(loss)

The heterogeneous convolution wrapper :class:`torch_geometric.nn.conv.HeteroConv` allows to define custom heterogeneous message and update functions to build arbitrary MP-GNNs for heterogeneous graphs from scratch.
While the automatic converter :meth:`~torch_geometric.nn.to_hetero` uses the same operator for all edge types, the wrapper allows to define different operators for different edge types.
Here, :class:`~torch_geometric.nn.conv.HeteroConv` takes a dictionary of submodules as input, one for each edge type in the graph data.
The following `example <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/hetero_conv_dblp.py>`__ shows how to apply it.

.. code-block:: python

    import torch_geometric.transforms as T
    from torch_geometric.datasets import OGB_MAG
    from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv, Linear


    dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
    data = dataset[0]

    class HeteroGNN(torch.nn.Module):
        def __init__(self, hidden_channels, out_channels, num_layers):
            super().__init__()

            self.convs = torch.nn.ModuleList()
            for _ in range(num_layers):
                conv = HeteroConv({
                    ('paper', 'cites', 'paper'): GCNConv(-1, hidden_channels),
                    ('author', 'writes', 'paper'): SAGEConv((-1, -1), hidden_channels),
                    ('paper', 'rev_writes', 'author'): GATConv((-1, -1), hidden_channels, add_self_loops=False),
                }, aggr='sum')
                self.convs.append(conv)

            self.lin = Linear(hidden_channels, out_channels)

        def forward(self, x_dict, edge_index_dict):
            for conv in self.convs:
                x_dict = conv(x_dict, edge_index_dict)
                x_dict = {key: x.relu() for key, x in x_dict.items()}
            return self.lin(x_dict['author'])

    model = HeteroGNN(hidden_channels=64, out_channels=dataset.num_classes,
                      num_layers=2)

We can initialize the model by calling it once (see :ref:`here<lazyinit>` for more details about lazy initialization)

.. code-block:: python

    with torch.no_grad():  # Initialize lazy modules.
         out = model(data.x_dict, data.edge_index_dict)

and run the standard training procedure as outlined :ref:`here<trainfunc>`.

Deploy Existing Heterogeneous Operators
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

:pyg:`PyG` provides operators (*e.g.*, :class:`torch_geometric.nn.conv.HGTConv`), which are specifically designed for heterogeneous graphs.
These operators can be directly used to build heterogeneous GNN models as can be seen in the following `example <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/hgt_dblp.py>`__:

.. code-block:: python

    import torch_geometric.transforms as T
    from torch_geometric.datasets import OGB_MAG
    from torch_geometric.nn import HGTConv, Linear


    dataset = OGB_MAG(root='./data', preprocess='metapath2vec', transform=T.ToUndirected())
    data = dataset[0]

    class HGT(torch.nn.Module):
        def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
            super().__init__()

            self.lin_dict = torch.nn.ModuleDict()
            for node_type in data.node_types:
                self.lin_dict[node_type] = Linear(-1, hidden_channels)

            self.convs = torch.nn.ModuleList()
            for _ in range(num_layers):
                conv = HGTConv(hidden_channels, hidden_channels, data.metadata(),
                               num_heads, group='sum')
                self.convs.append(conv)

            self.lin = Linear(hidden_channels, out_channels)

        def forward(self, x_dict, edge_index_dict):
            for node_type, x in x_dict.items():
                x_dict[node_type] = self.lin_dict[node_type](x).relu_()

            for conv in self.convs:
                x_dict = conv(x_dict, edge_index_dict)

            return self.lin(x_dict['author'])

    model = HGT(hidden_channels=64, out_channels=dataset.num_classes,
                num_heads=2, num_layers=2)

We can initialize the model by calling it once (see :ref:`here<lazyinit>` for more details about lazy initialization).

.. code-block:: python

    with torch.no_grad():  # Initialize lazy modules.
         out = model(data.x_dict, data.edge_index_dict)

and run the standard training procedure as outlined :ref:`here<trainfunc>`.

Heterogeneous Graph Samplers
----------------------------

:pyg:`PyG` provides various functionalities for sampling heterogeneous graphs, *i.e.* in the standard :class:`torch_geometric.loader.NeighborLoader` class  or in dedicated heterogeneous graph samplers such as :class:`torch_geometric.loader.HGTLoader`.
This is especially useful for efficient representation learning on large heterogeneous graphs, where processing the full number of neighbors is too computationally expensive.
Heterogeneous graph support for other samplers such as :class:`torch_geometric.loader.ClusterLoader` or :class:`torch_geometric.loader.GraphSAINTLoader` will be added soon.
Overall, all heterogeneous graph loaders will produce a :class:`~torch_geometric.data.HeteroData` object as output, holding a subset of the original data, and mainly differ in the way their sampling procedures works.
As such, only minimal code changes are required to convert the training procedure from :ref:`full-batch training<trainfunc>` to mini-batch training.

Performing neighbor sampling using :class:`~torch_geometric.loader.NeighborLoader` works as outlined in the following `example <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/to_hetero_mag.py>`__:

.. code-block:: python

    import torch_geometric.transforms as T
    from torch_geometric.datasets import OGB_MAG
    from torch_geometric.loader import NeighborLoader

    transform = T.ToUndirected()  # Add reverse edge types.
    data = OGB_MAG(root='./data', preprocess='metapath2vec', transform=transform)[0]

    train_loader = NeighborLoader(
        data,
        # Sample 15 neighbors for each node and each edge type for 2 iterations:
        num_neighbors=[15] * 2,
        # Use a batch size of 128 for sampling training nodes of type "paper":
        batch_size=128,
        input_nodes=('paper', data['paper'].train_mask),
    )

    batch = next(iter(train_loader))


Notably, :class:`~torch_geometric.loader.NeighborLoader` works for both homogeneous *and* heterogeneous graphs.
When operating in heterogeneous graphs, more fine-grained control over the amount of sampled neighbors of individual edge types is possible, but not necessary, *e.g.*:

.. code-block:: python

    num_neighbors = {key: [15] * 2 for key in data.edge_types}

Using the :obj:`input_nodes` argument, we further specify the type and indices of nodes from which we want to sample local neighborhoods, *i.e.* all the "paper" nodes marked as training nodes according to :obj:`data['paper'].train_mask`.

Printing :obj:`batch` then yields the following output:

.. code-block:: text

    HeteroData(
      paper={
        x=[20799, 256],
        y=[20799],
        train_mask=[20799],
        val_mask=[20799],
        test_mask=[20799],
        batch_size=128
      },
      author={ x=[4419, 128] },
      institution={ x=[302, 128] },
      field_of_study={ x=[2605, 128] },
      (author, affiliated_with, institution)={ edge_index=[2, 0] },
      (author, writes, paper)={ edge_index=[2, 5927] },
      (paper, cites, paper)={ edge_index=[2, 11829] },
      (paper, has_topic, field_of_study)={ edge_index=[2, 10573] },
      (institution, rev_affiliated_with, author)={ edge_index=[2, 829] },
      (paper, rev_writes, author)={ edge_index=[2, 5512] },
      (field_of_study, rev_has_topic, paper)={ edge_index=[2, 10499] }
    )

As such, :obj:`batch` holds a total of 28,187 nodes involved for computing the embeddings of 128 "paper" nodes.
Sampled nodes are always sorted based on the order in which they were sampled.
Thus, the first :obj:`batch['paper'].batch_size` nodes represent the set of original mini-batch nodes, making it easy to obtain the final output embeddings via slicing.

Training our heterogeneous GNN model in mini-batch mode is then similar to training it in full-batch mode, except that we now iterate over the mini-batches produced by :obj:`train_loader` and optimize model parameters based on individual mini-batches:

.. code-block:: python

    def train():
        model.train()

        total_examples = total_loss = 0
        for batch in train_loader:
            optimizer.zero_grad()
            batch = batch.to('cuda:0')
            batch_size = batch['paper'].batch_size
            out = model(batch.x_dict, batch.edge_index_dict)
            loss = F.cross_entropy(out['paper'][:batch_size],
                                   batch['paper'].y[:batch_size])
            loss.backward()
            optimizer.step()

            total_examples += batch_size
            total_loss += float(loss) * batch_size

        return total_loss / total_examples

Importantly, we only make use of the first 128 "paper" nodes during loss computation.
We do so by slicing both "paper" labels :obj:`batch['paper'].y` and "paper" output predictions :obj:`out['paper']` based on :obj:`batch['paper'].batch_size`, representing the labels and final output predictions of original mini-batch nodes, respectively.