dgl.graphbolt.gpu_cached_feature

dgl.graphbolt.gpu_cached_feature(fallback_features: Feature | Dict[FeatureKey, Feature], max_cache_size_in_bytes: int) → GPUCachedFeature | Dict[FeatureKey, GPUCachedFeature][source]

GPU 缓存特征，封装了一个回退特征。它使用最近最少使用 (LRU) 算法作为缓存淘汰策略。

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

参数：

  • fallback_features (Union[Feature, Dict[FeatureKey, Feature]]) – 回退特征。

  • max_cache_size_in_bytes (int) – GPU 缓存的容量（字节）。

返回值：

用 GPUCachedFeature 封装的特征。

返回类型：

Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]