gnn-layers 1.0.1
Some custom GNN layers for PyTorch Geometric
License: MIT
Requires Python: >=3
Keywords: gnn,graph-neural-network,convolution,pooling,pytorch,graph

# custom-gnn-layers

This is a little collection of the custom graph convolutional and pooling layers I've made for various projects. Everything here is built on the [PyTorch Geometric](https://github.com/rusty1s/pytorch_geometric) library and can be used like a regular PyTorch module.

## Installation

> Requires Python 3.X, PyTorch 1.2.X, and PyTorch Geometric

You can install these layers with pip:

`$ pip install gnn_layers`

## Convolutional Layers

### EdgeAttentionConv

`EdgeAttentionConv` is an edge-conditioned filter with an attention mechanism. It's the same as [`NNConv`](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.NNConv), except an attention coefficient for each message is calculated from the edge features. The idea is that messages from some neighbors may be more important than others, depending on their connection with the root node. Node embeddings are updated like so:

<p align="center">
  <img align="center" src="assets/edge_attention_eq.png" width="50%" />
</p>

where _W<sub>r</sub>_ and _W<sub>g</sub>_ are trainable weight matrices, and _h<sub>e</sub>_ is a neural network (e.g. a MLP). _W<sub>r</sub>_ is used to transform the root node features and _W<sub>g</sub>_ is used to calculate an attention coefficient.

**Parameters:**

- **in_channels** (_int_): Size of each input node embedding.
- **out_channels** (_int_): Size of each output node embedding.
- **edge_nn** (_torch.nn.Module_): A neural network _h<sub>e</sub>_ that maps edge features `edge_attr` of shape `[-1, num_edge_features]` to shape `[-1, in_channels * out_channels]`
- **root_weight** (_bool, optional_): If set to `False`, the layer will not add the transformed root node features to the output. (default: `True`)
- **bias** (_bool, optional_): If set to `False`, the layer will not learn an additive bias. (default: `True`)

**Example:**

```python
import torch
from gnn_layers import EdgeAttentionConv

# Convolutional layer
conv = EdgeAttentionConv(
  in_channels=1,
  out_channels=4,
  edge_nn=torch.nn.Linear(2, 1 * 4),
)

# Your input graph
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
edge_attr = torch.tensor([[0, 1], [1, 0], [1, 0], [0, 1]], dtype=torch.long)

# Your output graph
x = conv(x, edge_index, edge_attr) # Shape is now [3, 4]
```

### ImageConv

`ImageConv` is an edge-conditioned filter for graphs where the node features are n-dimensional matrices, such as 2D or 3D images, rather than vectors. The filter applies a (non-graph) convolution, i.e. `torch.nn.Conv2d` or `torch.nn.Conv3d`, to transform the node features. Node embeddings are updated like so:

<p align="center">
  <img src="assets/image_eq.png" width="47%" />
</p>

where _&phi;<sub>r</sub>_ and _&phi;<sub>m</sub>_ are convolutional layers, and _W<sub>e</sub>_ is a weight matrix.

**Parameters:**

- **in_channels** (_int_): Number of channels in the input node image.
- **out_channels** (_int_): Number of channels in the output node image.
- **image_dims** (_tuple_): Dimensions of the input node image as a tuple, e.g. for a 4x4 image, set to `(4, 4)`.
- **kernel_size** (_tuple_): Size of the convolving kernel.
- **num_edge_attr** (_int_): Number of edge features.
- **bias** (_bool, optional_): If set to `False`, the layer will not learn an additive bias. (default: `True`)
- **aggr** (_str, optional_): The aggregation scheme to use (`"add"`, `"mean"`, `"max"`). (default: `"add"`)
- **\*\*kwargs** (_optional_): Additional arguments for `torch.nn.Conv1d`, `torch.nn.Conv2d`, or `torch.nn.Conv3d`.

**Example:**

```python
import torch
from gnn_layers import ImageConv

# Convolutional layer
conv = ImageConv(
  in_channels=1,
  out_channels=4,
  image_dims=(8, 8),
  kernel_size=(2, 2),
  num_edge_attr=2
)

# Your input graph
x = torch.randn((3, 1, 8, 8), dtype=torch.float)
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
edge_attr = torch.tensor([[0, 1], [1, 0], [1, 0], [0, 1]], dtype=torch.long)

# Your output graph
x = conv(x, edge_index, edge_attr) # Shape is now [3, 4, 7, 7]
```

<!--This is useful for learning on images that have some sort of embedded relational structure. For example, ...-->

<!-- ## Pooling Layers

### SimPool

SimPool can either learn which nodes to merge together or which edges to drop.

### DiffPool

### GlobalEdgePool

Average the edge features and add them to the graph embedding. -->