Skip to content

Commit 6d2357a

Browse files
peri044peri044
committed
code Refactor
1 parent e0f399d commit 6d2357a

2 files changed

Lines changed: 38 additions & 23 deletions

File tree

‎TensorFlow/Classification/ConvNets/postprocess_ckpt.py‎

Lines changed: 37 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,28 +4,25 @@
44
import argparse
55
import os
66
import shutil
7-
8-
def main(args):
7+
8+
def process_checkpoint(input_ckpt, output_ckpt_path):
9+
"""
10+
This function loads a RN50 checkpoint with Dense layer as the final layer
11+
and transforms the final dense layer into a 1x1 convolution layer. The weights
12+
of the dense layer are reshaped into weights of 1x1 conv layer.
13+
Args:
14+
input_ckpt: Path to the input RN50 ckpt which has dense layer as classification layer.
15+
Returns:
16+
None. New checkpoint with 1x1 conv layer as classification layer is generated.
17+
"""
18+
919
with tf.Session() as sess:
10-
ckpt = args.ckpt
11-
new_ckpt=args.out
12-
output_dir = "./new_ckpt_dir"
13-
if os.path.isdir(output_dir):
14-
shutil.rmtree(output_dir)
15-
# Create an output directory
16-
os.mkdir(output_dir)
17-
18-
new_ckpt_path = os.path.join(output_dir, new_ckpt)
19-
with open(os.path.join(output_dir, "checkpoint"), 'w') as file:
20-
file.write("model_checkpoint_path: "+ "\"" + new_ckpt + "\"")
21-
file.close()
2220
# Load all the variables
2321
all_vars = tf.train.list_variables(ckpt)
2422
ckpt_reader = tf.train.load_checkpoint(ckpt)
2523
# Capture the dense layer weights and reshape them to a 4D tensor which would be
2624
# the weights of a 1x1 convolution layer. This code replaces the dense (FC) layer
2725
# to a 1x1 conv layer.
28-
dense_layer = 'resnet50_v1.5/output/dense/kernel'
2926
dense_layer_value=0.
3027
new_var_list=[]
3128
for var in all_vars:
@@ -34,18 +31,36 @@ def main(args):
3431
dense_layer_value = curr_var
3532
else:
3633
new_var_list.append(tf.Variable(curr_var, name=var[0]))
37-
38-
new_var_value = np.reshape(dense_layer_value, [1, 1, 2048, 1001])
34+
35+
dense_layer_shape = [1, 1, 2048, 1001]
36+
new_var_value = np.reshape(dense_layer_value, )
3937
new_var = tf.Variable(new_var_value, name=dense_layer)
4038
new_var_list.append(new_var)
4139

4240
sess.run(tf.global_variables_initializer())
43-
tf.train.Saver(var_list=new_var_list).save(sess, new_ckpt_path, write_meta_graph=False, write_state=False)
44-
print ("Rewriting checkpoints completed")
41+
tf.train.Saver(var_list=new_var_list).save(sess, output_ckpt_path, write_meta_graph=False, write_state=False)
42+
print ("Rewriting checkpoint completed")
4543

4644
if __name__=='__main__':
4745
parser = argparse.ArgumentParser()
48-
parser.add_argument('--ckpt', type=str, required=True)
49-
parser.add_argument('--out', type=str, default='./new.ckpt')
46+
parser.add_argument('--input', type=str, required=True, help='Path to pretrained RN50 checkpoint with dense layer')
47+
parser.add_argument('--dense_layer', type=str, default='resnet/output/dense/kernel')
48+
parser.add_argument('--output', type=str, default='output_dir', help="Output directory to store new checkpoint")
5049
args = parser.parse_args()
51-
main(args)
50+
main(args)
51+
52+
input_ckpt = args.input
53+
# Create an output directory
54+
os.mkdir(args.output)
55+
56+
new_ckpt='new.ckpt'
57+
new_ckpt_path = os.path.join(args.output, new_ckpt)
58+
with open(os.path.join(output_dir, "checkpoint"), 'w') as file:
59+
file.write("model_checkpoint_path: "+ "\"" + new_ckpt + "\"")
60+
61+
# Process the input checkpoint, apply transforms and generate a new checkpoint.
62+
process_checkpoint(input_ckpt, new_ckpt_path)
63+
64+
65+
66+

‎TensorFlow/Classification/ConvNets/resnet50v1.5/training/QAT/GPU1_RN50_QAT.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
# This script does Quantization aware training of Resnet-50 by finetuning on the pre-trained model using 1 GPU and a batch size of 32.
22
# Usage ./GPU1_RN50_QAT.sh <path to the pre-trained model> <path to dataset> <path to results directory>
33

4-
python main.py --mode=train_and_evaluate --batch_size=32 --lr_warmup_epochs=1 --label_smoothing 0.1 --lr_init=0.00005 --momentum=0.875 --weight_decay=3.0517578125e-05 --finetune_checkpoint=$1 --data_dir=$2 --results_dir=$3 --quantize --symmetric --num_iter 10 --data_format NHWC
4+
python main.py --mode=train_and_evaluate --batch_size=32 --lr_warmup_epochs=1 --quantize --symmetric --use_qdq --label_smoothing 0.1 --lr_init=0.00005 --momentum=0.875 --weight_decay=3.0517578125e-05 --finetune_checkpoint=$1 --data_dir=$2 --results_dir=$3 --num_iter 10 --data_format NHWC

0 commit comments

Comments
 (0)