MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / get_tfkeras_model

Function get_tfkeras_model

tools/tensorflow-quantization/examples/utils.py:37–76  ·  view source on GitHub ↗
(model_name: str = "mobilenet_v1", shape: Tuple = None)

Source from the content-addressed store, hash-verified

35
36
37def get_tfkeras_model(model_name: str = "mobilenet_v1", shape: Tuple = None) -> tf.keras.Model:
38 """
39 Creates a native tf.keras.applications model.
40
41 Args:
42 model_name (str): Options={model_name_options}.
43
44 Returns:
45 model (tf.keras.Model): model corresponding to 'model_name'.
46
47 Raises:
48 ValueError: raised when 'model_name' is not supported.
49 """.format(
50 model_name_options=list(MODELS_CLASSES_DICT.keys())
51 )
52 try:
53 model_class = MODELS_CLASSES_DICT[model_name]
54 except ValueError:
55 raise ValueError("Model {} was not found!".format(model_name))
56 print("Loading model as {}".format(model_class))
57
58 if shape is None:
59 shape = (
60 _DEFAULT_IMAGE_SIZE[model_name],
61 _DEFAULT_IMAGE_SIZE[model_name],
62 _NUM_CHANNELS,
63 )
64
65 input_img = tf.keras.layers.Input(shape=shape, name="input_1")
66 model = model_class(
67 include_top=True,
68 weights="imagenet",
69 input_tensor=input_img,
70 input_shape=None,
71 pooling=None,
72 classes=_NUM_CLASSES,
73 classifier_activation="softmax",
74 )
75
76 return model
77
78
79def print_model_weights_shapes(model):

Callers 6

mainFunction · 0.90
get_resnet_modelFunction · 0.90
mainFunction · 0.90

Calls 2

printFunction · 0.85
keysMethod · 0.45