Open3D 的 open3d.core.Tensor
是一个用于表示多维张量的类。该类支持在 GPU 上进行计算。在进行 CUDA 计算时,可以使用 open3d.torch.device
将 Tensor
对象放置在 CUDA 设备上。
创建 CUDA 张量
import open3d.core as o3c
device = o3c.torch.Device('cuda:0')
a = o3c.Tensor.ones((3, 4), o3c.float32, device)
在 CUDA 设备上进行计算
b = o3c.Tensor.rand((3, 4), o3c.float32, device)
c = o3c.core.ops.add(a, b)
在 CUDA 设备上的张量操作与在 CPU 上的张量操作类似,除了需要将张量放置在 CUDA 设备上。
# 创建
a = o3c.Tensor.ones((2, 3), o3c.float32, device=o3c.torch.Device('cuda:0'))
# 运算
b = o3c.Tensor.full((2, 3), 2.0, o3c.float32, device=o3c.torch.Device('cuda:0'))
c = o3c.core.ops.add(a, b)
# 转换
d = a.cpu()
o3c.torch.device
显式指定计算设备。