TGraphX Feature Store: Scaling Tensor Features Beyond RAM

Training a GNN on a graph where each node carries a vector of 128 floats is manageable even on a laptop. Training on a graph where each node carries a 3D volume of shape [32, 64, 64] — roughly 500,000 floats per node — is a different problem entirely. A graph with 100,000 such nodes requires about 200 GB of RAM to hold all node features in memory at once. Most machines do not have that, and even those that do would spend significant time waiting for data to transfer to the GPU.

The TGraphX feature store addresses this problem by keeping features on disk and loading only the subset needed for the current mini-batch. This article explains why tensor features exceed RAM budgets quickly, how memory-mapped arrays work, and how to use tgraphx.feature_store to fetch features for subgraph mini-batches efficiently.

What This Builds On

This article assumes you are familiar with mini-batch GNN training and the concept of neighbor sampling — loading a small subgraph per step rather than the full graph. The neighbor sampling article covers the subgraph loading pattern that the feature store integrates with. Basic familiarity with numpy's memmap is helpful but not required.

Why Tensor Features Exhaust RAM

For standard vector-feature GNNs, memory is dominated by model parameters and activations, not the node features themselves. A graph with 1 million nodes, each with a 64-dimensional feature vector, requires only about 240 MB in float32 — trivially manageable.

The situation changes dramatically with spatial or volumetric node features:

• Medical imaging graph: 50,000 patient nodes, each with a [1, 96, 96, 64] MRI patch → ~3.5 TB
• Remote sensing graph: 1 million tile nodes, each with a [4, 64, 64] spectral image → ~16 GB
• Point cloud graph: 500,000 region nodes, each with a [3, 32, 32, 32] occupancy volume → ~18 GB

These sizes exceed typical GPU VRAM by one to three orders of magnitude, and in many cases exceed system RAM as well. The only practical approach is to store features on fast SSD storage and load them on demand.

Memory-Mapped Arrays: The Core Mechanism

NumPy's memmap creates an array-like object backed by a file on disk. When you access a slice of the array, the operating system's virtual memory system reads only the relevant pages from disk — there is no up-front copy of the entire file into RAM. The OS also caches recently accessed pages in the page cache, so repeated access to the same region is fast.

import numpy as np
        
        # Create a memory-mapped feature array
        # Shape: [100000, 32, 64, 64] in float32
        mmap = np.memmap(
            '/data/node_features.bin',
            dtype='float32',
            mode='r',      # read-only
            shape=(100000, 32, 64, 64)
        )
        
        # Load features for a mini-batch of node indices
        node_ids = np.array([0, 5, 99, 1234, 5678])
        batch_features = mmap[node_ids]   # shape [5, 32, 64, 64]
        # Only these 5 entries are read from disk (and cached by the OS)
        

The page cache means that if your mini-batch sampling is spatially local — neighboring nodes in the graph tend to have nearby indices in the feature array — repeated access patterns will be served from RAM. If node indices are random across the whole range, each access will trigger a disk read. Index ordering can therefore dramatically affect throughput.

Creating a Feature Store with TGraphX

TGraphX's tgraphx.feature_store module wraps the memmap pattern in a convenient interface:

import torch
        import numpy as np
        from tgraphx.feature_store import FeatureStore
        
        # Create and write a feature store (done once, typically during dataset preparation)
        feature_shape = (100_000, 8, 16, 16)  # [N, C, H, W]
        dtype = np.float32
        
        store = FeatureStore.create(
            path='/data/features.bin',
            shape=feature_shape,
            dtype=dtype,
        )
        
        # Populate the store (streaming from your source data)
        # Typically done in chunks to avoid loading everything at once
        chunk_size = 1000
        for start in range(0, feature_shape[0], chunk_size):
            end = min(start + chunk_size, feature_shape[0])
            # Replace with your actual data loading logic
            chunk = np.random.randn(end - start, 8, 16, 16).astype(np.float32)
            store.write(start, chunk)
        
        store.close()
        

Loading the store for training:

store = FeatureStore.load(path='/data/features.bin', shape=feature_shape, dtype=dtype)
        
        # Fetch features for a batch of node IDs
        node_ids = torch.tensor([0, 15, 42, 100, 999], dtype=torch.long)
        features = store.fetch(node_ids)   # returns a torch.Tensor on CPU
        print(features.shape)  # [5, 8, 16, 16]
        

The fetch call reads only the requested rows from disk. The returned tensor is a CPU tensor. Transfer to GPU happens separately, which allows overlapping data loading with GPU computation using PyTorch's DataLoader prefetching.

The fetch_features_for_subgraph Pattern

When training with mini-batch neighbor sampling, the typical pattern is:

1. Sample a subgraph (a set of node IDs and their induced edges)
2. Fetch features for those node IDs from the feature store
3. Run the GNN forward pass on the subgraph
4. Backpropagate

import torch
        import torch.nn.functional as F
        from tgraphx.feature_store import FeatureStore
        from tgraphx.layers.sage import TensorGraphSAGELayer
        import torch.nn as nn
        
        # Assume 'store' is an open FeatureStore
        # Assume 'full_edge_index' is the full graph edge index
        
        def sample_subgraph(edge_index, seed_nodes, num_neighbors=10):
            """Minimal neighbor sampler — replace with tgraphx.graphsaint or cluster_gcn."""
            all_nodes = set(seed_nodes.tolist())
            row, col = edge_index
            # Find neighbors of seed nodes
            for sn in seed_nodes.tolist():
                neighbors_mask = col == sn
                neighbors = row[neighbors_mask]
                # Sample up to num_neighbors
                if neighbors.numel() > num_neighbors:
                    idx = torch.randperm(neighbors.numel())[:num_neighbors]
                    neighbors = neighbors[idx]
                all_nodes.update(neighbors.tolist())
            all_nodes = torch.tensor(sorted(all_nodes), dtype=torch.long)
            return all_nodes
        
        class TensorSAGEModel(nn.Module):
            def __init__(self, in_channels, hidden_channels, num_classes):
                super().__init__()
                self.sage1 = TensorGraphSAGELayer(in_channels, hidden_channels)
                self.sage2 = TensorGraphSAGELayer(hidden_channels, hidden_channels)
                # After two spatial layers, features are [N, hidden, H, W]
                # Flatten for classifier
                self.classifier = nn.Linear(hidden_channels * 16 * 16, num_classes)
        
            def forward(self, x, edge_index):
                x = F.relu(self.sage1(x, edge_index))
                x = self.sage2(x, edge_index)
                # Global average pool over spatial dims
                x = x.mean(dim=(-2, -1))   # [N, hidden_channels]
                return self.classifier(x)
        
        # Training loop sketch
        feature_shape = (100_000, 8, 16, 16)
        store = FeatureStore.load('/data/features.bin', feature_shape, np.float32)
        model = TensorSAGEModel(in_channels=8, hidden_channels=32, num_classes=5)
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
        
        labels = torch.randint(0, 5, (100_000,))   # example labels
        
        for step in range(200):
            seed_nodes = torch.randint(0, 100_000, (64,))
            subgraph_nodes = sample_subgraph(full_edge_index, seed_nodes)
        
            # Fetch only the features we need
            x = store.fetch(subgraph_nodes)   # [subgraph_size, 8, 16, 16]
            x = x.cuda()
        
            # Build local edge index for the subgraph
            # (remapping global node IDs to local 0-indexed IDs)
            node_id_map = {nid.item(): i for i, nid in enumerate(subgraph_nodes)}
            local_edges = []
            r, c = full_edge_index
            for i in range(r.size(0)):
                ri, ci = r[i].item(), c[i].item()
                if ri in node_id_map and ci in node_id_map:
                    local_edges.append([node_id_map[ri], node_id_map[ci]])
            if local_edges:
                local_edge_index = torch.tensor(local_edges, dtype=torch.long).T.cuda()
            else:
                local_edge_index = torch.zeros(2, 0, dtype=torch.long).cuda()
        
            seed_local_ids = torch.tensor([node_id_map[s.item()] for s in seed_nodes
                                           if s.item() in node_id_map], dtype=torch.long)
            logits = model(x, local_edge_index)
            seed_labels = labels[subgraph_nodes[seed_local_ids]].cuda()
            loss = F.cross_entropy(logits[seed_local_ids], seed_labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        

Performance Considerations

Several factors determine whether the feature store provides adequate throughput:

Storage speed: NVMe SSDs provide sequential read speeds of 3–7 GB/s. SATA SSDs provide 500–600 MB/s. A training step that requires reading 64 subgraph nodes × 8 × 16 × 16 × 4 bytes ≈ 2 MB will take under 1 ms on an NVMe drive. For larger features or more nodes, the math changes quickly.

Random vs sequential access: memmap access to non-contiguous node IDs involves random I/O, which is slower than sequential I/O even on SSDs. Sorting node IDs before fetching can improve throughput slightly. Reordering graph nodes so that structurally close nodes have nearby IDs (using BFS or Cuthill-McKee ordering) can dramatically improve cache hit rates.

Prefetching: Load the next batch's features while the GPU is processing the current batch. PyTorch's DataLoader with num_workers > 0 handles this automatically when the dataset's __getitem__ calls store.fetch.

Batch size: Larger batches amortize the fixed overhead of indexing and disk seeks. If individual fetch calls are fast, a batch of 256 nodes is more efficient than 4 batches of 64 nodes.

Organizing a Feature Store for Multi-Dataset Experiments

When working across several datasets or multiple graph splits, organizing feature stores on disk becomes important. A practical layout:

/data/
          cora/
            features.bin          # node features, shape stored in metadata.json
            metadata.json         # {"shape": [2708, 1433], "dtype": "float32", "num_nodes": 2708}
            edge_index.pt         # torch.save'd edge_index tensor
            labels.pt             # torch.save'd label tensor
          my_spatial_graph/
            features.bin          # [N, C, H, W] spatial features
            metadata.json
            edge_index.pt
            labels.pt
        

Storing metadata in a sidecar JSON file avoids having to hardcode the shape when reopening the store:

import json
        
        def save_store_metadata(path, shape, dtype):
            meta = {"shape": list(shape), "dtype": str(dtype), "num_nodes": shape[0]}
            with open(path.replace('.bin', '_metadata.json'), 'w') as f:
                json.dump(meta, f)
        
        def load_store_from_directory(directory):
            import os
            bin_file = os.path.join(directory, 'features.bin')
            meta_file = os.path.join(directory, 'metadata.json')
            with open(meta_file) as f:
                meta = json.load(f)
            shape = tuple(meta['shape'])
            dtype = np.dtype(meta['dtype'])
            return FeatureStore.load(bin_file, shape, dtype)
        

This pattern makes it easy to add new datasets to an experiment suite without changing any code — just add a new directory and the loader discovers the shape from metadata.

Limitations and Honest Notes

The feature store is most useful when features are fixed after dataset creation. It is not designed for use cases where node features are updated dynamically during training (for example, in online learning or continual learning settings).

The memmap approach is read-optimized. Writing to a memmap array that is actively being read from another process can cause data races. The feature store should be treated as read-only during training.

Memory-mapped I/O performance degrades on network-attached storage (NAS, NFS mounts). The OS's page cache cannot effectively cache randomly accessed pages from a network drive. For cloud training, copying the feature file to local NVMe storage before training begins is strongly recommended.

The feature store does not handle distributed training across multiple machines automatically. Each training process must be able to read the feature file. In a distributed setting, you need either a shared filesystem that all workers can access or a copy of the feature file on each worker's local storage.

NumPy's memmap does not support all dtypes for all shapes. For complex tensor types or features requiring preprocessing at load time, a custom Dataset backed by memmap with a transform applied in __getitem__ is more flexible than the feature store directly.

Frequently Asked Questions

Can I use the feature store with non-image data, like spectrograms or time series per node?
Yes. The feature store is agnostic to the meaning of the stored tensor. Any fixed-shape tensor per node can be stored, as long as all nodes share the same feature shape.

What happens if the node features don't all have the same shape?
NumPy's memmap requires a uniform dtype and shape. If nodes have variable-length features, you need either padding to a common length or a more flexible store like HDF5 (via h5py) with variable-length datasets.

Is the feature store compatible with PyTorch's DataLoader?
Yes. Wrap the feature store in a torch.utils.data.Dataset and implement __getitem__ to call store.fetch. The DataLoader handles batching and prefetching.