NCHW和NHWC网络格式转换Tensorflow模型

时间:2019-08-30 01:40:11

标签: python tensorflow

对于此Humanpose Tensorflow网络network_cmubase,它仅接受NHWC输入格式。 如果我以NCHW格式构建网络,则错误为

Depth of input (32) is not a multiple of input depth of filter (3) for 'conv1_1/Conv2D' (op: 'Conv2D') with input shapes: [1,3,24,32], [3,3,3,64]. 

我构建网络的代码是

import tensorflow as tf
import numpy as np
from network_cmu import CmuNetwork


def main():
    #print(tensor_util.MakeNdarray(n.attr['value'].tensor))
    placeholder_input = tf.placeholder(dtype=tf.float32, shape=(1, 3, 24, 32), name="image") 
    net = CmuNetwork({'image': placeholder_input}, trainable=False)
    # Add an op to initialize the variables.
    init_op = tf.global_variables_initializer()
    saver = tf.train.Saver()
    init_op = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init_op)
        #for n in tf.get_default_graph().as_graph_def().node:
        #   print(n.name)
        save_path = saver.save(sess, "cmuThreeOutputs/model.ckpt") 

if __name__ == '__main__':
   main()

我应该更改为具有NCHW格式的网络吗?

1 个答案:

答案 0 :(得分:1)

您可以使用tf.transpose将轴从NHWC移到NCHW

input_ = tf.convert_to_tensor(np.random.rand(1, 3, 24, 32))
a1 = tf.transpose(input_, perm=[0, 2, 3, 1])
print(a1.shape)  # 1, 24, 32, 3

您甚至可以使用tf.reshape

a2 = tf.reshape(input_, (-1, input_.shape[2], input_.shape[3], input_.shape[1]))
print(a2.shape)  # 1, 24, 32, 3