评价此页

GdsFile#

class torch.cuda.gds.GdsFile(filename, flags)[source]#

cuFile 的包装器。

cuFile 是 GPUDirect Storage (GDS) API 的类文件接口。

有关更多详细信息,请参阅 cufile 文档

参数
  • filename (str) – 要打开的文件名。

  • flags (int) – 打开文件时传递给 os.open 的标志。 os.O_DIRECT 将自动添加。

示例

>>> src1 = torch.randn(1024, device="cuda")
>>> src2 = torch.randn(2, 1024, device="cuda")
>>> file = torch.cuda.gds.GdsFile(f, os.O_CREAT | os.O_RDWR)
>>> file.save_storage(src1.untyped_storage(), offset=0)
>>> file.save_storage(src2.untyped_storage(), offset=src1.nbytes)
>>> dest1 = torch.empty(1024, device="cuda")
>>> dest2 = torch.empty(2, 1024, device="cuda")
>>> file.load_storage(dest1.untyped_storage(), offset=0)
>>> file.load_storage(dest2.untyped_storage(), offset=src1.nbytes)
>>> torch.equal(src1, dest1)
True
>>> torch.equal(src2, dest2)
True
deregister_handle()[source]#

从 cuFile 驱动程序注销文件描述符。

这是 cuFileHandleDeregister 的包装器。

load_storage(storage, offset=0)[source]#

将数据从文件加载到存储中。

这是 cuFileRead 的包装器。将从文件中的 offset 位置加载 storage.nbytes() 字节的数据到存储中。

参数
  • storage (Storage) – 要加载数据的存储。

  • offset (int, optional) – 文件中的偏移量,从该位置开始加载。 (默认: 0)

register_handle()[source]#

向 cuFile 驱动程序注册文件描述符。

这是 cuFileHandleRegister 的包装器。

save_storage(storage, offset=0)[source]#

将存储中的数据保存到文件中。

这是 cuFileWrite 的包装器。存储中的所有字节将写入文件中的 offset 位置。

参数
  • storage (Storage) – 要从中保存数据的存储。

  • offset (int, optional) – 文件中的偏移量,从该位置开始保存。 (默认: 0)