def multi_lbi_job(x: oft.Numpy.Placeholder((96, 96, 96))):
with flow.scope.placement(src_device_type, "0:0-" + str(src_device_num - 1)):
src_s0 = flow.identity(x.with_distribute(flow.distribute.split(0)))
src_s1 = flow.identity(x.with_distribute(flow.distribute.split(1)))
src_b = flow.identity(x.with_distribute(flow.distribute.split(1)))
(t0_0, t0_1, t0_2) = flow.identity_n((src_s0, src_s1, src_b))
with flow.scope.placement(dst_device_type, "0:0-" + str(dst_device_num - 1)):
t0_0 = t0_0.with_distribute(flow.distribute.split(1))
t0_1 = t0_1.with_distribute(flow.distribute.broadcast())
t0_2 = t0_2.with_distribute(flow.distribute.split(1))
(t1_0, t1_1, t1_2) = flow.identity_n((t0_0, t0_1, t0_2))
return t1_0, t1_1, t1_2
def _data_loader_ofrecord(data_dir,
data_part_num,
batch_size,
part_name_suffix_length=-1,
shuffle=True):
ofrecord = flow.data.ofrecord_reader(
data_dir,
batch_size=batch_size,
data_part_num=data_part_num,
part_name_suffix_length=part_name_suffix_length,
random_shuffle=shuffle,
shuffle_after_epoch=shuffle)
def _blob_decoder(bn, shape, dtype=flow.int32):
return flow.data.OFRecordRawDecoder(ofrecord,
bn,
shape=shape,
dtype=dtype)
labels = _blob_decoder("labels", (1, ))
dense_fields = _blob_decoder("dense_fields", (FLAGS.num_dense_fields, ),
flow.float)
wide_sparse_fields = _blob_decoder("wide_sparse_fields",
(FLAGS.num_wide_sparse_fields, ))
deep_sparse_fields = _blob_decoder("deep_sparse_fields",
(FLAGS.num_deep_sparse_fields, ))
return flow.identity_n(
[labels, dense_fields, wide_sparse_fields, deep_sparse_fields])
def _data_loader_synthetic(batch_size):
devices = ['{}:0-{}'.format(i, FLAGS.gpu_num_per_node - 1) for i in range(FLAGS.num_nodes)]
with flow.scope.placement("cpu", devices):
def _blob_random(shape, dtype=flow.int32, initializer=flow.zeros_initializer(flow.int32)):
return flow.data.decode_random(shape=shape, dtype=dtype, batch_size=batch_size,
initializer=initializer)
labels = _blob_random((1,), initializer=flow.random_uniform_initializer(dtype=flow.int32))
dense_fields = _blob_random((FLAGS.num_dense_fields,), dtype=flow.float,
initializer=flow.random_uniform_initializer())
wide_sparse_fields = _blob_random((FLAGS.num_wide_sparse_fields,))
deep_sparse_fields = _blob_random((FLAGS.num_deep_sparse_fields,))
print('use synthetic data')
return flow.identity_n([labels, dense_fields, wide_sparse_fields, deep_sparse_fields])
def _data_loader_ofrecord(data_dir,
data_part_num,
batch_size,
part_name_suffix_length=-1,
shuffle=True):
if data_dir:
ofrecord = flow.data.ofrecord_reader(
data_dir,
batch_size=batch_size,
data_part_num=data_part_num,
part_name_suffix_length=part_name_suffix_length,
random_shuffle=shuffle,
shuffle_after_epoch=shuffle)
def _blob_decoder(bn, shape, dtype=flow.int32):
return flow.data.OFRecordRawDecoder(ofrecord,
bn,
shape=shape,
dtype=dtype)
labels = _blob_decoder("labels", (1, ))
dense_fields = _blob_decoder("dense_fields",
(FLAGS.num_dense_fields, ), flow.float)
wide_sparse_fields = _blob_decoder("wide_sparse_fields",
(FLAGS.num_wide_sparse_fields, ))
deep_sparse_fields = _blob_decoder("deep_sparse_fields",
(FLAGS.num_deep_sparse_fields, ))
print('load data form', data_dir)
else:
def _blob_random(shape,
dtype=flow.int32,
initializer=flow.zeros_initializer(flow.int32)):
return flow.data.decode_random(shape=shape,
dtype=dtype,
batch_size=batch_size,
initializer=initializer)
labels = _blob_random(
(1, ),
initializer=flow.random_uniform_initializer(dtype=flow.int32))
dense_fields = _blob_random(
(FLAGS.num_dense_fields, ),
dtype=flow.float,
initializer=flow.random_uniform_initializer())
wide_sparse_fields = _blob_random((FLAGS.num_wide_sparse_fields, ))
deep_sparse_fields = _blob_random((FLAGS.num_deep_sparse_fields, ))
print('use synthetic data')
return flow.identity_n(
[labels, dense_fields, wide_sparse_fields, deep_sparse_fields])
def _data_loader_onerec(data_dir, batch_size, shuffle):
assert data_dir
print('load onerec data form', data_dir)
files = glob.glob(os.path.join(data_dir, '*.onerec'))
readdata = flow.data.onerec_reader(files=files, batch_size=batch_size, random_shuffle=shuffle,
verify_example=False,
shuffle_buffer_size=64,
shuffle_after_epoch=shuffle)
def _blob_decoder(bn, shape, dtype=flow.int32):
return flow.data.onerec_decoder(readdata, key=bn, shape=shape, dtype=dtype)
labels = _blob_decoder('labels', shape=(1,))
dense_fields = _blob_decoder("dense_fields", (FLAGS.num_dense_fields,), flow.float)
wide_sparse_fields = _blob_decoder("wide_sparse_fields", (FLAGS.num_wide_sparse_fields,))
deep_sparse_fields = _blob_decoder("deep_sparse_fields", (FLAGS.num_deep_sparse_fields,))
return flow.identity_n([labels, dense_fields, wide_sparse_fields, deep_sparse_fields])
def identity_n_job(xs: Tuple[(oft.Numpy.Placeholder((5, 2)), ) * 3]):
return flow.identity_n(xs)