Proposal for batch.filter(mask) method to enhance efficiency

[leonardcaquot94]

Hello PyG team!

I’d like to propose the addition of a batch.filter(mask) method to PyTorch Geometric, which would take a mask or a list of indices and efficiently filter the batch without using the new_batch = Batch.from_data_list(batch[mask]) conversion. Below is a sample comparison between two approaches, demonstrating significant performance improvement.

Code Example:

from datetime import datetime
from torch_geometric.data import HeteroData, Batch
import torch

# Create initial HeteroData object and batch
data = HeteroData()
data['paper'].x = torch.randn(5, 16)
batch_size = 50
batch = Batch.from_data_list([data]*batch_size)

def func1(batch):
    idx_to_keep = torch.rand(batch_size) < 0.5
    data_to_keep = batch[idx_to_keep]
    new_batch = Batch.from_data_list(data_to_keep)
    return new_batch

def func2(batch):
    idx_to_keep = torch.rand(batch_size) < 0.5
    nodes_to_keep = idx_to_keep[batch['paper'].batch]
    batch['paper'].x = batch['paper'].x[nodes_to_keep]
    batch_filtered = batch['paper'].batch[nodes_to_keep]
    _, batch['paper'].batch, counts = torch.unique_consecutive(batch_filtered, return_inverse=True, return_counts=True)
    batch['paper'].ptr = torch.cat((torch.zeros(1, dtype=torch.int), counts.cumsum(0)))
    batch._num_graphs = idx_to_keep.sum()
    return batch

def copy(batch):
    batch_1 = Batch.from_data_list(batch.to_data_list())
    batch_2 = Batch.from_data_list(batch.to_data_list())
    return batch_1, batch_2

t1 = 0
t2 = 0
for _ in range(1000):
    batch_1, batch_2 = copy(batch)

    t = datetime.now()
    x = func1(batch_1)
    t1 += (datetime.now() - t).microseconds

    t = datetime.now()
    y = func2(batch_2)
    t2 += (datetime.now() - t).microseconds

print(f'func1 : {t1*1e-6} seconds')
print(f'func2 : {t2*1e-6} seconds')

Execution Results:

func1 : 0.946913 seconds
func2 : 0.11657 seconds

Explanation:

• func1 uses the standard Batch.from_data_list() method to filter the batch, which is slower.
• func2 directly manipulates the batch object without converting to/from lists, resulting in approximately a 10x speedup.

Additional Suggestions:

To make this feature more efficient, it could be beneficial to reuse some of the existing code from the collate function to handle custom node and edge attributes iteratively. Additionally, it might be useful to provide an option to return a "negative sub-batch" (i.e., the elements that are excluded from the mask) alongside the positive sub-batch.

Usage :

Currently, when we filter batch data like sub_batch = batch[mask], it returns a data list. But I believe it would be more convenient if it kept the result as a Batch object. This way, users can maintain the Batch format and call sub_batch.to_data_list() only when they explicitly need a data list. This would streamline operations where batch structure needs to be preserved.

I believe these improvements could greatly enhance performance, especially in batch filtering scenarios.

Let me know what you think of this idea, and if any further clarifications are needed!

[erikhuck]

We absolutely need a more efficient Batch splicing method. +1. You might be able to propose a pull request to get the ball rolling on this.

[erikhuck]

@leonardcaquot94 I've created a partial solution here. If you or anyone else has feedback, please let me know. We need to update the edge_index and we need an option to update the edge_attr as well.

import torch
import torch_geometric.nn as tgnn
import torch_geometric.data as tdata
import torch_geometric.transforms as ttrans
import torch_geometric.nn.aggr as aggr
import tqdm

max_n_nodes = 60
node_dim = 50
n_graphs = 7000
bond_dim = 5

def create_graph(n_nodes, node_dim, bond_dim):
    x = torch.rand(n_nodes, node_dim).cuda()
    edge_index = torch.randint(0, n_nodes, (2, n_nodes * 2)).cuda()
    edge_attr = torch.rand(edge_index.shape[1], bond_dim).cuda()
    graph = tdata.Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
    graph = ttrans.ToUndirected()(graph)
    graph.validate(raise_on_error=True)
    return graph

graphs = [
    create_graph(n_nodes=torch.randint(2, max_n_nodes, (1,)).item(), node_dim=node_dim, bond_dim=bond_dim) for _ in range(n_graphs)]
edge_counts = [graph.edge_index.shape[1] for graph in graphs]
graphs = tdata.Batch.from_data_list(graphs)

edge_ptr = [0]
for count in edge_counts:
    edge_ptr.append(edge_ptr[-1] + count)
edge_ptr = torch.tensor(edge_ptr)
graphs.edge_ptr = edge_ptr.cuda()

graph_indices = torch.tensor([999, 500, 760, 499, 299, 234, 0, 1500, 235, 1700, 354, 6999, 5, 3, 10] * 250)

def _get_splice_indices(graph_indices, ptr):
    start_idx = ptr[graph_indices].unsqueeze(1)
    end_idx = ptr[graph_indices + 1].unsqueeze(1)
    lengths = end_idx - start_idx
    indices = start_idx + torch.arange(lengths.max().item(), device='cuda').expand(len(lengths), -1)
    mask = (indices < end_idx)
    indices = indices[mask].flatten()
    start_indices = (lengths.cumsum(dim=0) - lengths).flatten()
    new_ptr = torch.cat((start_indices.masked_select(lengths.flatten() > 0), torch.tensor([len(indices)], device='cuda')))
    return indices, new_ptr, start_idx, lengths

def splice_batch(batch, graph_indices):
    edge_indices, edge_ptr, _, edge_lengths = _get_splice_indices(graph_indices, batch.edge_ptr)
    node_indices, ptr, nodes_start_idx, _ = _get_splice_indices(graph_indices, batch.ptr)
    edge_offsets = nodes_start_idx.squeeze().repeat_interleave(edge_lengths.squeeze()).reshape(1, -1)
    new_edge_index = batch.edge_index[:, edge_indices] - edge_offsets
    edge_offsets = ptr[:-1].repeat_interleave(edge_lengths.squeeze()).reshape(1, -1)
    new_edge_index = new_edge_index + edge_offsets
    return tdata.Batch(
        x=batch.x[node_indices],
        edge_index=new_edge_index,
        edge_attr=batch.edge_attr[edge_indices],
        ptr=ptr, edge_ptr=edge_ptr
    )

def old_method(batch, graph_indices):
    return tdata.Batch.from_data_list(batch[graph_indices])

old_method_batch = old_method(graphs, graph_indices)
new_method_batch = splice_batch(graphs, graph_indices)

assert torch.equal(new_method_batch.x, old_method_batch.x)
assert torch.equal(new_method_batch.edge_attr, old_method_batch.edge_attr)
assert torch.equal(new_method_batch.edge_index, old_method_batch.edge_index)
assert torch.equal(new_method_batch.ptr, old_method_batch.ptr)

[leonardcaquot94]

@erikhuck thank you very much. I tried your code on HeteroData and I think we'll have to deal with more complex situations where data have more complex structures.
Perhaps we could try applying it to a batch of graphs like the one below:

batch_size = 50
data_list = []
for i in range(batch_size):
    num_nodes = (i % 5) + 3
    num_edges = num_nodes + 2 * (i % 3)
    data = HeteroData()
    data.custom_attr_1= i
    data.custom_attr_1 = [i]*3
    data['paper'].x = torch.tensor([i]).repeat((num_nodes, 4))
    data['paper'].custom_attr = torch.tensor([i]).repeat((num_nodes*2, 1))
    data['paper', 'to', 'paper'].edge_index = torch.randint(num_nodes, (2, num_edges))
    data_list.append(data)

I believe we’ll need to iterate over different data stores to manage each attribute, similar to the approach in the collate function.

[leonardcaquot94]

Here is a PR linked to this discussion : #9911