Update instructions and output node name

TensorFlow/Classification/ConvNets/postprocess_ckpt.py

@@ -1,11 +1,9 @@
 import tensorflow as tf
-import pdb
 import numpy as np
 import argparse
 import os
-import shutil
         
-def process_checkpoint(input_ckpt, output_ckpt_path):
+def process_checkpoint(input_ckpt, output_ckpt_path, dense_layer):
     """
     This function loads a RN50 checkpoint with Dense layer as the final layer 
     and transforms the final dense layer into a 1x1 convolution layer. The weights
@@ -15,25 +13,23 @@ def process_checkpoint(input_ckpt, output_ckpt_path):
     Returns:
         None. New checkpoint with 1x1 conv layer as classification layer is generated.
     """
-    
     with tf.Session() as sess:
         # Load all the variables
-        all_vars = tf.train.list_variables(ckpt)
-        ckpt_reader = tf.train.load_checkpoint(ckpt)
+        all_vars = tf.train.list_variables(input_ckpt)
         # Capture the dense layer weights and reshape them to a 4D tensor which would be 
         # the weights of a 1x1 convolution layer. This code replaces the dense (FC) layer
         # to a 1x1 conv layer. 
         dense_layer_value=0.
         new_var_list=[]
         for var in all_vars:
-            curr_var = tf.train.load_variable(ckpt, var[0])
+            curr_var = tf.train.load_variable(input_ckpt, var[0])
             if var[0]==dense_layer:
                 dense_layer_value = curr_var
             else:
                 new_var_list.append(tf.Variable(curr_var, name=var[0]))
-        
+ 
         dense_layer_shape = [1, 1, 2048, 1001]
-        new_var_value = np.reshape(dense_layer_value, )
+        new_var_value = np.reshape(dense_layer_value, dense_layer_shape)
         new_var = tf.Variable(new_var_value, name=dense_layer)
         new_var_list.append(new_var)
         
@@ -44,10 +40,9 @@ def process_checkpoint(input_ckpt, output_ckpt_path):
 if __name__=='__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--input', type=str, required=True, help='Path to pretrained RN50 checkpoint with dense layer')
-    parser.add_argument('--dense_layer', type=str, default='resnet/output/dense/kernel')
+    parser.add_argument('--dense_layer', type=str, default='resnet50/output/dense/kernel')
     parser.add_argument('--output', type=str, default='output_dir', help="Output directory to store new checkpoint")
     args = parser.parse_args()
-    main(args)
     
     input_ckpt = args.input
     # Create an output directory
@@ -55,12 +50,8 @@ if __name__=='__main__':
     
     new_ckpt='new.ckpt'
     new_ckpt_path = os.path.join(args.output, new_ckpt)
-    with open(os.path.join(output_dir, "checkpoint"), 'w') as file:
+    with open(os.path.join(args.output, "checkpoint"), 'w') as file:
         file.write("model_checkpoint_path: "+ "\"" + new_ckpt + "\"")
         
     # Process the input checkpoint, apply transforms and generate a new checkpoint.
-    process_checkpoint(input_ckpt, new_ckpt_path)
-    
-    
-
-    
+    process_checkpoint(input_ckpt, new_ckpt_path, args.dense_layer)

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

@@ -393,11 +393,12 @@ These operations are automatically added at weights and activation layers in the
 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.
 
 #### Post process checkpoint
-  * `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.
+  `postprocess_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.
   This script should be used after performing QAT to reshape the FC layer weights in the final checkpoint.
   Arguments:
-     * `--ckpt` : Path to the trained checkpoint of RN50.
-     * `--out` : Name of the new checkpoint file which has the FC layer weights reshaped into 1x1 conv layer weights.
+     * `--input` : Path to the trained checkpoint of RN50.
+     * `--output` : Name of the new checkpoint file which has the FC layer weights reshaped into 1x1 conv layer weights.
+     * `--dense_layer` : Name of the FC layer
 
 ### Exporting Frozen graphs
 To export frozen graphs (which can be used for inference with <a href="https://developer.nvidia.com/tensorrt">TensorRT</a>), use: