TorchBasedFeatureStore

class dgl.graphbolt.TorchBasedFeatureStore(feat_data: List[dgl.graphbolt.impl.ondisk_metadata.OnDiskFeatureData])

Bases: dgl.graphbolt.impl.basic_feature_store.BasicFeatureStore

A store to manage multiple pytorch based feature for access.

The feature stores are described by the feat_data. The feat_data is a list of OnDiskFeatureData.

For a feature store, its format must be either “pt” or “npy” for Pytorch or Numpy formats. If the format is “pt”, the feature store must be loaded in memory. If the format is “npy”, the feature store can be loaded in memory or on disk.

Parameters

feat_data (List[OnDiskFeatureData]) – The description of the feature stores.

Examples

>>> import torch
>>> import numpy as np
>>> from dgl import graphbolt as gb
>>> edge_label = torch.tensor([[1], [2], [3]])
>>> node_feat = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> torch.save(edge_label, "/tmp/edge_label.pt")
>>> np.save("/tmp/node_feat.npy", node_feat.numpy())
>>> feat_data = [
...     gb.OnDiskFeatureData(domain="edge", type="author:writes:paper",
...         name="label", format="torch", path="/tmp/edge_label.pt",
...         in_memory=True),
...     gb.OnDiskFeatureData(domain="node", type="paper", name="feat",
...         format="numpy", path="/tmp/node_feat.npy", in_memory=False),
... ]
>>> feature_sotre = gb.TorchBasedFeatureStore(feat_data)

pin_memory_()

In-place operation to copy the feature store to pinned memory.