Function get_resnet_model

tools/tensorflow-quantization/examples/resnet/utils.py:24–45

Creates a native tf.keras ResNet model. Args: resnet_depth (str): ResNet depth. Options=[50 (default), 101, 152]. resnet_version (str): ResNet version. Options=[v1 (default), v2]. Returns: model (tf.keras.Model): model corresponding to 'resnet_depth' and 'resne

(resnet_depth: str = "50", resnet_version: str = "v1")

def get_resnet_model(resnet_depth: str = "50", resnet_version: str = "v1") -> tf.keras.Model:
    """
    Creates a native tf.keras ResNet model.

    Args:
        resnet_depth (str): ResNet depth. Options=[50 (default), 101, 152].
        resnet_version (str): ResNet version. Options=[v1 (default), v2].

    Returns:
        model (tf.keras.Model): model corresponding to 'resnet_depth' and 'resnet_version'.
    """

    shape = (
        _DEFAULT_IMAGE_SIZE["resnet_{}".format(resnet_version)],
        _DEFAULT_IMAGE_SIZE["resnet_{}".format(resnet_version)],
        _NUM_CHANNELS,
    )

    model_name = "resnet_" + resnet_depth + resnet_version
    model = get_tfkeras_model(model_name=model_name, shape=shape)

    return model