Tensor Attributes¶

每个本地 oneflow.Tensor 都有一个 oneflow.dtype 和 oneflow.device ，而每个全局 oneflow.Tensor 都有一个 oneflow.dtype, oneflow.placement 和 oneflow.sbp 。

oneflow.device¶

class oneflow.device¶

该文档参考自： https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.device

oneflow.device 是一个对象，用于指代一个 oneflow.Tensor 将被分配到的设备。

oneflow.device 包括一个设备类型（ ‘cpu’ 或 ‘cuda’ ）和一个可选的设备类型下的设备序号。如果设备序号不存在，则此对象将总是指代当前设备类型下的当前设备。

oneflow.device 所指代的设备可以通过 Tensor.device 属性获取。

oneflow.device 可以通过一个 string 构建，也可以通过 string 和设备序号共同构建。

通过 string 构建的示例：

>>> import oneflow as flow
>>> flow.device('cuda:0')
device(type='cuda', index=0)

>>> flow.device('cpu')
device(type='cpu', index=0)

>>> flow.device('cuda')  # 当前 cuda 设备
device(type='cuda', index=0)

通过 string 和设备序号构建的示例：

>>> import oneflow as flow
>>> flow.device('cuda', 0)
device(type='cuda', index=0)

>>> flow.device('cpu', 0)
device(type='cpu', index=0)

Note

The oneflow.device 参数在函数中通常可以被一个字符串替代。这保证了代码的快速原型化。

>>> import oneflow as flow
>>> # 一个使用了 oneflow.device 的示例函数
>>> cuda0 = flow.device('cuda:0')
>>> x = flow.randn(2,3, device=cuda0)

>>> # 可以将 flow.device 用 string 替代
>>> x = flow.randn(2,3, device='cuda:0')

oneflow.placement¶

class oneflow.placement¶

oneflow.placement 是一个对象，用于指代一个 oneflow.Tensor 被分配或即将被分配到的设备组。oneflow.placement 包含一个设备类型 (‘cpu’ 或者 ‘cuda’) 和对应的设备序列。

oneflow.Tensor 的 placement 可以通过 Tensor.placement 访问。

oneflow.placement 可以通过以下几种方式构造：

>>> import oneflow as flow

>>> p = flow.placement("cuda", ranks=[0, 1, 2, 3])
>>> p
oneflow.placement(type="cuda", ranks=[0, 1, 2, 3])
>>> p = flow.placement("cuda", ranks=[[0, 1], [2, 3]])
>>> p
oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]])

oneflow.env.all_device_placement¶

oneflow.env.all_device_placement(device_type: str) → oneflow._oneflow_internal.placement¶

返回 env 下所有的机器的所有设备的 placement。

参数：

device_type (str): cuda 或者 cpu

示例：

# world_size = 4, node_size = 1
import oneflow as flow

p = flow.env.all_device_placement("cuda") # oneflow.placement(device_type="cuda", machine_device_ids={0 : [0, 1, 2, 3]}, hierarchy=(4,))
p = flow.env.all_device_placement("cpu") # oneflow.placement(device_type="cpu", machine_device_ids={0 : [0, 1, 2, 3]}, hierarchy=(4,))

oneflow.sbp.sbp¶

class oneflow.sbp.sbp¶

一个 sbp 是一个表示全局视角下的 Tensor 与物理设备上的 Tensor 之间的映射关系的对象。

sbp 有三种类型：

split
表示物理设备上的 Tensor 是将全局视角下的 Tensor 切分得到的。切分时需要指定切分的维度。物理设备上的 Tensor 如果根据切分的维度被拼接，可以还原得到全局视角的 Tensor 。
broadcast
表示物理设备上的 Tensor 是全局视角下的 Tensor 完全相同的拷贝。
partial_sum
表示全局视角下的 Tensor 与物理设备上的 Tensor 的形状相同，但是物理设备上的值，只是全局视角下 Tensor 的一部分。返回的全局视角下的 Tensor 是通过把物理设备上的 Tensor 按照对应的位置（按元素）相加得到的。

一个 sbp 可以以多种方式被构造：

>>> import oneflow as flow

>>> s = flow.sbp.split(0)
>>> s
oneflow.sbp.split(axis=0)
>>> b = flow.sbp.broadcast
>>> b
oneflow.sbp.broadcast
>>> p = flow.sbp.partial_sum
>>> p
oneflow.sbp.partial_sum