Try   HackMD

How to implement a GNN

import torch
from torch_scatter import scatter

We will use this graph:

num_nodes = 5
num_edges = 6
num_edge_types = 3

edge_index = torch.LongTensor([
    [0,1], [0,4], [2,0], [2,3], [3,1], [2,4]
]).t()

edge_index

​​​​tensor([[0, 0, 2, 2, 3, 2],
​​​​        [1, 4, 0, 3, 1, 4]])

embed_dim = 10

x = torch.randn(num_nodes, embed_dim)
x.size()

​​​​torch.Size([5, 10])

Define the MessagePassing class

import inspect
import torch.nn as nn


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

        # get the list of arguments as a list of `string`
        self.message_args = inspect.getfullargspec(self.message)[0][1:]
        self.update_args = inspect.getfullargspec(self.update)[0][2:]

    def propagate(self, aggr, edge_index, **kwargs):
        '''
        edge_index.size(): [2, number_of_edges]
        '''
        assert aggr in ['sum', 'mean', 'max']
        kwargs['edge_index'] = edge_index

        dim_size = None
        message_args = []
        for arg in self.message_args:
            if arg.endswith('_i') or arg.endswith('_j'):    # will almost always be `x_i` or `x_j`
                x = kwargs[arg[:-2]]                        # get `x`, has shape [number_of_nodes, embed_dim]
                dim_size = x.size(0)
                j = 0 if arg.endswith('_i') else 1
                x_j = x[edge_index[j]]
                message_args.append(x_j)
            else:
                message_args.append(kwargs[arg])

        update_args = [kwargs[arg] for arg in self.update_args]

        out = self.message(*message_args)
        out = scatter(out, edge_index[0], dim=0, dim_size=dim_size, reduce=aggr)
        return self.update(out, *update_args)
            
    def message(self):
        raise NotImplementedError

    def update(self):
        raise NotImplementedError

Create a nn.Module that inherits MessagePassing

class MyGNN(MessagePassing):
    def __init__(self, embed_dim=embed_dim):
        super(MyGNN, self).__init__()
        self.embed_dim = embed_dim
        self.message_output = None
        self.update_output = None
        
    def message(self, x_j, a_random_argument):
        '''
        `x_j` has shape (num_edges, embed_dim)
        `a_random_argument` can be any argument. In this case it is a tensor of shape `(embed_dim,)`
        '''
        print(f'\tInside message method. `x_j.size()` is {x_j.size()}')
        output = x_j * a_random_argument.view(1,-1)
        self.message_output = output
        return output
    
    def update(self, out):
        print(f'\tInside update method. `out.size()` is {out.size()}')
        output = out
        self.update_output = output
        return output
    
    def forward(self, x, edge_index):
        return self.propagate('mean', edge_index, x=x, a_random_argument=torch.ones(self.embed_dim)*2)

model = MyGNN()

model_output = model(x, edge_index)
model_output.size()

​​​​	Inside message method. `x_j.size()` is torch.Size([6, 10])
​​​​	Inside update method. `out.size()` is torch.Size([5, 10])


​​​​torch.Size([5, 10])

model.message_output

​​​​tensor([[-3.9773, -1.1126, -0.0140,  3.8670,  1.2074, -1.7092,  0.3560,  3.2655,
​​​​         -2.3335, -4.1227],
​​​​        [-1.3679, -3.1507, -1.1292,  1.1108, -2.8639,  1.0764, -1.2942, -2.1353,
​​​​         -1.2333, -0.1250],
​​​​        [ 0.2759, -3.5196, -0.0355, -0.4968,  0.6988,  0.3293, -1.4655,  1.2332,
​​​​         -0.5728, -0.9543],
​​​​        [ 0.1292,  0.8452, -1.5253,  4.7322, -0.6898, -0.2914, -0.8749,  0.1460,
​​​​          0.2724,  1.6101],
​​​​        [-3.9773, -1.1126, -0.0140,  3.8670,  1.2074, -1.7092,  0.3560,  3.2655,
​​​​         -2.3335, -4.1227],
​​​​        [-1.3679, -3.1507, -1.1292,  1.1108, -2.8639,  1.0764, -1.2942, -2.1353,
​​​​         -1.2333, -0.1250]])

model.update_output

​​​​tensor([[-2.6726, -2.1317, -0.5716,  2.4889, -0.8283, -0.3164, -0.4691,  0.5651,
​​​​         -1.7834, -2.1239],
​​​​        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
​​​​          0.0000,  0.0000],
​​​​        [-0.3209, -1.9417, -0.8967,  1.7820, -0.9516,  0.3714, -1.2115, -0.2520,
​​​​         -0.5112,  0.1769],
​​​​        [-3.9773, -1.1126, -0.0140,  3.8670,  1.2074, -1.7092,  0.3560,  3.2655,
​​​​         -2.3335, -4.1227],
​​​​        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
​​​​          0.0000,  0.0000]])

Now let's run these results manually

Before calling message, propagate moves messages from nodes to edges.

This can be done in two ways:

  1. For each edge, move the message (node embedding) of the source node: happens if x_i is an argument of message method
  2. For each edge, move the message (node embedding) of the target node: happens if x_j is an argument of message method

edge_index

​​​​tensor([[0, 0, 2, 2, 3, 2],
​​​​        [1, 4, 0, 3, 1, 4]])

torch.equal(x[edge_index[1]] * 2, model.message_output)

​​​​True

Now let's modify the argument of message to x_i

class MySecondGNN(MessagePassing):
    def __init__(self, embed_dim=embed_dim):
        super(MySecondGNN, self).__init__()
        self.embed_dim = embed_dim
        self.message_output = None
        self.update_output = None
        
    def message(self, x_i, a_random_argument):
        '''
        `x_j` has shape (num_edges, embed_dim)
        `a_random_argument` can be any argument. In this case it is a tensor of shape `(embed_dim,)`
        '''
        print(f'\tInside message method. `x_j.size()` is {x_i.size()}')
        output = x_i * a_random_argument.view(1,-1)
        self.message_output = output
        return output
    
    def update(self, out):
        print(f'\tInside update method. `out.size()` is {out.size()}')
        output = out
        self.update_output = output
        return output
    
    def forward(self, x, edge_index):
        return self.propagate('mean', edge_index, x=x, a_random_argument=torch.ones(self.embed_dim)*2)

model = MySecondGNN()

model_output = model(x, edge_index)
model_output.size()

​​​​	Inside message method. `x_j.size()` is torch.Size([6, 10])
​​​​	Inside update method. `out.size()` is torch.Size([5, 10])





​​​​torch.Size([5, 10])

torch.equal(x[edge_index[0]] * 2, model.message_output)

​​​​True

Now let's take a look at what happens before update is called: the messages are put back from edges to nodes

Note that this is not one to one - a node can have many edges. As a result it can receive messages from many edges. Thus we need a way to aggregate these messages. This is exactly what the aggr function does in scatter.

edge_index

​​​​tensor([[0, 0, 2, 2, 3, 2],
​​​​        [1, 4, 0, 3, 1, 4]])

Messages received by node 0:

message_output = x[edge_index[0]] * 2
torch.equal(message_output[[0,1]].mean(dim=0), model.update_output[0])

​​​​True

Messages received by node 2:

torch.equal(message_output[[2,3,5]].mean(dim=0), model.update_output[2])

​​​​True

We can pass messages to all the nodes using scatter

scatter(message_output, edge_index[0], dim=0, dim_size=num_nodes, reduce='mean')

​​​​tensor([[ 0.2759, -3.5196, -0.0355, -0.4968,  0.6988,  0.3293, -1.4655,  1.2332,
​​​​         -0.5728, -0.9543],
​​​​        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
​​​​          0.0000,  0.0000],
​​​​        [-0.5084,  2.1698,  2.5749, -1.7214, -2.0338,  0.1578, -0.6327, -0.6667,
​​​​          2.7229,  2.0444],
​​​​        [ 0.1292,  0.8452, -1.5253,  4.7322, -0.6898, -0.2914, -0.8749,  0.1460,
​​​​          0.2724,  1.6101],
​​​​        [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,
​​​​          0.0000,  0.0000]])

torch.equal(scatter(message_output, edge_index[0], dim=0, dim_size=num_nodes, reduce='mean'), model.update_output)

​​​​True
tags: pytorch, gnn
Last changed by 
 
vinsis
0
163

Read more

Metaformer is Actually What You Need for Vision [CVPR 2022]

Paper: https://arxiv.org/abs/2111.11418 Key idea: abstract the network architecture from high performing models like Transformers, MLP-Mixers etc. It is this network that gives good performance. They replace transformer, MLP-mixer etc with pooling to prove this statement. The main thing to understand is how pooling works: class Pooling(nn.Module): """ Implementation of pooling for PoolFormer --pool_size: pooling size """

Dec 9, 2022
Patches are all you need?

Main claim: patches are what lead to an improved performace at least to a certain extent Stem implementation self.stem = nn.Sequential( nn.Conv2d(in_chans, dim, kernel_size=patch_size, stride=patch_size), activation(), nn.BatchNorm2d(dim) ) Blocks implementation

Dec 9, 2022
MLP-Mixer: An All-MLP Architecture for Vision

Mostly the same people behind ViT paper. Adequate (84.15 top 1 on ImageNet by Mixer-L/16) but not SOTA. Benefits much more from scaling up. Common part with ViT: Divide an image into NxN patches, unroll each patch and do a linear transform. Some simple Mlp = Linear -> Activation -> Dropout -> Linear -> Dropout style MLP layers implemented here class Mlp(nn.Module): """ MLP as used in Vision Transformer, MLP-Mixer and related networks """ def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): super().__init__()

Dec 9, 2022
4. BeIT (Bert Encoding representations for Image Transformers)

Another self-supervised image representation model Key idea: During pre-training, we randomly mask some proportion of image patches, and feed the corrupted input to Transformer. The model learns to recover the visual tokens of the original image, instead of the raw pixels of masked patches. Image patches are flattened into vectors and linearly projected. In our experiments, we split each 224 × 224 image into a 14 × 14 grid of image patches, where each patch is 16 × 16. An image is broken down into patches. Each patch has two representations: patch representation via transformer encoder and tokens from a fixed vocabulary.

Dec 9, 2022
Read more from vinsis
Published on HackMD