INNER CODE UNIT · Python

load_model

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

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)
            for cls in custom_optimizers
        })

    if custom_objects is not None:
        byteps_objects.update(custom_objects)

    return keras.models.load_model(filepath, custom_objects=byteps_objects)

View source record →