GPUCachedFeatureο
- class dgl.graphbolt.GPUCachedFeature(fallback_feature: Feature, cache_size: int)[source]ο
Bases:
Feature
GPU cached feature wrapping a fallback feature.
Places the GPU cache to torch.cuda.current_device().
- Parameters:
Examples
>>> import torch >>> from dgl import graphbolt as gb >>> torch_feat = torch.arange(10).reshape(2, -1).to("cuda") >>> cache_size = 5 >>> fallback_feature = gb.TorchBasedFeature(torch_feat) >>> feature = gb.GPUCachedFeature(fallback_feature, cache_size) >>> feature.read() tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], device='cuda:0') >>> feature.read(torch.tensor([0]).to("cuda")) tensor([[0, 1, 2, 3, 4]], device='cuda:0') >>> feature.update(torch.tensor([[1 for _ in range(5)]]).to("cuda"), ... torch.tensor([1]).to("cuda")) >>> feature.read(torch.tensor([0, 1]).to("cuda")) tensor([[0, 1, 2, 3, 4], [1, 1, 1, 1, 1]], device='cuda:0') >>> feature.size() torch.Size([5])
- read(ids: Tensor | None = None)[source]ο
Read the feature by index.
The returned tensor is always in GPU memory, no matter whether the fallback feature is in memory or on disk.
- 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()[source]ο
Get the size of the feature.
- Returns:
The size of the feature.
- Return type:
torch.Size
- update(value: Tensor, ids: Tensor | None = None)[source]ο
Update the feature.
- 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.