Scaling Tensor Node Features: The TGraphX Feature Store

A flat-vector graph with a million nodes and 128 features fits comfortably in memory: that is half a gigabyte. Replace each node with a [16, 8, 8] feature map and the same million nodes need roughly sixteen times as much — and [3, 64, 64] image nodes blow past most GPUs entirely. Tensor-valued node features are simply bigger, so a tensor-aware framework needs a story for keeping features somewhere other than a single resident tensor. TGraphX answers with a small feature-store abstraction in tgraphx/feature_store.py. This tutorial shows what it offers and when to reach for it.

This pairs naturally with mini-batch sampling — see Neighbor Sampling for Large Graphs — because the point of a feature store is to fetch only the nodes a batch touches.

Two stores, one interface

The package ships two implementations with a deliberately similar surface:

Both expose put(name, tensor, …) to register a named feature block and get(...) to retrieve rows, plus contains, feature_names, and metadata helpers. InMemoryFeatureStore adds memory_estimate_bytes() and summary() so you can see what you are holding, and a to(device) method to move everything at once. A FeatureStoreError is raised for misuse rather than letting a bad access fail obscurely.

from tgraphx.feature_store import InMemoryFeatureStore
        
        store = InMemoryFeatureStore()
        store.put("x", node_features)            # node_features: [N, 16, 8, 8]
        print(store.summary())                    # names, shapes, estimated bytes
        rows = store.get("x", indices)            # fetch only the rows you need
        

Memory-mapping for out-of-core features

When the feature block is larger than memory, MemmapFeatureStore(root) writes each named block as a NumPy .npy file under root and reads slices lazily through numpy.memmap. The operating system pages in only the bytes you actually touch, so fetching a mini-batch of a few thousand nodes does not load the whole array. The trade-off is explicit in the source: this path requires NumPy (it raises ImportError otherwise), and disk reads are slower than resident memory — a classic capacity-for-latency exchange. Whether it pays off depends on your data size and access pattern, which you should profile rather than assume; the package ships the mechanism, not a benchmark that would justify a speed claim in the abstract.

The point: fetch only what a batch needs

The store earns its keep in combination with sampling. When NeighborLoader produces a GraphMiniBatch (covered in the batching internals article), the nodes in that batch are a small subset of the graph. tgraphx/loaders provides fetch_features_for_subgraph, which pulls exactly those nodes' features out of a store and maps global node ids to the batch's local ids. The full feature tensor never has to be resident; only the per-batch slice is materialised. That is the standard recipe for training on graphs whose feature tensors do not fit in memory.

Choosing a store

A simple decision rule:

• Everything fits in VRAM → keep features as a plain tensor on the graph; you do not need a store.
• Fits in RAM but not VRAM, or shared across many batches → InMemoryFeatureStore, then move per-batch slices to the device.
• Larger than RAM → MemmapFeatureStore, accept disk-read latency, and ensure NumPy is installed.

The capability map labels the feature store as Beta: tested and documented, with the API stable within the v1.x series. Treat it as a solid foundation rather than a fully tuned production data layer — for very large workloads you will still want to profile I/O and consider sharding.

A sizing rule of thumb

A quick calculation tells you which store you need. A node tensor of shape [C, H, W] in float32 occupies 4·C·H·W bytes; multiply by N nodes for the resident feature block. For one million [16, 8, 8] nodes that is 4 · 16 · 8 · 8 · 10⁶ ≈ 4.1 GB — comfortable in system RAM, tight on a consumer GPU, and the reason an InMemoryFeatureStore on CPU with per-batch device transfers often serves better than keeping everything resident on the accelerator. Change the node tensor to [3, 64, 64] and the same million nodes need roughly 49 GB, which is squarely MemmapFeatureStore territory. The rule of thumb: estimate 4·∏(s)·N bytes, compare it against your RAM and VRAM budgets, and pick the store accordingly. Because disk-backed access trades latency for capacity, profile the I/O on a representative batch before committing to memmap for a latency-sensitive loop — the package gives you the mechanism, and the source exposes the hooks, but it does not promise that disk reads will keep your accelerator fed.

As a final caveat, a feature store solves capacity, not bandwidth: if your bottleneck is moving features to the device each step rather than holding them in memory, no store fixes that, and you should profile the transfer before assuming memmap is the answer.

Related guides

• Neighbor Sampling for Large Graphs in TGraphX
• The TGraphX Graph and GraphBatch: Storage and Batching
• Getting Started with Tensor-Valued Nodes

Conclusion

Because tensor node features are large, TGraphX separates "where features live" from "which features a batch needs." InMemoryFeatureStore keeps them resident; MemmapFeatureStore pages them from disk; fetch_features_for_subgraph pulls just the mini-batch slice. The abstraction in tgraphx/feature_store.py is modest and Beta-labelled, but it is exactly the seam you need to train on graphs whose features would never fit in a single tensor.