e6ef08f367e0389e811d63eaa5afb16183a19e2b,src/sdk/pynni/nni/compression/tensorflow/compressor.py,,_locate_layers,#Any#Any#,240

Before Change

    // TODO: it can only find layers in `Model` and `list` for now.
    ret = {}

    if isinstance(model, tf.keras.Model):
        for key, value in model.__dict__.items():
            if isinstance(value, tf.keras.Model):
                ret.update(_locate_layers(value, cur_path + [key]))
            elif isinstance(value, list):
                ret.update(_locate_layers(value, cur_path + [key]))
            elif isinstance(value, tf.keras.layers.Layer):
                ret[id(value)] = LayerInfo(value, cur_path + [key])

    elif isinstance(model, list):
        for i, item in enumerate(model):
            if isinstance(item, tf.keras.Model):
                ret.update(_locate_layers(item, cur_path + [i]))
            elif isinstance(item, tf.keras.layers.Layer):
                ret[id(item)] = LayerInfo(item, cur_path + [i])

    else:
        raise ValueError("Unexpected model type: {}".format(type(model)))
    return ret

def _select_config(layer_info, config_list):
    // Find the last matching config block for given layer.

After Change

    // `cur_path`s format is documented in `LayerInfo.path`.
    // TODO: it can only find layers in `Model` and `list` for now.
    assert isinstance(model, tf.keras.Model)
    if isinstance(model, tf.keras.Sequential):
        _logger.warning("`Sequential` model is not supported yet, ignored.")
    ret = {}
    for key, value in model.__dict__.items():
        if isinstance(value, tf.keras.Model):
            ret.update(_locate_layers(value, cur_path + [key]))

Instances