python - 如何在 tensorflow 中制作 reshape 层？

Question

当我有形状 [2,2,2] 的数据时，例如:

a = np.array([[(1,2), (3,4)],
              [(5,6), (7,8)]
              ])

我希望该层输出 [2,2]，例如:

b = np.array([[1,0],
              [0,1]])

如何构建图层？我当前的设置返回 [2,2,1] 的形状，并且我似乎无法在图层的单位变量中指定尺寸:

tf_x = tf.placeholder(tf.float32, [None, 2, 2]) 

output = tf.layers.dense(tf_x, 1, tf.nn.relu) 

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())    
    pred = sess.run(output, {tf_x: a})

Answer 1

My current setup returns a shape of [2,2,1] and I cannot seem to be able to specify the dimensions in the units variable of the layer:

a = a.reshape(1, -1)

tf_x = tf.placeholder(tf.float32, [None, a.shape[-1]]) 

# now if you want the final array to have total 4 element, you can set it as number of output
output = tf.layers.dense(tf_x, 4, tf.nn.relu) 

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())    
    pred = sess.run(output, {tf_x: a})
pred = pred.reshape(2, 2)