深度学习将(N,H*W,C)展开为(N,H,W,C)的代码

可以如下实现：

import tensorflow as tf

def flatten(input_tensor):
    input_shape = input_tensor.get_shape().as_list()
    # flatten操作需要先知道input的shape，所以需要调用get_shape()方法
    # input_shape[-1] 表示最后一维的元素个数，即C
    # [-1] 表示不确定C的值，用于后续计算要self-consistent
    dim = input_shape[1] * input_shape[2] * input_shape[3]
    # 计算纵向的元素个数，即N * H * W
    output_tensor = tf.reshape(input_tensor, [-1, dim])
    # -1 表示不确定第一维的元素个数，可以根据input_shape的其他维度计算得到
    return output_tensor

使用方法：

input_tensor = tf.random.normal(shape=(128, 32, 32, 64)) # input shape: (N=128, H=32, W=32, C=64)
output_tensor = flatten(input_tensor) # output shape: (N=128, H*W*C=65536)