mini-batch VGAE with sparse tensor

[A11en0]

hi, thanks for your wonderful works. could you please give a demo that implements the mini-batch version VGAE with sparse tensor?

Or helps to re-implement the GraphBTM, it uses the VGAE to make a topic model, which has a raw code at https://github.com/valdersoul/GraphBTM.

Edit: Has to mention is that it takes the adjacent matrix $A$ as the node feature $X$ in the model. So we need both of $A$ and $X$ should be sparse tensors.

[rusty1s]

Do you mean a mini-batch VGAE version for multiple graphs or for a single large graph? I think the mini-batch version of VGAE for multiple graphs is pretty straightforward to add, and only involves utilizing torch_geometric.utils.to_dense_batch in the decoder part (no changes are necessary for the encoder):

def encode(self, z, batch):
    z = to_dense_batch(z, batch)
    return z @ z.transpose(1, 2)

[A11en0]

Thanks for your quick reply!

yes, I mean multiple graph mini-batch. I take your instruction to implement the decoder, and then I make a mini-batch version of GAE(I first try an easier one then convert it to VGAE), I used the UPFD dataset which is a multi-graph dataset to do this experiment.

To Summarize, I encounter these main problems:

• How to calculate the reconstruction loss in the batch style? (I don't sure whether my implementation is right or not.)
• The decoder output of my own datasets is a dense matrix, about the size of [64 * 10000 * 10000], it's hard to calculate the reconstruction loss, can I use a sparse matrix instead of it?
• the loss cannot be done

Following is my not finished works, it's really tough, hope you can help me to check it, thanks!

[A11en0]

import argparse
import os.path as osp
import torch
import torch.nn.functional as F
from torch_geometric.datasets import UPFD
from torch_geometric.data import DataLoader
from torch_geometric.transforms import ToUndirected
from torch_geometric.nn import GCNConv
from torch_geometric.utils import negative_sampling, to_dense_batch, to_dense_adj

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='politifact',
                    choices=['politifact', 'gossipcop'])
parser.add_argument('--feature', type=str, default='profile',
                    choices=['profile', 'spacy', 'bert', 'content'])
args = parser.parse_args()

path = osp.join(osp.dirname(osp.realpath(__file__)), '.', 'datasets', 'UPFD')
train_dataset = UPFD(path, args.dataset, args.feature, 'train', ToUndirected())
val_dataset = UPFD(path, args.dataset, args.feature, 'val', ToUndirected())
test_dataset = UPFD(path, args.dataset, args.feature, 'test', ToUndirected())

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)

class GraphBTM(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(GraphBTM, self).__init__()
        self.conv1 = GCNConv(in_channels, 128)
        self.conv2 = GCNConv(128, out_channels)

    def encode(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = x.relu()
        return self.conv2(x, edge_index)

    def decode(self, z, batch):
        z, mask = to_dense_batch(z, batch)
        z = z @ z.transpose(1, 2)
        z = z.view(z.size(0), -1)
        return z

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = GraphBTM(train_dataset.num_features, 64).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=0.01)

def train():
    model.train()
    loss = 0.0
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()

        z = model.encode(data.x, data.edge_index)
        link_preds = model.decode(z, data.batch)
        link_preds = link_preds.sigmoid()

        neg_edge_index = negative_sampling(data.edge_index, z.size(0))
        link_labels = to_dense_adj(neg_edge_index, data.batch)
        link_labels = link_labels.view(link_labels.size(0), -1)

        loss = F.binary_cross_entropy_with_logits(link_preds, link_labels)

        loss.backward()
        optimizer.step()

    return loss

for epoch in range(1, 101):
    loss = train()
    print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}')

[rusty1s]

I see. Converting via to_dense_batch makes only sense if your adjacency matrices are reasonable small which is not the case. As an alternative, you can decode in a sparse fashion, e.g. via:

def decode(self, z, edge_index):
    src, dst = edge_index
    return (z[src] * z[dst]).sum(dim=-1)

and apply your training loop as follows:

for data in train_loader:
    ...
    z = model.encode(data.x, data.edge_index)
    neg_edge_index = batched_negative_sampling(data.edge_index, data.batch, ...)

    out_pos = model.decode(z, data.edge_index)
    out_neg = model.decode(z, neg_edge_index)
    
    pos_loss = F.binary_cross_entropy_with_logits(out_pos, torch.ones_like(out_pos))
    neg_loss = F.binary_cross_entropy_with_logits(out_neg, torch.zeros_like(out_neg))
    loss = pos_loss + neg_loss

[rusty1s]

Happy to take a PR on this :)