# Bringing contextual word representations into your models

```python
__author__ = "Christopher Potts"
__version__ = "CS224u, Stanford, Fall 2020"
```

## Contents

1. [Overview](#Overview)
1. [General set-up](#General-set-up)
1. [Hugging Face BERT interfaces](#Hugging-Face-BERT-interfaces)
  1. [Hugging Face BERT set-up](#Hugging-Face-BERT-set-up)
  1. [Hugging Face BERT basics](#Hugging-Face-BERT-basics)
  1. [BERT featurization with Hugging Face](#BERT-featurization-with-Hugging-Face)
    1. [Simple feed-forward experiment](#Simple-feed-forward-experiment)
    1. [A feed-forward experiment with the sst module](#A-feed-forward-experiment-with-the-sst-module)
    1. [An RNN experiment with the sst module](#An-RNN-experiment-with-the-sst-module)
  1. [BERT fine-tuning with Hugging Face](#BERT-fine-tuning-with-Hugging-Face)
    1. [HfBertClassifier](#HfBertClassifier)
    1. [HfBertClassifier experiment](#HfBertClassifier-experiment)
1. [Using ELMo](#Using-ELMo)
  1. [ELMo Allen NLP set-up](#ELMo-Allen-NLP-set-up)
  1. [ELMo fine-tuning](#ELMo-fine-tuning)
    1. [AllenNLP ELMo interfaces](#AllenNLP-ELMo-interfaces)
    1. [ElmoRNNClassifier](#ElmoRNNClassifier)
    1. [ElmoRNNClassifier experiment](#ElmoRNNClassifier-experiment)

## Overview

This notebook provides a basic introduction to using pre-trained [BERT](https://github.com/google-research/bert) and [ELMo](https://allennlp.org/elmo) representations. It is meant as a practical companion to our lecture on contextual word representations. The goal of this notebook is just to help you use these representations in your own work. The BERT and ELMo teams have done amazing work to make these resources available to the community. Many projects can benefit from them, so it is probably worth your time to experiment.

This notebook should be considered an experimental extension to the regular course materials. It has some special requirements – libraries and data files – that are not part of the core requirements for this repository. All these tools are very new and being updated frequently, so you might need to do some fiddling to get all of this to work. As I said, though, it's probably worth the effort!

A number of the experiments in this notebook are resource intensive. I've included timing information for the expensive steps, to give you a sense for how long things are likely to take. I ran this notebook on a laptop with a single NVIDIA RTX 2080 GPU.

## General set-up

The following are requirements that you'll already have met if you've been working in this repository. As you can see, we'll use the [Stanford Sentiment Treebank](sst_01_overview.ipynb) for illustrations, and we'll try out a few different deep learning models.

```python
import os
import sst
from torch_shallow_neural_classifier import TorchShallowNeuralClassifier
from torch_rnn_classifier import TorchRNNModel
from torch_rnn_classifier import TorchRNNClassifier
from torch_rnn_classifier import TorchRNNClassifierModel
from torch_rnn_classifier import TorchRNNClassifier
from sklearn.metrics import classification_report
import utils
```

```python
utils.fix_random_seeds()
```

```python
SST_HOME = os.path.join("data", "trees")
```

## Hugging Face BERT interfaces

### Hugging Face BERT set-up

To install this library, run

```pip install transformers```

I've tested this code with versions 2.4, 2.5, and 2.11 of `transformers`. Try to get at least 2.5. It requires `pip >= 20` and, I think, a version of [Rust](https://www.rust-lang.org) at least as high as 1.21.1.

```python
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer
```

The `transformers` library does a lot of logging. To avoid ending up with a cluttered notebook, I am changing the logging level. You might want to skip this as you scale up to building production systems, since the logging is very good – it gives you a lot of insights into what the models and code are doing.

```python
import logging
logger = logging.getLogger()
logger.level = logging.ERROR
```

### Hugging Face BERT basics

To start, let's get a feel for the basic API that `transformers` provides. The first step is specifying the pretrained parameters we'll be using:

```python
hf_weights_name = 'bert-base-cased'
```

There are lots other options for pretrained weights. See [this section of the project README.md](https://github.com/huggingface/transformers#quick-tour) for a good overview and code that documents how these weights align with different Transformer model classes.

Next, we specify a tokenizer and a model that match both each other and our choice of pretrained weights:

```python
hf_tokenizer = BertTokenizer.from_pretrained(hf_weights_name)
```

```python
hf_model = BertModel.from_pretrained(hf_weights_name)
```

It's illuminating to see what the tokenizer does to example texts:

```python
hf_example_texts = [
    "Encode sentence 1. [SEP] And sentence 2!",
    "Bert knows Snuffleupagus"]
```

The `encode` method maps individual strings to indices into the underlying embedding used by the model:

```python
ex0_ids = hf_tokenizer.encode(hf_example_texts[0], add_special_tokens=True)

ex0_ids
```

    [101, 13832, 13775, 5650, 122, 119, 102, 1262, 5650, 123, 106, 102]

We can get a better feel for what these representations are like by mapping the indices back to "words":

```python
hf_tokenizer.convert_ids_to_tokens(ex0_ids)
```

    ['[CLS]',
     'En',
     '##code',
     'sentence',
     '1',
     '.',
     '[SEP]',
     'And',
     'sentence',
     '2',
     '!',
     '[SEP]']

For modeling, we will often need to pad (and perhaps truncate) token lists so that we can work with fixed-dimensional tensors: The `batch_encode_plus` has a lot of options for doing this:

```python
hf_example_ids = hf_tokenizer.batch_encode_plus(
    hf_example_texts,
    add_special_tokens=True,
    return_attention_mask=True,
    pad_to_max_length=True)
```

```python
hf_example_ids.keys()
```

    dict_keys(['input_ids', 'token_type_ids', 'attention_mask'])

The `token_type_ids` is used for multi-text inputs like NLI. The `'input_ids'` field gives the indices for each of the two examples:

```python
hf_example_ids['input_ids']
```

    [[101, 13832, 13775, 5650, 122, 119, 102, 1262, 5650, 123, 106, 102],
     [101, 15035, 3520, 156, 14787, 13327, 4455, 28026, 1116, 102, 0, 0]]

For fine-tuning, we want to avoid attending to padded tokens. The `'attention_mask'` captures the needed mask, which we'll be able to feed directly to the pretrained BERT model:

```python
hf_example_ids['attention_mask']
```

    [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]

Finally, we can run these indices and masks through the pretrained model:

```python
X_hf_example = torch.tensor(hf_example_ids['input_ids'])
X_hf_example_mask = torch.tensor(hf_example_ids['attention_mask'])

with torch.no_grad():
    hf_final_hidden_states, cls_output = hf_model(
        X_hf_example, attention_mask=X_hf_example_mask)
```

BERT representations are pretty large – this shows the shape of the tensor for 2 examples, with the second padded to the length of the larger one in the batch (12). The individual representations have dimensionality 768.

```python
hf_final_hidden_states.shape
```

    torch.Size([2, 12, 768])

Those are all the essential ingredients for working with these parameters in Hugging Face. Of course, the library has a lot of other functionality, but the above suffices to featurize and to fine-tune.

### BERT featurization with Hugging Face

To start, we'll use the Hugging Face interfaces just to featurize examples to create inputs to a separate model. In this setting, the BERT parameters are frozen. The heart of this approach is the following featurizer, which flattens an SST tree into a string, tokenizes it, and computes its hidden representations:

```python
def hugging_face_bert_phi(tree):
    s = " ".join(tree.leaves())
    input_ids = hf_tokenizer.encode(s, add_special_tokens=True)
    X = torch.tensor([input_ids])
    with torch.no_grad():
        final_hidden_states, cls_output = hf_model(X)
        return final_hidden_states.squeeze(0).numpy()
```

#### Simple feed-forward experiment

For a simple feed-forward experiment, we can get the representation of the `[CLS]` tokens and use them as the inputs to a shallow neural network:

```python
def hugging_face_bert_classifier_phi(tree):
    reps = hugging_face_bert_phi(tree)
    #return reps.mean(axis=0)  # Another good, easy option.
    return reps[0]
```

Next we read in the SST train and dev portions as a lists of `(tree, label)` pairs:

```python
hf_train = list(sst.train_reader(SST_HOME, class_func=sst.binary_class_func))

hf_dev = list(sst.dev_reader(SST_HOME, class_func=sst.binary_class_func))
```

Split the input/output pairs out into separate lists:

```python
X_hf_tree_train, y_hf_train = zip(*hf_train)

X_hf_tree_dev, y_hf_dev = zip(*hf_dev)
```

In the next step, we featurize all of the examples. These steps are likely to be the slowest in these experiments:

```python
%time X_hf_train = [hugging_face_bert_classifier_phi(tree)
                    for tree in X_hf_tree_train]
```

    CPU times: user 34min 55s, sys: 6.99 s, total: 35min 2s
    Wall time: 5min 52s

```python
%time X_hf_dev = [hugging_face_bert_classifier_phi(tree)
                  for tree in X_hf_tree_dev]
```

    CPU times: user 4min 13s, sys: 868 ms, total: 4min 13s
    Wall time: 42.6 s

Now that all the examples are featurized, we can fit a model and evaluate it:

```python
hf_mod = TorchShallowNeuralClassifier(
    early_stopping=True,
    hidden_dim=300)
```

```python
%time _ = hf_mod.fit(X_hf_train, y_hf_train)
```

    Stopping after epoch 38. Validation score did not improve by tol=1e-05 for more than 10 epochs. Final error is 2.0183719843626022

    CPU times: user 19.2 s, sys: 496 ms, total: 19.7 s
    Wall time: 5.28 s

```python
hf_preds = hf_mod.predict(X_hf_dev)
```

```python
print(classification_report(y_hf_dev, hf_preds, digits=3))
```

                  precision    recall  f1-score   support

        negative      0.850     0.834     0.842       428
        positive      0.843     0.858     0.850       444

        accuracy                          0.846       872
       macro avg      0.846     0.846     0.846       872
    weighted avg      0.846     0.846     0.846       872

#### A feed-forward experiment with the sst module

It is straightforward to conduct experiments like the above using `sst.experiment`, which will enable you to do a wider range of experiments without writing or copy-pasting a lot of code.

```python
def fit_hf_shallow_network(X, y):
    mod = TorchShallowNeuralClassifier(
        hidden_dim=300,
        early_stopping=True)
    mod.fit(X, y)
    return mod
```

```python
%%time
_ = sst.experiment(
    SST_HOME,
    hugging_face_bert_classifier_phi,
    fit_hf_shallow_network,
    train_reader=sst.train_reader,
    assess_reader=sst.dev_reader,
    class_func=sst.binary_class_func,
    vectorize=False)  # Pass in the BERT reps directly!
```

    Stopping after epoch 22. Validation score did not improve by tol=1e-05 for more than 10 epochs. Final error is 2.3845930993556976

                  precision    recall  f1-score   support

        negative      0.839     0.839     0.839       428
        positive      0.845     0.845     0.845       444

        accuracy                          0.842       872
       macro avg      0.842     0.842     0.842       872
    weighted avg      0.842     0.842     0.842       872

    CPU times: user 42min 27s, sys: 8.83 s, total: 42min 36s
    Wall time: 7min 8s

#### An RNN experiment with the sst module

We can also use BERT representations as the input to an RNN. There is just one key change from how we used these models before:

* Previously, we would feed in lists of tokens, and they would be converted to indices into a fixed embedding space. This presumes that all words have the same representation no matter what their context is.

* With BERT, we skip the embedding entirely and just feed in lists of BERT vectors, which means that the same word can be represented in different ways.

`TorchRNNClassifier` supports this via `use_embedding=False`. In turn, you needn't supply a vocabulary:

```python
def fit_hf_rnn(X, y):
    mod = TorchRNNClassifier(
        vocab=[],
        early_stopping=True,
        use_embedding=False)  # Pass in the BERT hidden states directly!
    mod.fit(X, y)
    return mod
```

```python
%%time
_ = sst.experiment(
    SST_HOME,
    hugging_face_bert_phi,
    fit_hf_rnn,
    train_reader=sst.train_reader,
    assess_reader=sst.dev_reader,
    class_func=sst.binary_class_func,
    vectorize=False)  # Pass in the BERT hidden states directly!
```

    Stopping after epoch 25. Validation score did not improve by tol=1e-05 for more than 10 epochs. Final error is 0.9420086964964867

                  precision    recall  f1-score   support

        negative      0.849     0.853     0.851       428
        positive      0.857     0.854     0.856       444

        accuracy                          0.853       872
       macro avg      0.853     0.853     0.853       872
    weighted avg      0.853     0.853     0.853       872

    CPU times: user 44min, sys: 31.2 s, total: 44min 31s
    Wall time: 7min 37s

### BERT fine-tuning with Hugging Face

The above experiments are quite successful – BERT gives us a reliable boost compared to other methods we've explored for the SST task. However, we might expect to do even better if we fine-tune the BERT parameters as part of fitting our SST classifier. To do that, we need to incorporate the Hugging Face BERT model into our classifier. This too is quite straightforward.

#### HfBertClassifier

The most important step is to create an `nn.Module` subclass that has, for its parameters, both the BERT model and parameters for our own classifier. Here we define a very simple fine-tuning set-up in which some layers built on top of the output corresponding to `[CLS]` are used as the basis for the SST classifier:

```python
class HfBertClassifierModel(nn.Module):
    def __init__(self, n_classes, weights_name='bert-base-cased'):
        super().__init__()
        self.n_classes = n_classes
        self.weights_name = weights_name
        self.bert = BertModel.from_pretrained(self.weights_name)
        self.bert.train()
        self.hidden_dim = self.bert.embeddings.word_embeddings.embedding_dim
        # The only new parameters -- the classifier:
        self.classifier_layer = nn.Linear(
            self.hidden_dim, self.n_classes)

    def forward(self, indices, mask):
        final_hidden_states, cls_output = self.bert(
            indices, attention_mask=mask)
        return self.classifier_layer(cls_output)
```

As you can see, `self.bert` does the heavy-lifting: it reads in all the pretrained BERT parameters, and I've specified `self.bert.train()` just to make sure that these parameters can be updated during our training process.

In `forward`, `self.bert` is used to process inputs, and then `cls_output` is fed into `self.classifier_layer`. Hugging Face has already added a layer on top of the actual output for `[CLS]`, so we can specify the model as

$$
\begin{align}
[h_{1}, \ldots, h_{n}] &= \text{BERT}([x_{1}, \ldots, x_{n}]) \\
h &= \tanh(h_{1}W_{hh} + b_{h}) \\
y &= \textbf{softmax}(hW_{hy} + b_{y})
\end{align}$$

for a tokenized input sequence $[x_{1}, \ldots, x_{n}]$.

The Hugging Face documentation somewhat amusingly says, of `cls_output`,

> This output is usually _not_ a good summary of the semantic content of the input, you're often better with averaging or pooling the sequence of hidden-states for the whole input sequence.

which is entirely reasonable, but it will require more resources, so we'll do the simpler thing here.

For the training and prediction interface, we can subclass `TorchShallowNeuralClassifier` so that we don't have to write any of our own data-handling, training, or prediction code. The central changes are using `HfBertClassifierModel` in `build_graph` and processing the data with `batch_encode_plus`.

```python
class HfBertClassifier(TorchShallowNeuralClassifier):
    def __init__(self, weights_name, *args, **kwargs):
        self.weights_name = weights_name
        self.tokenizer = BertTokenizer.from_pretrained(self.weights_name)
        super().__init__(*args, **kwargs)
        self.params += ['weights_name']

    def build_graph(self):
        return HfBertClassifierModel(self.n_classes_, self.weights_name)

    def build_dataset(self, X, y=None):
        data = self.tokenizer.batch_encode_plus(
            X,
            max_length=None,
            add_special_tokens=True,
            pad_to_max_length=True,
            return_attention_mask=True)
        indices = torch.tensor(data['input_ids'])
        mask = torch.tensor(data['attention_mask'])
        if y is None:
            dataset = torch.utils.data.TensorDataset(indices, mask)
        else:
            self.classes_ = sorted(set(y))
            self.n_classes_ = len(self.classes_)
            class2index = dict(zip(self.classes_, range(self.n_classes_)))
            y = [class2index[label] for label in y]
            y = torch.tensor(y)
            dataset = torch.utils.data.TensorDataset(indices, mask, y)
        return dataset
```

#### HfBertClassifier experiment

That's it! Let's see how we do on the SST binary, root-only problem. Because fine-tuning is expensive, we'll conduct a modest hyperparameter search and run the model for just one epoch per setting evaluation, as we did when [assessing NLI models](nli_02_models.ipynb).

```python
def bert_fine_tune_phi(tree):
    return " ".join(tree.leaves())
```

```python
def fit_hf_bert_classifier_with_hyperparameter_search(X, y):
    basemod = HfBertClassifier(
        weights_name='bert-base-cased',
        batch_size=8,  # Small batches to avoid memory overload.
        max_iter=1,  # We'll search based on 1 iteration for efficiency.
        n_iter_no_change=5,   # Early-stopping params are for the
        early_stopping=True)  # final evaluation.

    param_grid = {
        'gradient_accumulation_steps': [1, 4, 8],
        'eta': [0.00005, 0.0001, 0.001],
        'hidden_dim': [100, 200, 300]}

    bestmod = utils.fit_classifier_with_hyperparameter_search(
        X, y, basemod, cv=3, param_grid=param_grid)

    return bestmod
```

```python
%%time
hf_bert_classifier_xval = sst.experiment(
    SST_HOME,
    bert_fine_tune_phi,
    fit_hf_bert_classifier_with_hyperparameter_search,
    train_reader=sst.train_reader,
    assess_reader=sst.dev_reader,
    class_func=sst.binary_class_func,
    vectorize=False)  # Pass in the BERT hidden state directly!
```

    Finished epoch 1 of 1; error is 38.469218485057354

    Best params: {'eta': 5e-05, 'gradient_accumulation_steps': 8, 'hidden_dim': 300}
    Best score: 0.892
                  precision    recall  f1-score   support

        negative      0.953     0.808     0.875       428
        positive      0.839     0.962     0.896       444

        accuracy                          0.886       872
       macro avg      0.896     0.885     0.885       872
    weighted avg      0.895     0.886     0.886       872

    CPU times: user 1h 31min 39s, sys: 4min 6s, total: 1h 35min 46s
    Wall time: 1h 36min 43s

And now on to the final test-set evaluation, using the best model from above:

```python
optimized_hf_bert_classifier = hf_bert_classifier_xval['model']
```

```python
# Remove the rest of the experiment results to clear out some memory:
del hf_bert_classifier_xval
```

```python
def fit_optimized_hf_bert_classifier(X, y):
    optimized_hf_bert_classifier.max_iter = 1000
    optimized_hf_bert_classifier.fit(X, y)
    return optimized_hf_bert_classifier
```

```python
%%time
_ = sst.experiment(
    SST_HOME,
    bert_fine_tune_phi,
    fit_optimized_hf_bert_classifier,
    train_reader=(sst.train_reader, sst.dev_reader),
    assess_reader=sst.test_reader,
    class_func=sst.binary_class_func,
    vectorize=False)  # Pass in the BERT hidden state directly!
```

    Stopping after epoch 10. Validation score did not improve by tol=1e-05 for more than 5 epochs. Final error is 0.964973742607981

                  precision    recall  f1-score   support

        negative      0.890     0.928     0.908       912
        positive      0.924     0.884     0.904       909

        accuracy                          0.906      1821
       macro avg      0.907     0.906     0.906      1821
    weighted avg      0.907     0.906     0.906      1821

    CPU times: user 13min 12s, sys: 17.4 s, total: 13min 29s
    Wall time: 13min 29s

The above is just one of the many possible ways to fine-tune BERT using our course modules or new modules you write. The crux of it is creating an `nn.Module` that combines the BERT parameters with your model's new parameters.

## Using ELMo

### ELMo Allen NLP set-up

There are a number of ways to use pre-trained ELMo models. We'll use the simplest of the AllenNLP interfaces. Run the following to install [AllenNLP](https://allennlp.org):

```sh
pip install allennlp
```
I've tested this notebook with versions, 0.8.0, 0.9.0, and 1.0.0.

Mac users: If your installation fails, make sure your Xcode tools are up to date by running `xcode-select --install`.

```python
from allennlp.modules.elmo import Elmo, batch_to_ids
import torch
import torch.nn as nn
```

We'll use the following models, which will download from S3 to a local temp directory the first time you use them with `ElmoEmbedder` or `Elmo` as described below.

```python
elmo_file_path = "https://allennlp.s3.amazonaws.com/models/elmo/2x4096_512_2048cnn_2xhighway/"

options_file = elmo_file_path + "elmo_2x4096_512_2048cnn_2xhighway_options.json"

weights_file = elmo_file_path + "elmo_2x4096_512_2048cnn_2xhighway_weights.hdf5"
```

For more models:

https://allennlp.org/elmo

For additional details:

https://docs.allennlp.org/master/api/modules/elmo/

### ELMo fine-tuning

Fine-tuning ELMo proceeds in essentially the same way it did for BERT: we create an `nn.Module` that combines the parameters from ELMo with our task-specific parameters and then optimize everything on the new task. To illustrate, I'll define an RNN on top of the ELMo model using new subclasses of `TorchRNNClassifier` and `TorchRNNClassifierModel`.

#### AllenNLP ELMo interfaces

To start, let's get a feel for the primary interface, and then we'll write the classes that will allow us to use these components systematically.

The interface to the ELMo parameters in this context is the class `Elmo`:

```python
elmo = Elmo(options_file, weights_file, num_output_representations=1)
```

This model expects tokenized inputs:

```python
elmo_example_texts = [
    ["Encode", "sentence", "1", "."],
    ["ELMo", "knows" "Snuffleupagus"]]
```

The ELMo model processes its tokens at the character-level, creating convolutional representations for the words from various character n-gram combinations:

```python
elmo_character_ids = batch_to_ids(elmo_example_texts)

# First word of the first example:
elmo_character_ids[0][0]
```

    tensor([259,  70, 111, 100, 112, 101, 102, 260, 261, 261, 261, 261, 261, 261,
            261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261,
            261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261, 261,
            261, 261, 261, 261, 261, 261, 261, 261])

`elmo` embeds these at the word-level:

```python
elmo_embeddings = elmo(elmo_character_ids)
```

`elmo_embeddings` is a dict. The value of the key `'elmo_representations'` is a list of tensors corresponding to each layer of the model. In other words, each tensor in the list is a complete representation of the example. The final element of the list is the final representation layer. I specified `num_output_representations=1` when initializing `elmo` above, so we get a list of length 1:

```python
elmo_embeddings['elmo_representations']
```

    [tensor([[[-0.7944,  0.0000,  0.0000,  ..., -0.0000, -0.0000, -1.9283],
              [-0.0000,  1.1850,  0.6947,  ...,  0.0000, -0.9297, -0.2358],
              [ 0.0000,  0.5358,  0.7767,  ..., -0.6500, -0.0777, -0.4875],
              [-0.0000, -0.4965, -0.0000,  ..., -0.1605,  0.0000,  0.2256]],

             [[ 0.4553, -0.0000,  0.0000,  ..., -0.0000, -0.0000, -0.5380],
              [ 0.0000, -0.5309, -0.0000,  ..., -0.2244, -0.0000,  0.7476],
              [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
              [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],
            grad_fn=<MulBackward0>)]

#### ElmoRNNClassifier

The above are the representations we will be fine-tuning. There are many ways to cdo this. In my simple illustration, I just take the top layer, as we did in the simpler featurization example above, but now keeping each word representation separate for use in the input to the task-specific RNN. Here is the `nn.Module` built on `TorchRNNClassifierModel`:

```python
class ElmoRNNClassifierModel(TorchRNNClassifierModel):
    def __init__(self,
            options_file,
            weights_file,
            rnn,
            output_dim,
            classifier_activation):
        super().__init__(rnn, output_dim, classifier_activation)
        self.options_file = options_file
        self.weights_file = weights_file
        self.elmo = Elmo(
            self.options_file,
            self.weights_file,
            num_output_representations=1,
            dropout=0)
        self.elmo.train()

    def forward(self, X, seq_lengths):
        result = self.elmo(X)
        X = result['elmo_representations'][-1]
        outputs, state  = self.rnn(X, seq_lengths)
        state = self.get_batch_final_states(state)
        if self.rnn.bidirectional:
            state = torch.cat((state[0], state[1]), dim=1)
        h = self.classifier_activation(self.hidden_layer(state))
        logits = self.classifier_layer(h)
        return logits
```

And here is the subclass of `TorchRNNClassifier` that lets us take advantage of all the optimization and prediction methods of that class:

```python
class ElmoRNNClassifier(TorchRNNClassifier):
    def __init__(self, vocab, options_file, weights_file, **model_kwargs):
        self.options_file = options_file
        self.weights_file = weights_file
        # Values determined by using ELMo:
        model_kwargs['use_embedding'] = False
        model_kwargs['embedding'] = None
        model_kwargs['embed_dim'] = 1024
        super().__init__(vocab, **model_kwargs)
        self.params += ['options_file', 'weights_file']

    def build_graph(self):

        # The RNN is setup just as in a regular `TorchRNNClassifier`:
        rnn = TorchRNNModel(
            vocab_size=len(self.vocab),
            embedding=self.embedding,
            use_embedding=self.use_embedding,
            embed_dim=self.embed_dim,
            rnn_cell_class=self.rnn_cell_class,
            hidden_dim=self.hidden_dim,
            bidirectional=self.bidirectional,
            freeze_embedding=self.freeze_embedding)

        # The Classifier layer uses our new `ElmoRNNClassifierModel`:
        model = ElmoRNNClassifierModel(
            options_file=self.options_file,
            weights_file=self.weights_file,
            rnn=rnn,
            output_dim=self.n_classes_,
            classifier_activation=self.classifier_activation)

        return model

    def build_dataset(self, X, y=None):
        seq_lengths = [len(ex) for ex in X]
        seq_lengths = torch.tensor(seq_lengths)
        X = batch_to_ids(X)
        if y is None:
            return torch.utils.data.TensorDataset(X, seq_lengths)
        else:
            self.classes_ = sorted(set(y))
            self.n_classes_ = len(self.classes_)
            class2index = dict(zip(self.classes_, range(self.n_classes_)))
            y = [class2index[label] for label in y]
            y = torch.tensor(y)
            return torch.utils.data.TensorDataset(X, seq_lengths, y)
```

#### ElmoRNNClassifier experiment

And finally here is a self-contained evaluation involving a modest hyperparameter search:

```python
def elmo_fine_tune_phi(tree):
    return tree.leaves()
```

```python
def fit_elmo_rnn_with_hyperparameter_search(X, y):
    basemod = ElmoRNNClassifier(
        vocab=[],
        options_file=options_file,
        weights_file=weights_file,
        batch_size=8,  # Kept small so that we can explore large networks.
        max_iter=1,  # We'll search based on 1 iteration for efficiency.
        n_iter_no_change=5,   # Early-stopping params are for the
        early_stopping=True)  # final evalution.

    param_grid = {
        'gradient_accumulation_steps': [1, 4, 8],
        'eta': [0.001, 0.01, 0.05],
        'hidden_dim': [50, 100, 200]}

    bestmod = utils.fit_classifier_with_hyperparameter_search(
        X, y, basemod, cv=3, param_grid=param_grid)

    return bestmod
```

```python
%%time
elmo_rnn_xval = sst.experiment(
    SST_HOME,
    elmo_fine_tune_phi,
    fit_elmo_rnn_with_hyperparameter_search,
    train_reader=sst.train_reader,
    assess_reader=sst.dev_reader,
    class_func=sst.binary_class_func,
    vectorize=False)  # Pass in the ELMo reps directly!
```

    Finished epoch 1 of 1; error is 330.32408917695284

    Best params: {'eta': 0.001, 'gradient_accumulation_steps': 1, 'hidden_dim': 200}
    Best score: 0.856
                  precision    recall  f1-score   support

        negative      0.848     0.848     0.848       428
        positive      0.854     0.854     0.854       444

        accuracy                          0.851       872
       macro avg      0.851     0.851     0.851       872
    weighted avg      0.851     0.851     0.851       872

    CPU times: user 2h 4min 55s, sys: 4min 16s, total: 2h 9min 12s
    Wall time: 2h 1min 35s

And now we move to the test-set evaluation using the best model we found:

```python
optimized_elmo_rnn = elmo_rnn_xval['model']
```

```python
# Remove the unneeded experimental data to save memory:
del elmo_rnn_xval
```

```python
def fit_optimized_elmo_rnn(X, y):
    optimized_elmo_rnn.max_iter = 20
    optimized_elmo_rnn.fit(X, y)
    return optimized_elmo_rnn
```

```python
%%time
_ = sst.experiment(
    SST_HOME,
    elmo_fine_tune_phi,
    fit_optimized_elmo_rnn,
    train_reader=(sst.train_reader, sst.dev_reader),
    assess_reader=sst.test_reader,
    class_func=sst.binary_class_func,
    vectorize=False)  # Pass in the BERT hidden state directly!
```

    Stopping after epoch 10. Validation score did not improve by tol=1e-05 for more than 5 epochs. Final error is 27.70847176760435

                  precision    recall  f1-score   support

        negative      0.891     0.880     0.886       912
        positive      0.882     0.892     0.887       909

        accuracy                          0.886      1821
       macro avg      0.886     0.886     0.886      1821
    weighted avg      0.886     0.886     0.886      1821

    CPU times: user 15min 31s, sys: 34.3 s, total: 16min 6s
    Wall time: 16min


# Markdown command :

 Commands

 - Esc : command mode

- esc H -> show shortcuts
- esc Y -> cell to code
- esc M -> cell to markdown

- esc C -> copy cell
- esc shift-V / V -> paste cell above / below
- esc A/B -> insert cell above / below

- esc shift Q -> clear o/p
- esc O -> toggle cell o/p

- esc L / esc shift-L -> show line number / all line number
- esc F -> find and replace

- esc D-D -> delete cell
- esc Z -> undo delete cell

 Custom

- esc q-q -> shutdown

 - Edit mode:

- enter (from command mode)
- ctrl-shift-minus (split cell at cursor)

#### Markdwon and Latex

`| This | is   |
|------|------|
|   a  | table|`

---------------

| This | is   |
|------|------|
|   a  | table|

\`\`\`python (language) Your-text \`\`\`

```javascript
console.log("Hello World")
```

```python
console.log("Hello World")
```

```

<video controls src="../images/animation.m4v" />
```

Or

```python
from IPython.display import Image
Image(filename='output1.png')
```

- Inline expressions can be added by surrounding the latex code with \$:
- \$\$ <- moves to center

``` $e^{i\pi} + 1 = 0$```

$$e^{i\pi} + 1 = 0$$

```$$e^x=\sum_{i=0}^\infty \frac{1}{i!}x^i$$```

$$e^x=\sum_{i=0}^\infty \frac{1}{i!}x^i$$

```\begin{align}
\dot{x} & = \sigma(y-x) \\
\dot{y} & = \rho x - y - xz \\
\dot{z} & = -\beta z + xy
\end{align} ```

\begin{align}
\dot{x} & = \sigma(y-x) \\
\dot{y} & = \rho x - y - xz \\
\dot{z} & = -\beta z + xy
\end{align}

```\begin{equation*}
\left( \sum_{k=1}^n a_k b_k \right)^2 \leq \left( \sum_{k=1}^n a_k^2 \right) \left( \sum_{k=1}^n b_k^2 \right)
\end{equation*}```

\begin{equation*}
\left( \sum_{k=1}^n a_k b_k \right)^2 \leq \left( \sum_{k=1}^n a_k^2 \right) \left( \sum_{k=1}^n b_k^2 \right)
\end{equation*}

#### Matrix

```\begin{equation*}
\mathbf{V}_1 \times \mathbf{V}_2 =  \begin{vmatrix}
\mathbf{i} & \mathbf{j} & \mathbf{k} \\
\frac{\partial X}{\partial u} &  \frac{\partial Y}{\partial u} & 0 \\
\frac{\partial X}{\partial v} &  \frac{\partial Y}{\partial v} & 0
\end{vmatrix}
\end{equation*}```

\begin{equation*}
\mathbf{V}_1 \times \mathbf{V}_2 =  \begin{vmatrix}
\mathbf{i} & \mathbf{j} & \mathbf{k} \\
\frac{\partial X}{\partial u} &  \frac{\partial Y}{\partial u} & 0 \\
\frac{\partial X}{\partial v} &  \frac{\partial Y}{\partial v} & 0
\end{vmatrix}
\end{equation*}

- begin{vmatrix} -> matrix
- & separates the column and "\\\" separates row

#### Probability

```\begin{equation*}
P(E)   = {n \choose k} p^k (1-p)^{ n-k}
\end{equation*}```

\begin{equation*}
P(E)   = {n \choose k} p^k (1-p)^{ n-k}
\end{equation*}

#### Trival

```\begin{equation*}
\frac{1}{\Bigl(\sqrt{\phi \sqrt{5}}-\phi\Bigr) e^{\frac25 \pi}} =
1+\frac{e^{-2\pi}} {1+\frac{e^{-4\pi}} {1+\frac{e^{-6\pi}}
{1+\frac{e^{-8\pi}} {1+\ldots} } } }
\end{equation*}```

\begin{equation*}
\frac{1}{\Bigl(\sqrt{\phi \sqrt{5}}-\phi\Bigr) e^{\frac25 \pi}} =
1+\frac{e^{-2\pi}} {1+\frac{e^{-4\pi}} {1+\frac{e^{-6\pi}}
{1+\frac{e^{-8\pi}} {1+\ldots} } } }
\end{equation*}

#### Text in equation

```\begin{equation*}
1 +  \frac{q^2}{(1-q)}+\frac{q^6}{(1-q)(1-q^2)}+\cdots =
\prod_{j=0}^{\infty}\frac{1}{(1-q^{5j+2})(1-q^{5j+3})},
\quad\quad \text{for $|q|<1$}.
\end{equation*}```

\begin{equation*}
1 +  \frac{q^2}{(1-q)}+\frac{q^6}{(1-q)(1-q^2)}+\cdots =
\prod_{j=0}^{\infty}\frac{1}{(1-q^{5j+2})(1-q^{5j+3})},
\quad\quad \text{for $|q|<1$}.
\end{equation*}

#### Derivative eqn and matrix operators

```\begin{align}
\nabla \times \vec{\mathbf{B}} -\, \frac1c\, \frac{\partial\vec{\mathbf{E}}}{\partial t} & = \frac{4\pi}{c}\vec{\mathbf{j}} \\   \nabla \cdot \vec{\mathbf{E}} & = 4 \pi \rho \\
\nabla \times \vec{\mathbf{E}}\, +\, \frac1c\, \frac{\partial\vec{\mathbf{B}}}{\partial t} & = \vec{\mathbf{0}} \\
\nabla \cdot \vec{\mathbf{B}} & = 0
\end{align} ```

\begin{align}
\nabla \times \vec{\mathbf{B}} -\, \frac1c\, \frac{\partial\vec{\mathbf{E}}}{\partial t} & = \frac{4\pi}{c}\vec{\mathbf{j}} \\   \nabla \cdot \vec{\mathbf{E}} & = 4 \pi \rho \\
\nabla \times \vec{\mathbf{E}}\, +\, \frac1c\, \frac{\partial\vec{\mathbf{B}}}{\partial t} & = \vec{\mathbf{0}} \\
\nabla \cdot \vec{\mathbf{B}} & = 0
\end{align}

- vector -> vec
- bold -> mathbf
- \, (comma) -> space

```\begin{align}
\nabla \cdot \,\,\,\,\,\,\,\, \vec{\mathbf{B}}  &= 0
\end{align}```

\begin{align}
\nabla \cdot \,\,\,\,\,\,\,\, \vec{\mathbf{B}}  &= 0
\end{align}

#### Inline expression and sentencing
```This expression $\sqrt{3x-1}+(1+x)^2$ is an example of a TeX inline equation in a [Markdown-formatted](https://daringfireball.net/projects/markdown/) sentence.```

This expression $\sqrt{3x-1}+(1+x)^2$ is an example of a TeX inline equation in a [Markdown-formatted](https://daringfireball.net/projects/markdown/) sentence.

#### Other example
```$$x = {-b \pm \sqrt{b^2-4ac} \over 2a}$$```

$$x = {-b \pm \sqrt{b^2-4ac} \over 2a}$$

#### Other example

```$$x = {-b \pm \sqrt{b^2-4ac} \over 2a}$$```

$$x = {-b \pm \sqrt{b^2-4ac} \over 2a}$$

```$$f(a) = \frac{1}{2\pi i} \oint\frac{f(z)}{z-a}dz$$```

$$f(a) = \frac{1}{2\pi i} \oint\frac{f(z)}{z-a}dz$$

```$$\int_D ({\nabla\cdot} F)dV=\int_{\partial D} F\cdot ndS$$```

$$\int_D ({\nabla\cdot} F)dV=\int_{\partial D} F\cdot ndS$$

 - int_{lower}^{upper}

```$$\int_{a}^{b} x^2 dx$$```
$$\int_{a}^{b} x^2 dx$$

```$$\vec{\nabla} \times \vec{F} = \left( \frac{\partial F_z}{\partial y} - \frac{\partial F_y}{\partial z} \right) \mathbf{i}
        + \left( \frac{\partial F_x}{\partial z} - \frac{\partial F_z}{\partial x} \right) \mathbf{j} + \left( \frac{\partial
        F_y}{\partial x} - \frac{\partial F_x}{\partial y} \right) \mathbf{k}$$```

$$\vec{\nabla} \times \vec{F} = \left( \frac{\partial F_z}{\partial y} - \frac{\partial F_y}{\partial z} \right) \mathbf{i}
        + \left( \frac{\partial F_x}{\partial z} - \frac{\partial F_z}{\partial x} \right) \mathbf{j} + \left( \frac{\partial
        F_y}{\partial x} - \frac{\partial F_x}{\partial y} \right) \mathbf{k}$$

 - (left / right) <- brackets
$$\left( hello \right)$$

```python

```


```python
%matplotlib inline
```

Tutorial: Generative models of graphs
===========================================

**Author**: `Mufei Li <https://github.com/mufeili>`_,
`Lingfan Yu <https://github.com/ylfdq1118>`_, Zheng Zhang

In this tutorial, you learn how to train and generate one graph at
a time. You also explore parallelism within the graph embedding operation, which is an
essential building block. The tutorial ends with a simple optimization that
delivers double the speed by batching across graphs.

Earlier tutorials showed how embedding a graph or
a node enables you to work on tasks such as `semi-supervised classification for nodes
<http://docs.dgl.ai/tutorials/models/1_gcn.html#sphx-glr-tutorials-models-1-gcn-py>`__
or `sentiment analysis
<http://docs.dgl.ai/tutorials/models/3_tree-lstm.html#sphx-glr-tutorials-models-3-tree-lstm-py>`__.
Wouldn't it be interesting to predict the future evolution of the graph and
perform the analysis iteratively?

To address the evolution of the graphs, you generate a variety of graph samples. In other words, you need
**generative models** of graphs. In-addition to learning
node and edge features, you would need to model the distribution of arbitrary graphs.
While general generative models can model the density function explicitly and
implicitly and generate samples at once or sequentially, you only focus
on explicit generative models for sequential generation here. Typical applications
include drug or materials discovery, chemical processes, or proteomics.

Introduction
--------------------
The primitive actions of mutating a graph in Deep Graph Library (DGL) are nothing more than ``add_nodes``
and ``add_edges``. That is, if you were to draw a circle of three nodes,

.. figure:: https://user-images.githubusercontent.com/19576924/48313438-78baf000-e5f7-11e8-931e-cd00ab34fa50.gif
   :alt:

you can write the code as follows.

```python
import dgl

g = dgl.DGLGraph()
g.add_nodes(1)              # Add node 0
g.add_nodes(1)              # Add node 1

# Edges in DGLGraph are directed by default.
# For undirected edges, add edges for both directions.
g.add_edges([1, 0], [0, 1]) # Add edges (1, 0), (0, 1)
g.add_nodes(1)              # Add node 2
g.add_edges([2, 1], [1, 2]) # Add edges (2, 1), (1, 2)
g.add_edges([2, 0], [0, 2]) # Add edges (2, 0), (0, 2)
```

    Using backend: pytorch
    /home/viktor/anaconda3/envs/dgl/lib/python3.8/site-packages/dgl/base.py:45: DGLWarning: Recommend creating graphs by `dgl.graph(data)` instead of `dgl.DGLGraph(data)`.
      return warnings.warn(message, category=category, stacklevel=1)

Real-world graphs are much more complex. There are many families of graphs,
with different sizes, topologies, node types, edge types, and the possibility
of multigraphs. Besides, a same graph can be generated in many different
orders. Regardless, the generative process entails a few steps.

- Encode a changing graph.
- Perform actions stochastically.
- If you are training, collect error signals and optimize the model parameters.

When it comes to implementation, another important aspect is speed. How do you
parallelize the computation, given that generating a graph is fundamentally a
sequential process?

Note

To be sure, this is not necessarily a hard constraint. Subgraphs can be
   built in parallel and then get assembled. But we
   will restrict ourselves to the sequential processes for this tutorial.

DGMG: The main flow
--------------------
For this tutorial, you use
`Deep Generative Models of Graphs <https://arxiv.org/abs/1803.03324>`__
) (DGMG) to implement a graph generative model using DGL. Its algorithmic
framework is general but also challenging to parallelize.

Note

While it's possible for DGMG to handle complex graphs with typed nodes,
   typed edges, and multigraphs, here you use a simplified version of it
   for generating graph topologies.

DGMG generates a graph by following a state machine, which is basically a
two-level loop. Generate one node at a time and connect it to a subset of
the existing nodes, one at a time. This is similar to language modeling. The
generative process is an iterative one that emits one word or character or sentence
at a time, conditioned on the sequence generated so far.

At each time step, you either:
     - Add a new node to the graph
     - Select two existing nodes and add an edge between them

.. figure:: https://user-images.githubusercontent.com/19576924/48605003-7f11e900-e9b6-11e8-8880-87362348e154.png
   :alt:

The Python code will look as follows. In fact, this is *exactly* how inference
with DGMG is implemented in DGL.

```python
def forward_inference(self):
    stop = self.add_node_and_update()
    while (not stop) and (self.g.number_of_nodes() < self.v_max + 1):
        num_trials = 0
        to_add_edge = self.add_edge_or_not()
        while to_add_edge and (num_trials < self.g.number_of_nodes() - 1):
            self.choose_dest_and_update()
            num_trials += 1
            to_add_edge = self.add_edge_or_not()
        stop = self.add_node_and_update()

    return self.g
```

Assume you have a pre-trained model for generating cycles of nodes 10-20.
How does it generate a cycle on-the-fly during inference? Use the code below
to create an animation with your own model.

::

    import torch
    import matplotlib.animation as animation
    import matplotlib.pyplot as plt
    import networkx as nx
    from copy import deepcopy

    if __name__ == '__main__':
        # pre-trained model saved with path ./model.pth
        model = torch.load('./model.pth')
        model.eval()
        g = model()

        src_list = g.edges()[1]
        dest_list = g.edges()[0]

        evolution = []

        nx_g = nx.Graph()
        evolution.append(deepcopy(nx_g))

        for i in range(0, len(src_list), 2):
            src = src_list[i].item()
            dest = dest_list[i].item()
            if src not in nx_g.nodes():
                nx_g.add_node(src)
                evolution.append(deepcopy(nx_g))
            if dest not in nx_g.nodes():
                nx_g.add_node(dest)
                evolution.append(deepcopy(nx_g))
            nx_g.add_edges_from([(src, dest), (dest, src)])
            evolution.append(deepcopy(nx_g))

        def animate(i):
            ax.cla()
            g_t = evolution[i]
            nx.draw_circular(g_t, with_labels=True, ax=ax,
                             node_color=['#FEBD69'] * g_t.number_of_nodes())

        fig, ax = plt.subplots()
        ani = animation.FuncAnimation(fig, animate,
                                      frames=len(evolution),
                                      interval=600)

.. figure:: https://user-images.githubusercontent.com/19576924/48928548-2644d200-ef1b-11e8-8591-da93345382ad.gif
   :alt:

DGMG: Optimization objective
------------------------------
Similar to language modeling, DGMG trains the model with *behavior cloning*,
or *teacher forcing*. Assume for each graph there exists a sequence of
*oracle actions* $a_{1},\cdots,a_{T}$ that generates it. What the model
does is to follow these actions, compute the joint probabilities of such
action sequences, and maximize them.

By chain rule, the probability of taking $a_{1},\cdots,a_{T}$ is:

\begin{align}p(a_{1},\cdots, a_{T}) = p(a_{1})p(a_{2}|a_{1})\cdots p(a_{T}|a_{1},\cdots,a_{T-1}).\\\end{align}

The optimization objective is then simply the typical MLE loss:

\begin{align}-\log p(a_{1},\cdots,a_{T})=-\sum_{t=1}^{T}\log p(a_{t}|a_{1},\cdots, a_{t-1}).\\\end{align}

```python
def forward_train(self, actions):
    """
    - actions: list
        - Contains a_1, ..., a_T described above
    - self.prepare_for_train()
        - Initializes self.action_step to be 0, which will get
          incremented by 1 every time it is called.
        - Initializes objects recording log p(a_t|a_1,...a_{t-1})

    Returns
    -------
    - self.get_log_prob(): log p(a_1, ..., a_T)
    """
    self.prepare_for_train()

    stop = self.add_node_and_update(a=actions[self.action_step])
    while not stop:
        to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
        while to_add_edge:
            self.choose_dest_and_update(a=actions[self.action_step])
            to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
        stop = self.add_node_and_update(a=actions[self.action_step])

    return self.get_log_prob()
```

The key difference between ``forward_train`` and ``forward_inference`` is
that the training process takes oracle actions as input and returns log
probabilities for evaluating the loss.

DGMG: The implementation
--------------------------
The ``DGMG`` class
``````````````````````````
Below you can find the skeleton code for the model. You gradually
fill in the details for each function.

```python
import torch.nn as nn

class DGMGSkeleton(nn.Module):
    def __init__(self, v_max):
        """
        Parameters
        ----------
        v_max: int
            Max number of nodes considered
        """
        super(DGMGSkeleton, self).__init__()

        # Graph configuration
        self.v_max = v_max

    def add_node_and_update(self, a=None):
        """Decide if to add a new node.
        If a new node should be added, update the graph."""
        return NotImplementedError

    def add_edge_or_not(self, a=None):
        """Decide if a new edge should be added."""
        return NotImplementedError

    def choose_dest_and_update(self, a=None):
        """Choose destination and connect it to the latest node.
        Add edges for both directions and update the graph."""
        return NotImplementedError

    def forward_train(self, actions):
        """Forward at training time. It records the probability
        of generating a ground truth graph following the actions."""
        return NotImplementedError

    def forward_inference(self):
        """Forward at inference time.
        It generates graphs on the fly."""
        return NotImplementedError

    def forward(self, actions=None):
        # The graph you will work on
        self.g = dgl.DGLGraph()

        # If there are some features for nodes and edges,
        # zero tensors will be set for those of new nodes and edges.
        self.g.set_n_initializer(dgl.frame.zero_initializer)
        self.g.set_e_initializer(dgl.frame.zero_initializer)

        if self.training:
            return self.forward_train(actions=actions)
        else:
            return self.forward_inference()
```

Encoding a dynamic graph
``````````````````````````
All the actions generating a graph are sampled from probability
distributions. In order to do that, you project the structured data,
namely the graph, onto an Euclidean space. The challenge is that such
process, called *embedding*, needs to be repeated as the graphs mutate.

Graph embedding
''''''''''''''''''''''''''
Let $G=(V,E)$ be an arbitrary graph. Each node $v$ has an
embedding vector $\textbf{h}_{v} \in \mathbb{R}^{n}$. Similarly,
the graph has an embedding vector $\textbf{h}_{G} \in \mathbb{R}^{k}$.
Typically, $k > n$ since a graph contains more information than
an individual node.

The graph embedding is a weighted sum of node embeddings under a linear
transformation:

\begin{align}\textbf{h}_{G} =\sum_{v\in V}\text{Sigmoid}(g_m(\textbf{h}_{v}))f_{m}(\textbf{h}_{v}),\\\end{align}

The first term, $\text{Sigmoid}(g_m(\textbf{h}_{v}))$, computes a
gating function and can be thought of as how much the overall graph embedding
attends on each node. The second term $f_{m}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{k}$
maps the node embeddings to the space of graph embeddings.

Implement graph embedding as a ``GraphEmbed`` class.

```python
import torch

class GraphEmbed(nn.Module):
    def __init__(self, node_hidden_size):
        super(GraphEmbed, self).__init__()

        # Setting from the paper
        self.graph_hidden_size = 2 * node_hidden_size

        # Embed graphs
        self.node_gating = nn.Sequential(
            nn.Linear(node_hidden_size, 1),
            nn.Sigmoid()
        )
        self.node_to_graph = nn.Linear(node_hidden_size,
                                       self.graph_hidden_size)

    def forward(self, g):
        if g.number_of_nodes() == 0:
            return torch.zeros(1, self.graph_hidden_size)
        else:
            # Node features are stored as hv in ndata.
            hvs = g.ndata['hv']
            return (self.node_gating(hvs) *
                    self.node_to_graph(hvs)).sum(0, keepdim=True)
```

Update node embeddings via graph propagation
'''''''''''''''''''''''''''''''''''''''''''''

The mechanism of updating node embeddings in DGMG is similar to that for
graph convolutional networks. For a node $v$ in the graph, its
neighbor $u$ sends a message to it with

\begin{align}\textbf{m}_{u\rightarrow v}=\textbf{W}_{m}\text{concat}([\textbf{h}_{v}, \textbf{h}_{u}, \textbf{x}_{u, v}]) + \textbf{b}_{m},\\\end{align}

where $\textbf{x}_{u,v}$ is the embedding of the edge between
$u$ and $v$.

After receiving messages from all its neighbors, $v$ summarizes them
with a node activation vector

\begin{align}\textbf{a}_{v} = \sum_{u: (u, v)\in E}\textbf{m}_{u\rightarrow v}\\\end{align}

and use this information to update its own feature:

\begin{align}\textbf{h}'_{v} = \textbf{GRU}(\textbf{h}_{v}, \textbf{a}_{v}).\\\end{align}

Performing all the operations above once for all nodes synchronously is
called one round of graph propagation. The more rounds of graph propagation
you perform, the longer distance messages travel throughout the graph.

With DGL, you implement graph propagation with ``g.update_all``.
The message notation here can be a bit confusing. Researchers can refer
to $\textbf{m}_{u\rightarrow v}$ as messages, however the message function
below only passes $\text{concat}([\textbf{h}_{u}, \textbf{x}_{u, v}])$.
The operation $\textbf{W}_{m}\text{concat}([\textbf{h}_{v}, \textbf{h}_{u}, \textbf{x}_{u, v}]) + \textbf{b}_{m}$
is then performed across all edges at once for efficiency consideration.

```python
from functools import partial

class GraphProp(nn.Module):
    def __init__(self, num_prop_rounds, node_hidden_size):
        super(GraphProp, self).__init__()

        self.num_prop_rounds = num_prop_rounds

        # Setting from the paper
        self.node_activation_hidden_size = 2 * node_hidden_size

        message_funcs = []
        node_update_funcs = []
        self.reduce_funcs = []

        for t in range(num_prop_rounds):
            # input being [hv, hu, xuv]
            message_funcs.append(nn.Linear(2 * node_hidden_size + 1,
                                           self.node_activation_hidden_size))

            self.reduce_funcs.append(partial(self.dgmg_reduce, round=t))
            node_update_funcs.append(
                nn.GRUCell(self.node_activation_hidden_size,
                           node_hidden_size))

        self.message_funcs = nn.ModuleList(message_funcs)
        self.node_update_funcs = nn.ModuleList(node_update_funcs)

    def dgmg_msg(self, edges):
        """For an edge u->v, return concat([h_u, x_uv])"""
        return {'m': torch.cat([edges.src['hv'],
                                edges.data['he']],
                               dim=1)}

    def dgmg_reduce(self, nodes, round):
        hv_old = nodes.data['hv']
        m = nodes.mailbox['m']
        message = torch.cat([
            hv_old.unsqueeze(1).expand(-1, m.size(1), -1), m], dim=2)
        node_activation = (self.message_funcs[round](message)).sum(1)

        return {'a': node_activation}

    def forward(self, g):
        if g.number_of_edges() > 0:
            for t in range(self.num_prop_rounds):
                g.update_all(message_func=self.dgmg_msg,
                             reduce_func=self.reduce_funcs[t])
                g.ndata['hv'] = self.node_update_funcs[t](
                     g.ndata['a'], g.ndata['hv'])
```

Actions
``````````````````````````
All actions are sampled from distributions parameterized using neural networks
and here they are in turn.

Action 1: Add nodes
''''''''''''''''''''''''''

Given the graph embedding vector $\textbf{h}_{G}$, evaluate

\begin{align}\text{Sigmoid}(\textbf{W}_{\text{add node}}\textbf{h}_{G}+b_{\text{add node}}),\\\end{align}

which is then used to parametrize a Bernoulli distribution for deciding whether
to add a new node.

If a new node is to be added, initialize its feature with

\begin{align}\textbf{W}_{\text{init}}\text{concat}([\textbf{h}_{\text{init}} , \textbf{h}_{G}])+\textbf{b}_{\text{init}},\\\end{align}

where $\textbf{h}_{\text{init}}$ is a learnable embedding module for
untyped nodes.

```python
import torch.nn.functional as F
from torch.distributions import Bernoulli

def bernoulli_action_log_prob(logit, action):
    """Calculate the log p of an action with respect to a Bernoulli
    distribution. Use logit rather than prob for numerical stability."""
    if action == 0:
        return F.logsigmoid(-logit)
    else:
        return F.logsigmoid(logit)

class AddNode(nn.Module):
    def __init__(self, graph_embed_func, node_hidden_size):
        super(AddNode, self).__init__()

        self.graph_op = {'embed': graph_embed_func}

        self.stop = 1
        self.add_node = nn.Linear(graph_embed_func.graph_hidden_size, 1)

        # If to add a node, initialize its hv
        self.node_type_embed = nn.Embedding(1, node_hidden_size)
        self.initialize_hv = nn.Linear(node_hidden_size + \
                                       graph_embed_func.graph_hidden_size,
                                       node_hidden_size)

        self.init_node_activation = torch.zeros(1, 2 * node_hidden_size)

    def _initialize_node_repr(self, g, node_type, graph_embed):
        """Whenver a node is added, initialize its representation."""
        num_nodes = g.number_of_nodes()
        hv_init = self.initialize_hv(
            torch.cat([
                self.node_type_embed(torch.LongTensor([node_type])),
                graph_embed], dim=1))
        g.nodes[num_nodes - 1].data['hv'] = hv_init
        g.nodes[num_nodes - 1].data['a'] = self.init_node_activation

    def prepare_training(self):
        self.log_prob = []

    def forward(self, g, action=None):
        graph_embed = self.graph_op['embed'](g)

        logit = self.add_node(graph_embed)
        prob = torch.sigmoid(logit)

        if not self.training:
            action = Bernoulli(prob).sample().item()
        stop = bool(action == self.stop)

        if not stop:
            g.add_nodes(1)
            self._initialize_node_repr(g, action, graph_embed)

        if self.training:
            sample_log_prob = bernoulli_action_log_prob(logit, action)

            self.log_prob.append(sample_log_prob)

        return stop
```

Action 2: Add edges
''''''''''''''''''''''''''

Given the graph embedding vector $\textbf{h}_{G}$ and the node
embedding vector $\textbf{h}_{v}$ for the latest node $v$,
you evaluate

\begin{align}\text{Sigmoid}(\textbf{W}_{\text{add edge}}\text{concat}([\textbf{h}_{G}, \textbf{h}_{v}])+b_{\text{add edge}}),\\\end{align}

which is then used to parametrize a Bernoulli distribution for deciding
whether to add a new edge starting from $v$.

```python
class AddEdge(nn.Module):
    def __init__(self, graph_embed_func, node_hidden_size):
        super(AddEdge, self).__init__()

        self.graph_op = {'embed': graph_embed_func}
        self.add_edge = nn.Linear(graph_embed_func.graph_hidden_size + \
                                  node_hidden_size, 1)

    def prepare_training(self):
        self.log_prob = []

    def forward(self, g, action=None):
        graph_embed = self.graph_op['embed'](g)
        src_embed = g.nodes[g.number_of_nodes() - 1].data['hv']

        logit = self.add_edge(torch.cat(
            [graph_embed, src_embed], dim=1))
        prob = torch.sigmoid(logit)

        if self.training:
            sample_log_prob = bernoulli_action_log_prob(logit, action)
            self.log_prob.append(sample_log_prob)
        else:
            action = Bernoulli(prob).sample().item()

        to_add_edge = bool(action == 0)
        return to_add_edge
```

Action 3: Choose a destination
'''''''''''''''''''''''''''''''''

When action 2 returns `True`, choose a destination for the
latest node $v$.

For each possible destination $u\in\{0, \cdots, v-1\}$, the
probability of choosing it is given by

\begin{align}\frac{\text{exp}(\textbf{W}_{\text{dest}}\text{concat}([\textbf{h}_{u}, \textbf{h}_{v}])+\textbf{b}_{\text{dest}})}{\sum_{i=0}^{v-1}\text{exp}(\textbf{W}_{\text{dest}}\text{concat}([\textbf{h}_{i}, \textbf{h}_{v}])+\textbf{b}_{\text{dest}})}\\\end{align}

```python
from torch.distributions import Categorical

class ChooseDestAndUpdate(nn.Module):
    def __init__(self, graph_prop_func, node_hidden_size):
        super(ChooseDestAndUpdate, self).__init__()

        self.graph_op = {'prop': graph_prop_func}
        self.choose_dest = nn.Linear(2 * node_hidden_size, 1)

    def _initialize_edge_repr(self, g, src_list, dest_list):
        # For untyped edges, only add 1 to indicate its existence.
        # For multiple edge types, use a one-hot representation
        # or an embedding module.
        edge_repr = torch.ones(len(src_list), 1)
        g.edges[src_list, dest_list].data['he'] = edge_repr

    def prepare_training(self):
        self.log_prob = []

    def forward(self, g, dest):
        src = g.number_of_nodes() - 1
        possible_dests = range(src)

        src_embed_expand = g.nodes[src].data['hv'].expand(src, -1)
        possible_dests_embed = g.nodes[possible_dests].data['hv']

        dests_scores = self.choose_dest(
            torch.cat([possible_dests_embed,
                       src_embed_expand], dim=1)).view(1, -1)
        dests_probs = F.softmax(dests_scores, dim=1)

        if not self.training:
            dest = Categorical(dests_probs).sample().item()

        if not g.has_edge_between(src, dest):
            # For undirected graphs, add edges for both directions
            # so that you can perform graph propagation.
            src_list = [src, dest]
            dest_list = [dest, src]

            g.add_edges(src_list, dest_list)
            self._initialize_edge_repr(g, src_list, dest_list)

            self.graph_op['prop'](g)

        if self.training:
            if dests_probs.nelement() > 1:
                self.log_prob.append(
                    F.log_softmax(dests_scores, dim=1)[:, dest: dest + 1])
```

Putting it together
``````````````````````````

You are now ready to have a complete implementation of the model class.

```python
class DGMG(DGMGSkeleton):
    def __init__(self, v_max, node_hidden_size,
                 num_prop_rounds):
        super(DGMG, self).__init__(v_max)

        # Graph embedding module
        self.graph_embed = GraphEmbed(node_hidden_size)

        # Graph propagation module
        self.graph_prop = GraphProp(num_prop_rounds,
                                    node_hidden_size)

        # Actions
        self.add_node_agent = AddNode(
            self.graph_embed, node_hidden_size)
        self.add_edge_agent = AddEdge(
            self.graph_embed, node_hidden_size)
        self.choose_dest_agent = ChooseDestAndUpdate(
            self.graph_prop, node_hidden_size)

        # Forward functions
        self.forward_train = partial(forward_train, self=self)
        self.forward_inference = partial(forward_inference, self=self)

    @property
    def action_step(self):
        old_step_count = self.step_count
        self.step_count += 1

        return old_step_count

    def prepare_for_train(self):
        self.step_count = 0

        self.add_node_agent.prepare_training()
        self.add_edge_agent.prepare_training()
        self.choose_dest_agent.prepare_training()

    def add_node_and_update(self, a=None):
        """Decide if to add a new node.
        If a new node should be added, update the graph."""

        return self.add_node_agent(self.g, a)

    def add_edge_or_not(self, a=None):
        """Decide if a new edge should be added."""

        return self.add_edge_agent(self.g, a)

    def choose_dest_and_update(self, a=None):
        """Choose destination and connect it to the latest node.
        Add edges for both directions and update the graph."""

        self.choose_dest_agent(self.g, a)

    def get_log_prob(self):
        add_node_log_p = torch.cat(self.add_node_agent.log_prob).sum()
        add_edge_log_p = torch.cat(self.add_edge_agent.log_prob).sum()
        choose_dest_log_p = torch.cat(self.choose_dest_agent.log_prob).sum()
        return add_node_log_p + add_edge_log_p + choose_dest_log_p
```

Below is an animation where a graph is generated on the fly
after every 10 batches of training for the first 400 batches. You
can see how the model improves over time and begins generating cycles.

.. figure:: https://user-images.githubusercontent.com/19576924/48929291-60fe3880-ef22-11e8-832a-fbe56656559a.gif
   :alt:

For generative models, you can evaluate performance by checking the percentage
of valid graphs among the graphs it generates on the fly.

```python
import torch.utils.model_zoo as model_zoo

# Download a pre-trained model state dict for generating cycles with 10-20 nodes.
state_dict = model_zoo.load_url('https://data.dgl.ai/model/dgmg_cycles-5a0c40be.pth')
model = DGMG(v_max=20, node_hidden_size=16, num_prop_rounds=2)
model.load_state_dict(state_dict)
model.eval()

def is_valid(g):
    # Check if g is a cycle having 10-20 nodes.
    def _get_previous(i, v_max):
        if i == 0:
            return v_max
        else:
            return i - 1

    def _get_next(i, v_max):
        if i == v_max:
            return 0
        else:
            return i + 1

    size = g.number_of_nodes()

    if size < 10 or size > 20:
        return False

    for node in range(size):
        neighbors = g.successors(node)

        if len(neighbors) != 2:
            return False

        if _get_previous(node, size - 1) not in neighbors:
            return False

        if _get_next(node, size - 1) not in neighbors:
            return False

    return True

num_valid = 0
for i in range(100):
    g = model()
    num_valid += is_valid(g)

del model
print('Among 100 graphs generated, {}% are valid.'.format(num_valid))
```

    Downloading: "https://data.dgl.ai/model/dgmg_cycles-5a0c40be.pth" to /home/viktor/.cache/torch/hub/checkpoints/dgmg_cycles-5a0c40be.pth
    100.0%
    <ipython-input-8-582f9090524b>:34: DeprecationWarning: an integer is required (got type float).  Implicit conversion to integers using __int__ is deprecated, and may be removed in a future version of Python.
      self.node_type_embed(torch.LongTensor([node_type])),
    /home/viktor/anaconda3/envs/dgl/lib/python3.8/site-packages/dgl/base.py:45: DGLWarning: DGLGraph.has_edge_between is deprecated. Please use DGLGraph.has_edges_between
      return warnings.warn(message, category=category, stacklevel=1)

    Among 100 graphs generated, 97% are valid.

For the complete implementation, see the `DGL DGMG example
<https://github.com/dmlc/dgl/tree/master/examples/pytorch/dgmg>`__.


```python
%matplotlib inline
```

.. currentmodule:: dgl

Working with Heterogeneous Graphs
=================================

**Author**: Quan Gan, `Minjie Wang <https://jermainewang.github.io/>`_, Mufei Li,
George Karypis, Zheng Zhang

In this tutorial, you learn about:

* Examples of heterogenous graph data and typical applications.

* Creating and manipulating a heterogenous graph in DGL.

* Implementing `Relational-GCN <https://arxiv.org/abs/1703.06103>`_, a popular GNN model,
  for heterogenous graph input.

* Training a model to solve a node classification task.

Heterogeneous graphs, or *heterographs* for short, are graphs that contain
different types of nodes and edges. The different types of nodes and edges tend
to have different types of attributes that are designed to capture the
characteristics of each node and edge type. Within the context of
graph neural networks, depending on their complexity, certain node and edge types
might need to be modeled with representations that have a different number of dimensions.

DGL supports graph neural network computations on such heterogeneous graphs, by
using the heterograph class and its associated API.

Examples of heterographs
-----------------------
Many graph datasets represent relationships among various types of entities.
This section provides an overview for several graph use-cases that show such relationships
and can have their data represented as heterographs.

Citation graph
~~~~~~~~~~~~~~~
The Association for Computing Machinery publishes an `ACM dataset <https://aminer.org/citation>`_ that contains two
million papers, their authors, publication venues, and the other papers
that were cited. This information can be represented as a heterogeneous graph.

The following diagram shows several entities in the ACM dataset and the relationships among them
(taken from `Shi et al., 2015 <https://arxiv.org/pdf/1511.04854.pdf>`_).

.. figure:: https://data.dgl.ai/tutorial/hetero/acm-example.png#

This graph has three types of entities that correspond to papers, authors, and publication venues.
It also contains three types of edges that connect the following:

* Authors with papers corresponding to *written-by* relationships

* Papers with publication venues corresponding to *published-in* relationships

* Papers with other papers corresponding to *cited-by* relationships

Recommender systems
~~~~~~~~~~~~~~~~~~~~
The datasets used in recommender systems often contain
interactions between users and items. For example, the data could include the
ratings that users have provided to movies. Such interactions can be modeled
as heterographs.

The nodes in these heterographs will have two types, *users* and *movies*. The edges
will correspond to the user-movie interactions. Furthermore, if an interaction is
marked with a rating, then each rating value could correspond to a different edge type.
The following diagram shows an example of user-item interactions as a heterograph.

.. figure:: https://data.dgl.ai/tutorial/hetero/recsys-example.png

Knowledge graph
~~~~~~~~~~~~~~~~
Knowledge graphs are inherently heterogenous. For example, in
Wikidata, Barack Obama (item Q76) is an instance of a human, which could be viewed as
the entity class, whose spouse (item P26) is Michelle Obama (item Q13133) and
occupation (item P106) is politician (item Q82955). The relationships are shown in the following.
diagram.

.. figure:: https://data.dgl.ai/tutorial/hetero/kg-example.png

Creating a heterograph in DGL
-----------------------------
You can create a heterograph in DGL using the :func:`dgl.heterograph` API.
The argument to :func:`dgl.heterograph` is a dictionary. The keys are tuples
in the form of ``(srctype, edgetype, dsttype)`` specifying the relation name
and the two entity types it connects. Such tuples are called *canonical edge
types*. The values are data to initialize the graph structures, that is, which
nodes the edges actually connect.

For instance, the following code creates the user-item interactions heterograph shown earlier.

```python
# Each value of the dictionary is a pair of source and destination arrays.
# Nodes are integer IDs starting from zero. Nodes IDs of different types have
# separate countings.
import dgl
import numpy as np

ratings = dgl.heterograph(
    {('user', '+1', 'movie') : (np.array([0, 0, 1]), np.array([0, 1, 0])),
     ('user', '-1', 'movie') : (np.array([2]), np.array([1]))})
```

DGL supports creating a graph from a variety of data sources. The following
code creates the same graph as the above.

Creating from scipy matrix

```python
import scipy.sparse as sp
plus1 = sp.coo_matrix(([1, 1, 1], ([0, 0, 1], [0, 1, 0])), shape=(3, 2))
minus1 = sp.coo_matrix(([1], ([2], [1])), shape=(3, 2))
ratings = dgl.heterograph(
    {('user', '+1', 'movie') : plus1,
     ('user', '-1', 'movie') : minus1})

# Creating from networkx graph
import networkx as nx
plus1 = nx.DiGraph()
plus1.add_nodes_from(['u0', 'u1', 'u2'], bipartite=0)
plus1.add_nodes_from(['m0', 'm1'], bipartite=1)
plus1.add_edges_from([('u0', 'm0'), ('u0', 'm1'), ('u1', 'm0')])
# To simplify the example, reuse the minus1 object.
# This also means that you could use different sources of graph data
# for different relationships.
ratings = dgl.heterograph(
    {('user', '+1', 'movie') : plus1,
     ('user', '-1', 'movie') : minus1})
```

Manipulating heterograph
------------------------
You can create a more realistic heterograph using the ACM dataset. To do this, first
download the dataset as follows:

```python
import scipy.io
import urllib.request

data_url = 'https://data.dgl.ai/dataset/ACM.mat'
data_file_path = '/tmp/ACM.mat'

urllib.request.urlretrieve(data_url, data_file_path)
data = scipy.io.loadmat(data_file_path)
print(list(data.keys()))
```

The dataset stores node information by their types: ``P`` for paper, ``A``
for author, ``C`` for conference, ``L`` for subject code, and so on. The relationships
are stored as SciPy sparse matrix under key ``XvsY``, where ``X`` and ``Y``
could be any of the node type code.

The following code prints out some statistics about the paper-author relationships.

```python
print(type(data['PvsA']))
print('#Papers:', data['PvsA'].shape[0])
print('#Authors:', data['PvsA'].shape[1])
print('#Links:', data['PvsA'].nnz)
```

Converting this SciPy matrix to a heterograph in DGL is straightforward.

```python
pa_g = dgl.heterograph({('paper', 'written-by', 'author') : data['PvsA']})
# equivalent (shorter) API for creating heterograph with two node types:
pa_g = dgl.bipartite(data['PvsA'], 'paper', 'written-by', 'author')
```

You can easily print out the type names and other structural information.

```python
print('Node types:', pa_g.ntypes)
print('Edge types:', pa_g.etypes)
print('Canonical edge types:', pa_g.canonical_etypes)

# Nodes and edges are assigned integer IDs starting from zero and each type has its own counting.
# To distinguish the nodes and edges of different types, specify the type name as the argument.
print(pa_g.number_of_nodes('paper'))
# Canonical edge type name can be shortened to only one edge type name if it is
# uniquely distinguishable.
print(pa_g.number_of_edges(('paper', 'written-by', 'author')))
print(pa_g.number_of_edges('written-by'))
print(pa_g.successors(1, etype='written-by'))  # get the authors that write paper #1

# Type name argument could be omitted whenever the behavior is unambiguous.
print(pa_g.number_of_edges())  # Only one edge type, the edge type argument could be omitted
```

A homogeneous graph is just a special case of a heterograph with only one type
of node and edge. In this case, all the APIs are exactly the same as in
:class:`DGLGraph`.

```python
# Paper-citing-paper graph is a homogeneous graph
pp_g = dgl.heterograph({('paper', 'citing', 'paper') : data['PvsP']})
# equivalent (shorter) API for creating homogeneous graph
pp_g = dgl.graph(data['PvsP'], 'paper', 'cite')

# All the ntype and etype arguments could be omitted because the behavior is unambiguous.
print(pp_g.number_of_nodes())
print(pp_g.number_of_edges())
print(pp_g.successors(3))
```

Create a subset of the ACM graph using the paper-author, paper-paper,
and paper-subject relationships.  Meanwhile, also add the reverse
relationship to prepare for the later sections.

```python
G = dgl.heterograph({
        ('paper', 'written-by', 'author') : data['PvsA'],
        ('author', 'writing', 'paper') : data['PvsA'].transpose(),
        ('paper', 'citing', 'paper') : data['PvsP'],
        ('paper', 'cited', 'paper') : data['PvsP'].transpose(),
        ('paper', 'is-about', 'subject') : data['PvsL'],
        ('subject', 'has', 'paper') : data['PvsL'].transpose(),
    })

print(G)
```

**Metagraph** (or network schema) is a useful summary of a heterograph.
Serving as a template for a heterograph, it tells how many types of objects
exist in the network and where the possible links exist.

DGL provides easy access to the metagraph, which could be visualized using
external tools.

```python
# Draw the metagraph using graphviz.
import pygraphviz as pgv
def plot_graph(nxg):
    ag = pgv.AGraph(strict=False, directed=True)
    for u, v, k in nxg.edges(keys=True):
        ag.add_edge(u, v, label=k)
    ag.layout('dot')
    ag.draw('graph.png')

plot_graph(G.metagraph)
```

Learning tasks associated with heterographs
-------------------------------------------
Some of the typical learning tasks that involve heterographs include:

* *Node classification and regression* to predict the class of each node or
  estimate a value associated with it.

* *Link prediction* to predict if there is an edge of a certain
  type between a pair of nodes, or predict which other nodes a particular
  node is connected with (and optionally the edge types of such connections).

* *Graph classification/regression* to assign an entire
  heterograph into one of the target classes or to estimate a numerical
  value associated with it.

In this tutorial, we designed a simple example for the first task.

A semi-supervised node classification example
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Our goal is to predict the publishing conference of a paper using the ACM
academic graph we just created. To further simplify the task, we only focus
on papers published in three conferences: *KDD*, *ICML*, and *VLDB*. All
the other papers are not labeled, making it a semi-supervised setting.

The following code extracts those papers from the raw dataset and prepares
the training, validation, testing split.

```python
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

pvc = data['PvsC'].tocsr()
# find all papers published in KDD, ICML, VLDB
c_selected = [0, 11, 13]  # KDD, ICML, VLDB
p_selected = pvc[:, c_selected].tocoo()
# generate labels
labels = pvc.indices
labels[labels == 11] = 1
labels[labels == 13] = 2
labels = torch.tensor(labels).long()

# generate train/val/test split
pid = p_selected.row
shuffle = np.random.permutation(pid)
train_idx = torch.tensor(shuffle[0:800]).long()
val_idx = torch.tensor(shuffle[800:900]).long()
test_idx = torch.tensor(shuffle[900:]).long()
```

Relational-GCN on heterograph
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We use `Relational-GCN <https://arxiv.org/abs/1703.06103>`_ to learn the
representation of nodes in the graph. Its message-passing equation is as
follows:

\begin{align}h_i^{(l+1)} = \sigma\left(\sum_{r\in \mathcal{R}}
   \sum_{j\in\mathcal{N}_r(i)}W_r^{(l)}h_j^{(l)}\right)\end{align}

Breaking down the equation, you see that there are two parts in the
computation.

(i) Message computation and aggregation within each relation $r$

(ii) Reduction that merges the results from multiple relationships

Following this intuition, perform message passing on a heterograph in
two steps.

(i) Per-edge-type message passing

(ii) Type wise reduction

```python
import dgl.function as fn

class HeteroRGCNLayer(nn.Module):
    def __init__(self, in_size, out_size, etypes):
        super(HeteroRGCNLayer, self).__init__()
        # W_r for each relation
        self.weight = nn.ModuleDict({
                name : nn.Linear(in_size, out_size) for name in etypes
            })

    def forward(self, G, feat_dict):
        # The input is a dictionary of node features for each type
        funcs = {}
        for srctype, etype, dsttype in G.canonical_etypes:
            # Compute W_r * h
            Wh = self.weight[etype](feat_dict[srctype])
            # Save it in graph for message passing
            G.nodes[srctype].data['Wh_%s' % etype] = Wh
            # Specify per-relation message passing functions: (message_func, reduce_func).
            # Note that the results are saved to the same destination feature 'h', which
            # hints the type wise reducer for aggregation.
            funcs[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'h'))
        # Trigger message passing of multiple types.
        # The first argument is the message passing functions for each relation.
        # The second one is the type wise reducer, could be "sum", "max",
        # "min", "mean", "stack"
        G.multi_update_all(funcs, 'sum')
        # return the updated node feature dictionary
        return {ntype : G.nodes[ntype].data['h'] for ntype in G.ntypes}
```

Create a simple GNN by stacking two ``HeteroRGCNLayer``. Since the
nodes do not have input features, make their embeddings trainable.

```python
class HeteroRGCN(nn.Module):
    def __init__(self, G, in_size, hidden_size, out_size):
        super(HeteroRGCN, self).__init__()
        # Use trainable node embeddings as featureless inputs.
        embed_dict = {ntype : nn.Parameter(torch.Tensor(G.number_of_nodes(ntype), in_size))
                      for ntype in G.ntypes}
        for key, embed in embed_dict.items():
            nn.init.xavier_uniform_(embed)
        self.embed = nn.ParameterDict(embed_dict)
        # create layers
        self.layer1 = HeteroRGCNLayer(in_size, hidden_size, G.etypes)
        self.layer2 = HeteroRGCNLayer(hidden_size, out_size, G.etypes)

    def forward(self, G):
        h_dict = self.layer1(G, self.embed)
        h_dict = {k : F.leaky_relu(h) for k, h in h_dict.items()}
        h_dict = self.layer2(G, h_dict)
        # get paper logits
        return h_dict['paper']
```

Train and evaluate
~~~~~~~~~~~~~~~~~~
Train and evaluate this network.

```python
# Create the model. The output has three logits for three classes.
model = HeteroRGCN(G, 10, 10, 3)

opt = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

best_val_acc = 0
best_test_acc = 0

for epoch in range(100):
    logits = model(G)
    # The loss is computed only for labeled nodes.
    loss = F.cross_entropy(logits[train_idx], labels[train_idx])

    pred = logits.argmax(1)
    train_acc = (pred[train_idx] == labels[train_idx]).float().mean()
    val_acc = (pred[val_idx] == labels[val_idx]).float().mean()
    test_acc = (pred[test_idx] == labels[test_idx]).float().mean()

    if best_val_acc < val_acc:
        best_val_acc = val_acc
        best_test_acc = test_acc

    opt.zero_grad()
    loss.backward()
    opt.step()

    if epoch % 5 == 0:
        print('Loss %.4f, Train Acc %.4f, Val Acc %.4f (Best %.4f), Test Acc %.4f (Best %.4f)' % (
            loss.item(),
            train_acc.item(),
            val_acc.item(),
            best_val_acc.item(),
            test_acc.item(),
            best_test_acc.item(),
        ))
```

What's next?
------------
* Check out our full implementation in PyTorch
  `here <https://github.com/dmlc/dgl/tree/master/examples/pytorch/rgcn-hetero>`_.

* We also provide the following model examples:

  * `Graph Convolutional Matrix Completion <https://arxiv.org/abs/1706.02263>_`,
    which we implement in MXNet
    `here <https://github.com/dmlc/dgl/tree/v0.4.0/examples/mxnet/gcmc>`_.

  * `Heterogeneous Graph Attention Network <https://arxiv.org/abs/1903.07293>`_
    requires transforming a heterograph into a homogeneous graph according to
    a given metapath (i.e. a path template consisting of edge types).  We
    provide :func:`dgl.transform.metapath_reachable_graph` to do this.  See full
    implementation
    `here <https://github.com/dmlc/dgl/tree/master/examples/pytorch/han>`_.

  * `Metapath2vec <https://dl.acm.org/citation.cfm?id=3098036>`_ requires
    generating random walk paths according to a given metapath.  Please
    refer to the full metapath2vec implementation
    `here <https://github.com/dmlc/dgl/tree/master/examples/pytorch/metapath2vec>`_.

* :doc:`Full heterograph API reference <../../api/python/heterograph>`.