TorchBasedFeatureStoreΒΆ

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

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)
is_pinned()[source]ΒΆ

Returns True if all the stored features are pinned.

pin_memory_()[source]ΒΆ

In-place operation to copy the feature store to pinned memory. Returns the same object modified in-place.

to(device)[source]ΒΆ

Copy TorchBasedFeatureStore to the specified device.