Python中的Broadcast机制

A      (4d array):  8 x 1 x 6 x 1
B      (3d array):      7 x 1 x 5
Result (4d array):  8 x 7 x 6 x 5
A      (2d array):  5 x 4
B      (1d array):      1
Result (2d array):  5 x 4
A      (2d array):  15 x 3 x 5
B      (1d array):  15 x 1 x 5
Result (2d array):  15 x 3 x 5

矩阵 (3d array)   : 256 x 256 x 3
向量 (1d array)   :             3
结果 (3d array)   : 256 x 256 x 3

import tensorflow as tf
import numpy as np
if __name__ == "__main__":
	input0_shape = [1,1,3,1]
	input1_shape = [3]
	#维度扩展
	input_len = len(input0_shape) - len(input1_shape)
	for i in range(input_len):
		input1_shape.insert(0,1)
	print input1_shape
	#获取broadcast shape
	broadcast_shape = [0] * len(input0_shape)
	for i in range(len(input0_shape)):
		broadcast_shape[i] = max(input0_shape[i],input1_shape[i])
	print broadcast_shape
	data_a = np.random.random(input0_shape)	#hwcn
	data_b = np.random.random(input1_shape) #h,w,c_out,c_in
	a = tf.placeholder("float")
	b = tf.placeholder("float")
	c = tf.add(a,b)
	with tf.Session() as sess:
		sess.run(tf.global_variables_initializer())
		out = sess.run(c, feed_dict={a: data_a,b:data_b})
		#print data_a
		print data_b
		print out.shape
		#print out - data_a
	res_pre = out - data_b    	#获取input0的扩展结果，用于验证实际值
	out_tf = res_pre.reshape(broadcast_shape[0]*broadcast_shape[1]*broadcast_shape[2]*broadcast_shape[3])
	data_b_tmp = data_a.reshape(input0_shape[0]*input0_shape[1]*input0_shape[2]*input0_shape[3])
	print "out_tf"
	print out_tf
	f_dets = open("pre_data.dat", "w")
	for k in out_tf:
		b = float(k)
		a = '{:.10f}'.format(b)
		f_dets.write(str(a) + '\n')
	f_dets.close()
	out_res = [0]*broadcast_shape[0]*broadcast_shape[1]*broadcast_shape[2]*broadcast_shape[3]
	#进行数据扩展
	for i in range(broadcast_shape[0]):
		for j in range(broadcast_shape[1]):
			for k in range(broadcast_shape[2]):
				for m in range(broadcast_shape[3]):
					tmp_idx0 = i*broadcast_shape[1]*broadcast_shape[2]*broadcast_shape[3]  \
					          + j*broadcast_shape[2]*broadcast_shape[3] + k*broadcast_shape[3] + m
					ii = 0
					jj = 0
					kk = 0
					mm = 0
					if i >= input0_shape[0]:
						ii = input0_shape[0] -1
					else:
						ii = i
					if j >= input0_shape[1]:
						jj = input0_shape[1] -1
					else:
						jj = j
					if k >= input0_shape[2]:
						kk = input0_shape[2] -1
					else:
						kk = k
					if m >= input0_shape[3]:
						mm = input0_shape[3] -1
					else:
						mm = m
					tmp_idx1 = ii*input0_shape[1]*input0_shape[2]*input0_shape[3] \
								+ jj*input0_shape[2]*input0_shape[3] + kk*input0_shape[3] + mm
					#print mm
					out_res[tmp_idx0] = data_b_tmp[tmp_idx1]
	f_dets = open("aft_data.dat", "w")
	for k in out_res:
		b = float(k)
		a = '{:.10f}'.format(b)
		f_dets.write(str(a) + '\n')
	f_dets.close()
	#对比
	print "compare"
	print out_res - out_tf