Skip to content

Commit e0f399d

Browse files
peri044peri044
committed
Update frozen graph script and instructions
1 parent c4f90be commit e0f399d

3 files changed

Lines changed: 35 additions & 11 deletions

File tree

‎TensorFlow/Classification/ConvNets/export_frozen_graph.py‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,10 @@
66
import tensorflow as tf
77

88
import horovod.tensorflow as hvd
9-
from model import resnet_v1_5
9+
from model import resnet
1010

1111
tf.app.flags.DEFINE_string(
12-
'model_name', 'resnet50_v1.5', 'The name of the architecture to save. The default name was being '
12+
'model_name', 'resnet50', 'The name of the architecture to save. The default name was being '
1313
'used to train the model')
1414

1515
tf.app.flags.DEFINE_integer(
@@ -26,10 +26,10 @@
2626
'be specified at model runtime.')
2727

2828

29-
tf.app.flags.DEFINE_string('input_format', 'NHWC',
29+
tf.app.flags.DEFINE_string('input_format', 'NCHW',
3030
'The dataformat used by the layers in the model')
3131

32-
tf.app.flags.DEFINE_string('compute_format', 'NHWC',
32+
tf.app.flags.DEFINE_string('compute_format', 'NCHW',
3333
'The dataformat used by the layers in the model')
3434

3535
tf.app.flags.DEFINE_string('checkpoint', '',
@@ -72,15 +72,24 @@ def main(_):
7272
else:
7373
input_shape = [FLAGS.batch_size, FLAGS.image_size, FLAGS.image_size, 3]
7474
input_images = tf.placeholder(name='input', dtype=tf.float32, shape=input_shape)
75-
network = resnet_v1_5.ResnetModel(FLAGS.model_name, FLAGS.num_classes, FLAGS.compute_format, FLAGS.input_format)
75+
76+
resnet50_config = resnet.model_architectures[FLAGS.model_name]
77+
network = resnet.ResnetModel(FLAGS.model_name,
78+
FLAGS.num_classes,
79+
resnet50_config['layers'],
80+
resnet50_config['widths'],
81+
resnet50_config['expansions'],
82+
FLAGS.compute_format,
83+
FLAGS.input_format)
7684
probs, logits = network.build_model(
7785
input_images,
7886
training=False,
7987
reuse=False,
8088
use_final_conv=FLAGS.use_final_conv)
81-
89+
8290
if FLAGS.quantize:
83-
tf.contrib.quantize.experimental_create_eval_graph(symmetric=FLAGS.symmetric, use_qdq=FLAGS.use_qdq)
91+
tf.contrib.quantize.experimental_create_eval_graph(symmetric=FLAGS.symmetric,
92+
use_qdq=FLAGS.use_qdq)
8493

8594
# Define the saver and restore the checkpoint
8695
saver = tf.train.Saver()
@@ -101,4 +110,4 @@ def main(_):
101110

102111

103112
if __name__ == '__main__':
104-
tf.app.run()
113+
tf.app.run()

‎TensorFlow/Classification/ConvNets/model/layers/conv2d.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@ def conv2d(
3131
use_bias=True,
3232
kernel_initializer=tf.variance_scaling_initializer(),
3333
bias_initializer=tf.zeros_initializer(),
34-
trainable=True
34+
trainable=True,
35+
name=None
3536
):
3637

3738
if data_format not in ['NHWC', 'NCHW']:
@@ -52,7 +53,8 @@ def conv2d(
5253
kernel_initializer=kernel_initializer,
5354
bias_initializer=bias_initializer,
5455
trainable=trainable,
55-
activation=None
56+
activation=None,
57+
name=name
5658
)
5759

5860
return net

‎TensorFlow/Classification/ConvNets/resnet50v1.5/README.md‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -373,7 +373,7 @@ It is recommended to finetune a model with quantization nodes rather than train
373373

374374
For QAT network, we use <a href="https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/quantization/quantize_and_dequantize">tf.quantization.quantize_and_dequantize operation</a>.
375375
These operations are automatically added at weights and activation layers in the RN50 by using `tf.contrib.quantize.experimental_create_training_graph` utility. Support for using `tf.quantization.quantize_and_dequantize`
376-
operations for `tf.contrib.quantize.experimental_create_training_graph has been added in <a href="https://ngc.nvidia.com/catalog/containers/nvidia:tensorflow">TensorFlow 20.01-py3 NGC container</a> and later versions, which is required for this task.
376+
operations for `tf.contrib.quantize.experimental_create_training_graph` has been added in <a href="https://ngc.nvidia.com/catalog/containers/nvidia:tensorflow">TensorFlow 20.01-py3 NGC container</a> and later versions, which is required for this task.
377377

378378
#### Post process checkpoint
379379
* `post_process_ckpt.py` is a utility to convert the final classification FC layer into a 1x1 convolution layer using the same weights. This is required to ensure TensorRT can parse QAT models successfully.
@@ -382,6 +382,19 @@ operations for `tf.contrib.quantize.experimental_create_training_graph has been
382382
* `--ckpt` : Path to the trained checkpoint of RN50.
383383
* `--out` : Name of the new checkpoint file which has the FC layer weights reshaped into 1x1 conv layer weights.
384384

385+
### Exporting Frozen graphs
386+
To export frozen graphs (which can be used for inference with <a href="https://developer.nvidia.com/tensorrt">TensorRT</a>), use:
387+
388+
`python export_frozen_graph.py --checkpoint <path_to_checkpoint> --quantize --use_final_conv --use_qdq --symmetric --input_format NCHW --compute_format NCHW --output_file=<output_file_name>`
389+
390+
Arguments:
391+
392+
* `--checkpoint` : Optional argument to export the model with checkpoint weights.
393+
* `--quantize` : Optional flag to export quantized graphs.
394+
* `--use_qdq` : Use quantize_and_dequantize (QDQ) op instead of FakeQuantWithMinMaxVars op for quantization. QDQ does only scaling.
395+
* `--input_format` : Data format of input tensor (Default: NCHW). Use NCHW format to optimize the graph with TensorRT.
396+
* `--compute_format` : Data format of the operations in the network (Default: NCHW). Use NCHW format to optimize the graph with TensorRT.
397+
385398
### Inference process
386399
To run inference on a single example with a checkpoint and a model script, use:
387400

0 commit comments

Comments
 (0)