[经验日记]tensorflow的一种广播机制

[经验日记]tensorflow的一种广播机制

问题:将shape=(2,)维的Tensor A和shape=(2,2,2)的Tensor B以某种方式相乘,这里它们的第一个维度2代表Batch_size, A可以理解为sigmoid输出, B为两张2x2的图像分割结果,要将A中的数分别对应乘上B中的图片(先验二分类)

A = tf.constant([1,2])
B = tf.constant([
    [
        [1,0],
        [0,1]
    ],
    
    [
        [1,1],
        [0,0]
    ]
])

直接A*B的结果是这样的,可以发现其实是A与B的后两维做乘法,不满足我们的需求

A = tf.expand_dims(A,axis=-1)
A = tf.expand_dims(A,axis=-1)

这样在A的后面扩增两维后,A的shape为(2,1,1), 这样A与B维度相同,且第一个维度都是2, A的后两维便广播成2x2与B的后两维匹配,这样就能让对应的数乘到对应的图片上。
（只扩增一次的话还是与第一次结果一样，只是最后两维相乘）