GPU缓存特性

class dgl.graphbolt.GPUCachedFeature(fallback_feature: Feature, max_cache_size_in_bytes: int)

Bases: Feature

GPU缓存的特性包装了一个备用特性。

将GPU缓存放置在torch.cuda.current_device()上。

Parameters:

• fallback_feature (Feature) – 备用特征。

• max_cache_size_in_bytes (int) – GPU缓存的容量，以字节为单位。

示例

>>> 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)

通过索引读取特征。

返回的张量始终在GPU内存中，无论回退功能是在内存中还是在磁盘上。

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:

读取功能。

Return type:

torch.Tensor

size()

获取特征的大小。

Returns:

特征的大小。

Return type:

torch.Size

update(value: Tensor, ids: Tensor | None = None)

更新功能。

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.