ContextQMD
Back to Pytorch Geometric

Loading Graphs from CSV

docs/source/tutorial/load_csv.rst

2.7.010.5 KB
Original Source

Loading Graphs from CSV

In this example, we will show how to load a set of :obj:*.csv files as input and construct a heterogeneous graph from it, which can be used as input to a heterogeneous graph model <heterogeneous.html>__. This tutorial is also available as an executable example script <https://github.com/pyg-team/pytorch_geometric/tree/master/examples/hetero/load_csv.py>_ in the :obj:examples/hetero directory.

We are going to use the MovieLens dataset <https://grouplens.org/datasets/movielens/>_ collected by the GroupLens research group. This toy dataset describes 5-star rating and tagging activity from MovieLens. The dataset contains approximately 100k ratings across more than 9k movies from more than 600 users. We are going to use this dataset to generate two node types holding data for movies and users, respectively, and one edge type connecting users and movies, representing the relation of how a user has rated a specific movie.

First, we download the dataset to an arbitrary folder (in this case, the current directory):

.. code-block:: python

from torch_geometric.data import download_url, extract_zip

url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'
extract_zip(download_url(url, '.'), '.')

movie_path = './ml-latest-small/movies.csv'
rating_path = './ml-latest-small/ratings.csv'

Before we create the heterogeneous graph, let's take a look at the data.

.. code-block:: python

import pandas as pd

print(pd.read_csv(movie_path).head())
print(pd.read_csv(rating_path).head())

.. list-table:: Head of :obj:movies.csv :widths: 5 40 60 :header-rows: 1

* - movieId
  - title
  - genres
* - 1
  - Toy Story (1995)
  - Adventure|Animation|Children|Comedy|Fantasy
* - 2
  - Jumanji (1995)
  - Adventure|Children|Fantasy
* - 3
  - Grumpier Old Men (1995)
  - Comedy|Romance
* - 4
  - Waiting to Exhale (1995)
  - Comedy|Drama|Romance
* - 5
  - Father of the Bride Part II (1995)
  - Comedy

We see that the :obj:movies.csv file provides three columns: :obj:movieId assigns a unique identifier to each movie, while the :obj:title and :obj:genres columns represent title and genres of the given movie. We can make use of those two columns to define a feature representation that can be easily interpreted by machine learning models.

.. list-table:: Head of :obj:ratings.csv :widths: 5 5 10 30 :header-rows: 1

* - userId
  - movieId
  - rating
  - timestamp
* - 1
  - 1
  - 4.0
  - 964982703
* - 1
  - 3
  - 4.0
  - 964981247
* - 1
  - 6
  - 4.0
  - 964982224
* - 1
  - 47
  - 5.0
  - 964983815
* - 1
  - 50
  - 5.0
  - 964982931

The :obj:ratings.csv data connects users (as given by :obj:userId) and movies (as given by :obj:movieId), and defines how a given user has rated a specific movie (:obj:rating). Due to simplicity, we do not make use of the additional :obj:timestamp information.

For representing this data in the :pyg:PyG data format, we first define a method :meth:load_node_csv that reads in a :obj:*.csv file and returns a node-level feature representation :obj:x of shape :obj:[num_nodes, num_features]:

.. code-block:: python

import torch

def load_node_csv(path, index_col, encoders=None, **kwargs):
    df = pd.read_csv(path, index_col=index_col, **kwargs)
    mapping = {index: i for i, index in enumerate(df.index.unique())}

    x = None
    if encoders is not None:
        xs = [encoder(df[col]) for col, encoder in encoders.items()]
        x = torch.cat(xs, dim=-1)

    return x, mapping

Here, :meth:load_node_csv reads the :obj:*.csv file from :obj:path, and creates a dictionary :obj:mapping that maps its index column to a consecutive value in the range :obj:{ 0, ..., num_rows - 1 }. This is needed as we want our final data representation to be as compact as possible, e.g., the representation of a movie in the first row should be accessible via :obj:x[0].

We further utilize the concept of encoders, which define how the values of specific columns should be encoded into a numerical feature representation. For example, we can define a sentence encoder that encodes raw column strings into low-dimensional embeddings. For this, we make use of the excellent sentence-transformers <https://www.sbert.net/>_ library which provides a large number of state-of-the-art pretrained NLP embedding models:

.. code-block:: bash

pip install sentence-transformers

.. code-block:: python

class SequenceEncoder:
    def __init__(self, model_name='all-MiniLM-L6-v2', device=None):
        self.device = device
        self.model = SentenceTransformer(model_name, device=device)

    @torch.no_grad()
    def __call__(self, df):
        x = self.model.encode(df.values, show_progress_bar=True,
                              convert_to_tensor=True, device=self.device)
        return x.cpu()

The :class:SequenceEncoder class loads a pre-trained NLP model as given by :obj:model_name, and uses it to encode a list of strings into a :pytorch:PyTorch tensor of shape :obj:[num_strings, embedding_dim]. We can use this :class:SequenceEncoder to encode the :obj:title of the :obj:movies.csv file.

In a similar fashion, we can create another encoder that converts the genres of movies, e.g., :obj:Adventure|Children|Fantasy, into categorical labels. For this, we first need to find all existing genres present in the data, create a feature representation :obj:x of shape :obj:[num_movies, num_genres], and assign a :obj:1 to :obj:x[i, j] in case the genre :obj:j is present in movie :obj:i:

.. code-block:: python

class GenresEncoder:
    def __init__(self, sep='|'):
        self.sep = sep

    def __call__(self, df):
        genres = set(g for col in df.values for g in col.split(self.sep))
        mapping = {genre: i for i, genre in enumerate(genres)}

        x = torch.zeros(len(df), len(mapping))
        for i, col in enumerate(df.values):
            for genre in col.split(self.sep):
                x[i, mapping[genre]] = 1
        return x

With this, we can obtain our final representation of movies via:

.. code-block:: python

movie_x, movie_mapping = load_node_csv(
    movie_path, index_col='movieId', encoders={
        'title': SequenceEncoder(),
        'genres': GenresEncoder()
    })

Similarly, we can utilize :meth:load_node_csv for obtaining a user mapping from :obj:userId to consecutive values as well. However, there is no additional feature information for users present in this dataset. As such, we do not define any encoders:

.. code-block:: python

_, user_mapping = load_node_csv(rating_path, index_col='userId')

With this, we are ready to initialize our :class:~torch_geometric.data.HeteroData object and pass two node types into it:

.. code-block:: python

from torch_geometric.data import HeteroData

data = HeteroData()

data['user'].num_nodes = len(user_mapping)  # Users do not have any features.
data['movie'].x = movie_x

print(data)
HeteroData(
  user={ num_nodes=610 },
  movie={ x[9742, 404] }
)

As users do not have any node-level information, we solely define its number of nodes. As a result, we likely need to learn distinct user embeddings via :class:torch.nn.Embedding in an end-to-end fashion during training of a heterogeneous graph model.

Next, we take a look at connecting users with movies as defined by their ratings. For this, we define a method :meth:load_edge_csv that returns the final :obj:edge_index representation of shape :obj:[2, num_ratings] from :obj:ratings.csv, as well as any additional features present in the raw :obj:*.csv file:

.. code-block:: python

def load_edge_csv(path, src_index_col, src_mapping, dst_index_col, dst_mapping,
                  encoders=None, **kwargs):
    df = pd.read_csv(path, **kwargs)

    src = [src_mapping[index] for index in df[src_index_col]]
    dst = [dst_mapping[index] for index in df[dst_index_col]]
    edge_index = torch.tensor([src, dst])

    edge_attr = None
    if encoders is not None:
        edge_attrs = [encoder(df[col]) for col, encoder in encoders.items()]
        edge_attr = torch.cat(edge_attrs, dim=-1)

    return edge_index, edge_attr

Here, :obj:src_index_col and :obj:dst_index_col define the index columns of source and destination nodes, respectively. We further make use of the node-level mappings :obj:src_mapping and :obj:dst_mapping to ensure that raw indices are mapped to the correct consecutive indices in our final representation. For every edge defined in the file, it looks up the forward indices in :obj:src_mapping and :obj:dst_mapping, and moves the data appropriately.

Similarly to :meth:load_node_csv, encoders are used to return additional edge-level feature information. For example, for loading the ratings from the :obj:rating column in :obj:ratings.csv, we can define an :class:IdentityEncoder that simply converts a list of floating-point values into a :pytorch:PyTorch tensor:

.. code-block:: python

class IdentityEncoder:
    def __init__(self, dtype=None):
        self.dtype = dtype

    def __call__(self, df):
        return torch.from_numpy(df.values).view(-1, 1).to(self.dtype)

With this, we are ready to finalize our :class:~torch_geometric.data.HeteroData object:

.. code-block:: python

edge_index, edge_label = load_edge_csv(
    rating_path,
    src_index_col='userId',
    src_mapping=user_mapping,
    dst_index_col='movieId',
    dst_mapping=movie_mapping,
    encoders={'rating': IdentityEncoder(dtype=torch.long)},
)

data['user', 'rates', 'movie'].edge_index = edge_index
data['user', 'rates', 'movie'].edge_label = edge_label

print(data)
HeteroData(
  user={ num_nodes=610 },
  movie={ x=[9742, 404] },
  (user, rates, movie)={
    edge_index=[2, 100836],
    edge_label=[100836, 1]
  }
)

This :class:~torch_geometric.data.HeteroData object is the native format of heterogeneous graphs in :pyg:PyG and can be used as input for heterogeneous graph models <heterogeneous.html>__.

.. note::

Click `here <https://github.com/pyg-team/pytorch_geometric/tree/master/examples/hetero/load_csv.py>`_ to see the final example script.