GPUCachedFeature

class dgl.graphbolt.GPUCachedFeature(fallback_feature: Feature, cache: GPUFeatureCache, offset: int = 0)

基类: Feature

GPU 缓存特征，包装了一个备用特征。它使用最近最少使用 (LRU) 算法作为缓存淘汰策略。使用 gpu_cached_feature 来构造此类的实例。

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

参数:

• fallback_feature (Feature) – 备用特征。

• cache (GPUFeatureCache) – 用作缓存后端的 GPUFeatureCache 实例。

• offset (int, optional) – 在使用缓存之前，添加到给定 ids 的偏移值。如果多个 GPUCachedFeature 共享一个 GPUFeatureCache 对象，此参数很有用。

示例

>>> 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.gpu_cached_feature(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])

count()

获取特征的计数。

返回:

特征的计数。

返回类型:

int

read(ids: Tensor | None = None)

按索引读取特征。

返回的张量始终位于 GPU 内存中，无论备用特征是在内存中还是在磁盘上。

参数:

ids (torch.Tensor, optional) – 特征的索引。如果指定，则仅读取特征的指定索引。如果为 None，则返回整个特征。

返回:

读取的特征。

返回类型:

torch.Tensor

read_async(ids: Tensor)

按索引异步读取特征。

参数:

ids (torch.Tensor) – 特征的索引。仅读取特征的指定索引。

返回:

返回的生成器对象在第 read_async_num_stages(ids.device) 次调用时返回一个 future。可以通过调用返回的 future 对象的 .wait() 来访问返回结果。多次调用 .wait() 是未定义行为。

返回类型:

一个生成器对象。

示例

>>> import dgl.graphbolt as gb
>>> feature = gb.Feature(...)
>>> ids = torch.tensor([0, 2])
>>> for stage, future in enumerate(feature.read_async(ids)):
...     pass
>>> assert stage + 1 == feature.read_async_num_stages(ids.device)
>>> result = future.wait()  # result contains the read values.

read_async_num_stages(ids_device: device)

read_async 操作的阶段数。请参阅 read_async 函数了解其使用说明。当 read_async 用于位于 ids_device 上的张量时，此函数需要返回 yield 操作的数量。

参数:

ids_device (torch.device) – 传入 read_async 的 ids 参数的设备。

返回:

read_async 操作的阶段数。

返回类型:

int

size()

获取特征的大小。

返回:

特征的大小。

返回类型:

torch.Size

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

更新特征。

参数:

• value (torch.Tensor) – 特征的更新值。

• ids (torch.Tensor, optional) – 要更新的特征索引。如果指定，则仅更新特征的指定索引。对于该特征，ids[i] 行将更新为 value[i]。因此索引和值必须具有相同的长度。如果为 None，则将更新整个特征。

property cache_size_in_bytes

返回缓存占用的字节大小。

property miss_rate

返回自创建以来的缓存未命中率。