Function data_parallelism

src_code/thumt1_code/thumt/utils/parallel.py:23–54

(devices, fn, *args, **kwargs)

# Data-level parallelism
def data_parallelism(devices, fn, *args, **kwargs):
    num_worker = len(devices)
    devices = ["gpu:%d" % d for d in devices]

    # Replicate args and kwargs
    if args:
        new_args = [_maybe_repeat(arg, num_worker) for arg in args]
        # Transpose
        new_args = [list(x) for x in zip(*new_args)]
    else:
        new_args = [[] for _ in range(num_worker)]

    new_kwargs = [{} for _ in range(num_worker)]

    for k, v in six.iteritems(kwargs):
        vals = _maybe_repeat(v, num_worker)

        for i in range(num_worker):
            new_kwargs[i][k] = vals[i]

    fns = _maybe_repeat(fn, num_worker)

    # Now make the parallel call.
    outputs = []

    for i in range(num_worker):
        with tf.variable_scope(tf.get_variable_scope(), reuse=(i != 0)):
            with tf.name_scope("parallel_%d" % i):
                with tf.device(devices[i]):
                    outputs.append(fns[i](*new_args[i], **new_kwargs[i]))

    return outputs


def shard_features(features, device_list):