def ReplaceDeviceTag(parallel_desc_symbol, device_tag, builder=None):
assert parallel_desc_symbol.device_tag != device_tag
parallel_conf = placement_cfg.ParallelConf()
parallel_conf.set_device_tag(device_tag)
for device_name in parallel_desc_symbol.parallel_conf.device_name():
parallel_conf.add_device_name(device_name)
if builder is None:
return oneflow_api.PlacementSymbol(parallel_desc_symbol.symbol_id,
parallel_conf)
else:
return builder.GetParallelDescSymbol(parallel_conf)
def ReplaceDeviceTag(parallel_desc_symbol, device_tag, builder=None):
assert parallel_desc_symbol.device_tag != device_tag
parallel_conf = placement_cfg.ParallelConf()
parallel_conf.set_device_tag(device_tag)
for device_name in parallel_desc_symbol.parallel_conf.device_name():
parallel_conf.add_device_name(device_name)
hierarchy = shape_proto_cfg.ShapeProto()
for dim in parallel_desc_symbol.hierarchy:
hierarchy.add_dim(dim)
assert hierarchy.dim_size() > 0
parallel_conf.mutable_hierarchy().CopyFrom(hierarchy)
if builder is None:
return oneflow_api.PlacementSymbol(parallel_desc_symbol.symbol_id,
parallel_conf)
else:
return builder.GetParallelDescSymbol(parallel_conf)
def RandomParallelIdPerMachine(parallel_desc_symbol,
device_tag=None,
builder=None):
if device_tag is None:
device_tag = parallel_desc_symbol.parallel_conf.device_tag()
assert device_tag is not None
parallel_conf = placement_cfg.ParallelConf()
parallel_conf.set_device_tag(device_tag)
for machine_id, dev_ids in parallel_desc_symbol.machine_id2device_id_list.items(
):
dev_id = dev_ids[random.randint(0, len(dev_ids) - 1)]
parallel_conf.add_device_name("@%s:%s" % (machine_id, dev_id))
if builder is None:
return oneflow_api.PlacementSymbol(parallel_desc_symbol.symbol_id,
parallel_conf)
else:
return builder.GetParallelDescSymbol(parallel_conf)