Merge pull request YunYang1994#416 from KoapT/master

YunYang1994 · web-flow · commit fdefb7dba46b · 2019-11-25T19:22:01.000+08:00
Convert darknet weights to ckpt&amp;pb file.
diff --git a/from_darknet_weights_to_ckpt.py b/from_darknet_weights_to_ckpt.py
@@ -0,0 +1,75 @@
+import tensorflow as tf
+from core.yolov3 import YOLOV3
+
+iput_size = 412
+darknet_weights = '<your darknet weights' path>'
+ckpt_file = './checkpoint/yolov3.ckpt'
+
+def load_weights(var_list, weights_file):
+    """
+    Loads and converts pre-trained weights.
+    :param var_list: list of network variables.
+    :param weights_file: name of the binary file.
+    :return: list of assign ops
+    """
+    with open(weights_file, "rb") as fp:
+        _ = np.fromfile(fp, dtype=np.int32, count=5)
+        weights = np.fromfile(fp, dtype=np.float32)  # np.ndarray
+    print('weights_num:', weights.shape[0])
+    ptr = 0
+    i = 0
+    assign_ops = []
+    while i < len(var_list) - 1:
+        var1 = var_list[i]
+        var2 = var_list[i + 1]
+        # do something only if we process conv layer
+        if 'conv' in var1.name.split('/')[-2]:
+            # check type of next layer
+            if 'batch_normalization' in var2.name.split('/')[-2]:
+                # load batch norm params
+                gamma, beta, mean, var = var_list[i + 1:i + 5]
+                batch_norm_vars = [beta, gamma, mean, var]
+                for vari in batch_norm_vars:
+                    shape = vari.shape.as_list()
+                    num_params = np.prod(shape)
+                    vari_weights = weights[ptr:ptr + num_params].reshape(shape)
+                    ptr += num_params
+                    assign_ops.append(
+                        tf.assign(vari, vari_weights, validate_shape=True))
+                i += 4
+            elif 'conv' in var2.name.split('/')[-2]:
+                # load biases
+                bias = var2
+                bias_shape = bias.shape.as_list()
+                bias_params = np.prod(bias_shape)
+                bias_weights = weights[ptr:ptr +
+                                           bias_params].reshape(bias_shape)
+                ptr += bias_params
+                assign_ops.append(
+                    tf.assign(bias, bias_weights, validate_shape=True))
+                i += 1
+            shape = var1.shape.as_list()
+            num_params = np.prod(shape)
+
+            var_weights = weights[ptr:ptr + num_params].reshape(
+                (shape[3], shape[2], shape[0], shape[1]))
+            # remember to transpose to column-major
+            var_weights = np.transpose(var_weights, (2, 3, 1, 0))
+            ptr += num_params
+            assign_ops.append(
+                tf.assign(var1, var_weights, validate_shape=True))
+            i += 1
+    print('ptr:', ptr)
+    return assign_ops
+    
+with tf.name_scope('input'):
+    input_data = tf.placeholder(dtype=tf.float32,shape=(None, iput_size, iput_size, 3), name='input_data')
+model = YOLOV3(input_data, trainable=False)
+load_ops = load_weights(tf.global_variables(), darknet_weights)
+
+saver = tf.train.Saver(tf.global_variables())
+
+with tf.Session() as sess:
+    sess.run(load_ops)
+    save_path = saver.save(sess, save_path=ckpt_file)
+    print('Model saved in path: {}'.format(save_path))
diff --git a/from_darknet_weights_to_pb.py b/from_darknet_weights_to_pb.py
@@ -0,0 +1,26 @@
+import tensorflow as tf
+from core.yolov3 import YOLOV3
+from from_darknet_weights_to_ckpt import load_weights
+
+input_size = 416
+darknet_weights = '<your darknet weights file path>'
+pb_file = './yolov3.pb'
+output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"]
+
+with tf.name_scope('input'):
+    input_data = tf.placeholder(dtype=tf.float32, shape=(None, input_size, input_size, 3), name='input_data')
+model = YOLOV3(input_data, trainable=False)
+load_ops = load_weights(tf.global_variables(), darknet_weights)
+
+with tf.Session() as sess:
+    sess.run(load_ops)
+    output_graph_def = tf.graph_util.convert_variables_to_constants(
+        sess,
+        tf.get_default_graph().as_graph_def(),
+        output_node_names=output_node_names
+    )
+
+    with tf.gfile.GFile(output_graph, "wb") as f:
+        f.write(output_graph_def.SerializeToString())
+
+    print("{} ops written to {}.".format(len(output_graph_def.node), output_graph))
