Graph Neural Networks con PyTorch Geometric

En este post vamos a ver una introducción a las Graph Neural Networks (GNNs) y la librería PyTorch Geometric, que nos permite aplicar el deep learning a datos no estructurados, como grafos, empleando este tipo de modelos. Las GNNs permiten aplicar redes neuronales a grafos, teniendo en cuenta su estructura basada en nodos conectados entre sí. Para ello, cada estado x del nodo v del grafo G=(V,E) se va actualizando iterativamente agregando información de los nodos vecinos N(v):

Lo primero que debemos conocer de PyTorch Geometric es la clase torch_geometric.data.Data que almacena los datos del grafo. Esta clase tiene los siguientes atributos:

  • data.x: La matriz con los features de cada nodo del grafo. Tiene dimensión [num_nodes, num_node_features]
  • data.edge_index: La conectividad del grafo en formato COO. Para cada índice, representa el nodo origen y el destino. Tiene dimensión [2, num_edges].
  • data.edge_attr: La matriz con los features de los enlaces del grafo.
  • data.y: Las etiquetas con las que entrenar. P.e. una etiqueta para cada nodo o para una para el grafo completo.

Vamos a ver un ejemplo de esta clase. Primero vamos a instalar las librerías de PyTorch Geometric.

!pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.8.0+cu101.html
!pip install torch-geometric

De los datasets que vienen con la librería, vamos a importar Cora, una red de citas donde los nodos representan documentos, y ver sus características.

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='/tmp/Cora', name='Cora')

data = dataset[0]
print(data)

print(f'Número de nodos: {data.num_nodes}')
print(f'Número de features por nodo: {data.num_node_features}')
print(f'Número de clases: {dataset.num_classes}')
print(f'Número de enlaces: {data.num_edges}')
print(f'Grado medio de los nodos: {data.num_edges / data.num_nodes:.2f}')
print(f'Número de nodos de entrenamiento: {data.train_mask.sum()}')
print(f'Número de nodos de validación: {data.val_mask.sum()}')
print(f'Número de nodos de tests: {data.test_mask.sum()}')
print(f'Contiene nodos aislados: {data.contains_isolated_nodes()}')
print(f'Contiene bucles: {data.contains_self_loops()}')
print(f'No es dirigido: {data.is_undirected()}')
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.x
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.tx
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.allx
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.y
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ty
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.ally
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.graph
Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.cora.test.index
Processing...
Done!
Data(edge_index=[2, 10556], test_mask=[2708], train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])
Número de nodos: 2708
Número de features por nodo: 1433
Número de clases: 7
Número de enlaces: 10556
Grado medio de los nodos: 3.90
Número de nodos de entrenamiento: 140
Número de nodos de validación: 500
Número de nodos de tests: 1000
Contiene nodos aislados: False
Contiene bucles: False
No es dirigido: True

#Mostramos la matriz de enlaces en formato COO 
data.edge_index.t()  
tensor([[   0,  633],         [   0, 1862],         [   0, 2582],         ...,         [2707,  598],         [2707, 1473],         [2707, 2706]])
#Mostramos las etiquetas de los primeros 100 nodos 
data.y[0:100] 
tensor([3, 4, 4, 0, 3, 2, 0, 3, 3, 2, 0, 0, 4, 3, 3, 3, 2, 3, 1, 3, 5, 3, 4, 6,         3, 3, 6, 3, 2, 4, 3, 6, 0, 4, 2, 0, 1, 5, 4, 4, 3, 6, 6, 4, 3, 3, 2, 5,         3, 4, 5, 3, 0, 2, 1, 4, 6, 3, 2, 2, 0, 0, 0, 4, 2, 0, 4, 5, 2, 6, 5, 2,         2, 2, 0, 4, 5, 6, 4, 0, 0, 0, 4, 2, 4, 1, 4, 6, 0, 4, 2, 4, 6, 6, 0, 0,         6, 5, 0, 6])     
#Mostramos la máscara que indica que nodos son para entrenamiento viendo #que son los primeros 140 
data.train_mask[0:150] 
tensor([ True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True,  True,  True,  True,
        False, False, False, False, False, False, False, False, False, False])

Ahora, vamos a definir un modelo para realizar una clasificación de los nodos. Para ello, vamos a usar dos capas GCNConv que implementarán la Graph Neural Network. Después de la primera GCN (convierte de la dimensión número de features al número de canales 16) añadimos un ReLU y después de la segunda (convierte de 16 al número de clases) un softmax sobre el número de clases. Como se puede ver, las capas se aplican sobre los datos con los features de cada nodo y sobre edge_index, que contiene la estructura del grafo.

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

model = GCN()
print(model)
GCN(
  (conv1): GCNConv(1433, 16)
  (conv2): GCNConv(16, 7)
)

Ahora, vamos a entrenar el modelo usando 250 epochs (rondas) de los datos. Como se puede observar, usamos la máscara de entrenamiento para decir cuáles son los nodos que se tienen que usar para entrenar el modelo.

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

model.train()
for epoch in range(250):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

Por último, evaluamos el modelo usando la máscara que indica los nodos de test y vemos que el modelo tiene una buena tasa de acierto.

model.eval()
_, pred = model(data).max(dim=1)
correct = int(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / int(data.test_mask.sum())
print('Accuracy: {:.4f}'.format(acc))
Accuracy: 0.8050

Deja un comentario

Tu dirección de correo electrónico no será publicada. Los campos obligatorios están marcados con *

Orgullosamente ofrecido por WordPress | Tema: Baskerville 2 por Anders Noren.

Subir ↑

A %d blogueros les gusta esto: