Tensor Attributes

每个本地 oneflow.Tensor 都有一个 oneflow.dtypeoneflow.device ,而每个全局 oneflow.Tensor 都有一个 oneflow.dtype, oneflow.placementoneflow.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 有三种类型:

  1. split

    表示物理设备上的 Tensor 是将全局视角下的 Tensor 切分得到的。切分时需要指定切分的维度。物理设备上的 Tensor 如果根据切分的维度被拼接,可以还原得到全局视角的 Tensor 。

  2. broadcast

    表示物理设备上的 Tensor 是全局视角下的 Tensor 完全相同的拷贝。

  3. 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