tf算子计算流程统计

tf.depth_to_space() 

import numpy as np
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

if __name__ == "__main__":
	block_size = 3
	i1 = tf.linspace(1.0,288.0,288)
	i2 = tf.reshape(i1,[1,4,4,18])
	i3 = tf.depth_to_space(i2,block_size=block_size)

	with tf.Session() as sess:
		sess.run(tf.global_variables_initializer())
		print i2.eval()
		print i2.shape

		print i3.eval()
		print i3.shape

		i4 = i2.eval()
		i5 = np.array([0]*i3.shape[1] * i3.shape[2] * i3.shape[3])
		i6 = i5.reshape([1,i3.shape[1],i3.shape[2],i3.shape[3]])
		#print i6
		for out_h in range(i3.shape[1]):
			for out_w in range(i3.shape[2]):
				for out_c in range(i3.shape[3]):
					w_in = out_w / block_size
					h_in = out_h / block_size
					c_in = (out_w % block_size) * i3.shape[3] + out_c % i3.shape[3] + (out_h % block_size) * i3.shape[3] * block_size
                                        //onnx crd model
                                        //c_in = c * block_size * block_size + (h % block_size) * block_size + w % block_size
					i6[0][out_h][out_w][out_c] = i4[0][h_in][w_in][c_in]

		print i6 - i3.eval()