Repository: pyg-team/pytorch_geometric

Issue:
An example of heterogeneous link prediction via `RandomLinkSplit`



Discussion:
Hello! I'm thinking of picking this one up, if that's helpful. Did you have a specific dataset to use in mind already?
Cool :) I thought about extending the "Loading CSV" tutorial to showcase how one would apply a GNN on this one. I already started integrating the random link split behaviour, see [here](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/load_csv.py#L109-L127). The next task would be to create a heterogeneous GNN model, and train it in a supervised fashion against ratings in the training set. WDYT?
Ah yes, that is indeed a nice start already! :)

This example feels a bit different from "typical" link prediction statements, in that I don't think you can really have a contrastive loss with negative edges, as a missing edge in this graph just means we want to predict what rating there should be for each edge of type `('user', 'rates', 'movie')`. So we don't want to train the algorithm to give good separation between "likely" and "unlikely" edges. I think that's fine though, I see this as an edge classification problem and it seems a relevant example, reminiscent of the problem statement of predicting how users would rate products on online stores. 

Can you check that my plan for this fits your idea about what you would like?
- In the encoder step I perform heterogeneous message passing, maybe using a `GNN` with 2 layers
- In the decoder step I use a `DistMult` decoder to get scores for each edge label (from 0 to 5), so here I am - kind of - treating the edge labels as six different edge types. If `DistMult` doesn't do the trick, I could try a bilinear, `RESCAL` type decoder.
- Then I apply a `softmax` on the 6 scores I got from the decoder per training supervision edge to get something that "looks like" class probabilities, pick the class with the highest probability and use a loss suitable for a multiclass problem statement (such as `torch.nn.NLLLoss`)
- Then backpropagate etc...

I'm happy to try this approach, I just wanted to check if you already had some kind of plan that is quite different from mine, so we don't waste too much time. Thanks! :)
You are right, it's more of an edge classification problem in which no negative sampling is needed. Nonetheless, the model should be able to predict the ratings of unknown users/movies. Your approach sounds correct, and matches with the one I have in mind. Let me know how it goes :)
Alright, just a quick status update: I've put something together and it is learning but the performance is not amazing so I want to improve it a bit. The average test accuracy gets to about 40% after 400 epochs, which I guess is better than random for a 6 class problem but there are a few things I want to try to make it better before sharing it.

I'm afraid I only have time to do this in my evenings so progress is perhaps a bit slow. Hope that's ok.
Sure, please feel free to submit a PR early, so I can help with it :)

Resolving patch:
diff --git a/examples/hetero/hetero_link_pred.py b/examples/hetero/hetero_link_pred.py
new file mode 100644
--- /dev/null
+++ b/examples/hetero/hetero_link_pred.py
@@ -0,0 +1,132 @@
+import os.path as osp
+import argparse
+
+import torch
+from torch.nn import Linear
+import torch.nn.functional as F
+
+import torch_geometric.transforms as T
+from torch_geometric.datasets import MovieLens
+from torch_geometric.nn import SAGEConv, to_hetero
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--use_weighted_loss', action='store_true',
+                    help='Whether to use weighted MSE loss.')
+args = parser.parse_args()
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/MovieLens')
+dataset = MovieLens(path, model_name='all-MiniLM-L6-v2')
+data = dataset[0].to(device)
+
+# Add user node features for message passing:
+data['user'].x = torch.eye(data['user'].num_nodes, device=device)
+del data['user'].num_nodes
+
+# Add a reverse ('movie', 'rev_rates', 'user') relation for message passing:
+data = T.ToUndirected()(data)
+del data['movie', 'rev_rates', 'user'].edge_label  # Remove "reverse" label.
+
+# Perform a link-level split into training, validation, and test edges:
+train_data, val_data, test_data = T.RandomLinkSplit(
+    num_val=0.1,
+    num_test=0.1,
+    neg_sampling_ratio=0.0,
+    edge_types=[('user', 'rates', 'movie')],
+    rev_edge_types=[('movie', 'rev_rates', 'user')],
+)(data)
+
+# We have an unbalanced dataset with many labels for rating 3 and 4, and very
+# few for 0 and 1. Therefore we use a weighted MSE loss.
+if args.use_weighted_loss:
+    weight = torch.bincount(train_data['user', 'movie'].edge_label)
+    weight = weight.max() / weight
+else:
+    weight = None
+
+
+def weighted_mse_loss(pred, target, weight=None):
+    weight = 1. if weight is None else weight[target].to(pred.dtype)
+    return (weight * (pred - target.to(pred.dtype)).pow(2)).mean()
+
+
+class GNNEncoder(torch.nn.Module):
+    def __init__(self, hidden_channels, out_channels):
+        super().__init__()
+        self.conv1 = SAGEConv((-1, -1), hidden_channels)
+        self.conv2 = SAGEConv((-1, -1), out_channels)
+
+    def forward(self, x, edge_index):
+        x = self.conv1(x, edge_index).relu()
+        x = self.conv2(x, edge_index)
+        return x
+
+
+class EdgeDecoder(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
+        self.lin2 = Linear(hidden_channels, 1)
+
+    def forward(self, z_dict, edge_label_index):
+        row, col = edge_label_index
+        z = torch.cat([z_dict['user'][row], z_dict['movie'][col]], dim=-1)
+
+        z = self.lin1(z).relu()
+        z = self.lin2(z)
+        return z.view(-1)
+
+
+class Model(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
+        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
+        self.decoder = EdgeDecoder(hidden_channels)
+
+    def forward(self, x_dict, edge_index_dict, edge_label_index):
+        z_dict = self.encoder(x_dict, edge_index_dict)
+        return self.decoder(z_dict, edge_label_index)
+
+
+model = Model(hidden_channels=32).to(device)
+
+# Due to lazy initialization, we need to run one model step so the number
+# of parameters can be inferred:
+with torch.no_grad():
+    model.encoder(train_data.x_dict, train_data.edge_index_dict)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+
+def train():
+    model.train()
+    optimizer.zero_grad()
+    pred = model(train_data.x_dict, train_data.edge_index_dict,
+                 train_data['user', 'movie'].edge_label_index)
+    target = train_data['user', 'movie'].edge_label
+    loss = weighted_mse_loss(pred, target, weight)
+    loss.backward()
+    optimizer.step()
+    return float(loss)
+
+
+@torch.no_grad()
+def test(data):
+    model.eval()
+    pred = model(data.x_dict, data.edge_index_dict,
+                 data['user', 'movie'].edge_label_index)
+    pred = pred.clamp(min=0, max=5)
+    target = data['user', 'movie'].edge_label.float()
+    rmse = F.mse_loss(pred, target).sqrt()
+    return float(rmse)
+
+
+for epoch in range(1, 301):
+    loss = train()
+    train_rmse = test(train_data)
+    val_rmse = test(val_data)
+    test_rmse = test(test_data)
+    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_rmse:.4f}, '
+          f'Val: {val_rmse:.4f}, Test: {test_rmse:.4f}')
diff --git a/torch_geometric/datasets/__init__.py b/torch_geometric/datasets/__init__.py
--- a/torch_geometric/datasets/__init__.py
+++ b/torch_geometric/datasets/__init__.py
@@ -50,6 +50,7 @@
 from .actor import Actor
 from .ogb_mag import OGB_MAG
 from .dblp import DBLP
+from .movie_lens import MovieLens
 from .imdb import IMDB
 from .last_fm import LastFM
 from .hgb_dataset import HGBDataset
@@ -126,6 +127,7 @@
     'Actor',
     'OGB_MAG',
     'DBLP',
+    'MovieLens',
     'IMDB',
     'LastFM',
     'HGBDataset',
diff --git a/torch_geometric/datasets/movie_lens.py b/torch_geometric/datasets/movie_lens.py
new file mode 100644
--- /dev/null
+++ b/torch_geometric/datasets/movie_lens.py
@@ -0,0 +1,93 @@
+from typing import Optional, Callable, List
+
+import os
+import os.path as osp
+
+import torch
+
+from torch_geometric.data import (InMemoryDataset, HeteroData, download_url,
+                                  extract_zip)
+
+
+class MovieLens(InMemoryDataset):
+    r"""A heterogeneous rating dataset, assembled by GroupLens Research from
+    the `MovieLens web site <https://movielens.org>`_, consisting of nodes of
+    type :obj:`"movie"` and :obj:`"user"`.
+    User ratings for movies are available as ground truth labels for the edges
+    between the users and the movies :obj:`("user", "rates", "movie")`.
+
+    Args:
+        root (string): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            every access. (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        model_name (str): Name of model used to transform movie titles to node
+            features. The model comes from the`Huggingface SentenceTransformer
+            <https://huggingface.co/sentence-transformers>`_.
+    """
+
+    url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'
+
+    def __init__(self, root, transform: Optional[Callable] = None,
+                 pre_transform: Optional[Callable] = None,
+                 model_name: Optional[str] = "all-MiniLM-L6-v2"):
+        self.model_name = model_name
+        super().__init__(root, transform, pre_transform)
+        self.data, self.slices = torch.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return [
+            osp.join('ml-latest-small', 'movies.csv'),
+            osp.join('ml-latest-small', 'ratings.csv'),
+        ]
+
+    @property
+    def processed_file_names(self) -> str:
+        return f'data_{self.model_name}.pt'
+
+    def download(self):
+        path = download_url(self.url, self.raw_dir)
+        extract_zip(path, self.raw_dir)
+        os.remove(path)
+
+    def process(self):
+        import pandas as pd
+        from sentence_transformers import SentenceTransformer
+
+        data = HeteroData()
+
+        df = pd.read_csv(self.raw_paths[0], index_col='movieId')
+        movie_mapping = {idx: i for i, idx in enumerate(df.index)}
+
+        genres = df['genres'].str.get_dummies('|').values
+        genres = torch.from_numpy(genres).to(torch.float)
+
+        model = SentenceTransformer(self.model_name)
+        with torch.no_grad():
+            emb = model.encode(df['title'].values, show_progress_bar=True,
+                               convert_to_tensor=True).cpu()
+
+        data['movie'].x = torch.cat([emb, genres], dim=-1)
+
+        df = pd.read_csv(self.raw_paths[1])
+        user_mapping = {idx: i for i, idx in enumerate(df['userId'].unique())}
+        data['user'].num_nodes = len(user_mapping)
+
+        src = [user_mapping[idx] for idx in df['userId']]
+        dst = [movie_mapping[idx] for idx in df['movieId']]
+        edge_index = torch.tensor([src, dst])
+
+        rating = torch.from_numpy(df['rating'].values).to(torch.long)
+        data['user', 'rates', 'movie'].edge_index = edge_index
+        data['user', 'rates', 'movie'].edge_label = rating
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        torch.save(self.collate([data]), self.processed_paths[0])


Test patch: