TorchBasedFeature

class dgl.graphbolt.TorchBasedFeature(torch_feature: torch.Tensor, metadata: Optional[Dict] = None)

Bases: dgl.graphbolt.feature_store.Feature

A wrapper of pytorch based feature.

Initialize a torch based feature store by a torch feature. Note that the feature can be either in memory or on disk.

Parameters

torch_feature (torch.Tensor) – The torch feature. Note that the dimension of the tensor should be greater than 1.

Examples

>>> import torch
>>> from dgl import graphbolt as gb

1. The feature is in memory.

>>> torch_feat = torch.arange(10).reshape(2, -1)
>>> feature = gb.TorchBasedFeature(torch_feat)
>>> feature.read()
tensor([[0, 1, 2, 3, 4],
        [5, 6, 7, 8, 9]])
>>> feature.read(torch.tensor([0]))
tensor([[0, 1, 2, 3, 4]])
>>> feature.update(torch.tensor([[1 for _ in range(5)]]),
...                      torch.tensor([1]))
>>> feature.read(torch.tensor([0, 1]))
tensor([[0, 1, 2, 3, 4],
        [1, 1, 1, 1, 1]])
>>> feature.size()
torch.Size([5])

2. The feature is on disk.

>>> import numpy as np
>>> arr = np.array([[1, 2], [3, 4]])
>>> np.save("/tmp/arr.npy", arr)
>>> torch_feat = torch.from_numpy(np.load("/tmp/arr.npy", mmap_mode="r+"))
>>> feature = gb.TorchBasedFeature(torch_feat)
>>> feature.read()
tensor([[1, 2],
        [3, 4]])
>>> feature.read(torch.tensor([0]))
tensor([[1, 2]])

3. Pinned CPU feature. >>> torch_feat = torch.arange(10).reshape(2, -1).pin_memory() >>> feature = gb.TorchBasedFeature(torch_feat) >>> feature.read().device device(type=’cuda’, index=0) >>> feature.read(torch.tensor([0]).cuda()).device device(type=’cuda’, index=0)

metadata()

Get the metadata of the feature.

Returns

The metadata of the feature.

Return type

Dict

pin_memory_()

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

read(ids: Optional[torch.Tensor] = None)

Read the feature by index.

If the feature is on pinned CPU memory and ids is on GPU or pinned CPU memory, it will be read by GPU and the returned tensor will be on GPU. Otherwise, the returned tensor will be on CPU.

Parameters

ids (torch.Tensor, optional) – The index of the feature. If specified, only the specified indices of the feature are read. If None, the entire feature is returned.

Returns

The read feature.

Return type

torch.Tensor

size()

Get the size of the feature.

Returns

The size of the feature.

Return type

torch.Size

update(value: torch.Tensor, ids: Optional[torch.Tensor] = None)

Update the feature store.

Parameters

• value (torch.Tensor) – The updated value of the feature.

• ids (torch.Tensor, optional) – The indices of the feature to update. If specified, only the specified indices of the feature will be updated. For the feature, the ids[i] row is updated to value[i]. So the indices and value must have the same length. If None, the entire feature will be updated.