INNER CODE UNIT · Python

push_pull

bytedance/byteps · byteps/_keras/__init__.py:97

def push_pull(backend, value, name, average):
    return _eval(backend,  bps.push_pull(tf.constant(value, name=name), average=average))


def broadcast(backend, value, root_rank, name):
    return _eval(backend, bps.broadcast(tf.constant(value, name=name), root_rank, is_variable=False))


def load_model(keras, wrap_optimizer, optimizer_modules, filepath, custom_optimizers, custom_objects):
    byteps_objects = {
        subclass.__name__.lower(): wrap_optimizer(subclass)
        for subclass in keras.optimizers.Optimizer.__subclasses__()
        if subclass.__module__ in optimizer_modules
    }

    if custom_optimizers is not None:
        byteps_objects.update({
            cls.__name__: wrap_optimizer(cls)

View source record →