oneflow.placement

OneFlow placement Class

class oneflow.placement

A oneflow.palcement is an object representing the device group on which a oneflow.Tensor is or will be allocated. The oneflow.palcement contains a device type (‘cpu’ or ‘cuda’) and corresponding device sequence.

A oneflow.Tensor’s palcement can be accessed via the Tensor.palcement property.

A oneflow.palcement can be constructed in several ways:

>>> import oneflow as flow

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

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

Return a placement containing all devices of all machines under env.

Parameters

device_type (str) – cuda or cpu

For examples:

# 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,))