TensorFlow如何沿指定轴对特定切片执行乘法并保留其余元素不变

实现方法 #

这里推荐两种常用的实现方案，都可以满足你的需求：

方案1：广播系数乘法（最简洁，推荐）

利用TensorFlow的广播机制，构造一个对应待修改维度的系数向量，需要做乘法的位置填b，其余位置填1，直接和原张量相乘即可，全程不需要修改张量结构，性能最优：

import tensorflow as tf

a = tf.ones([2,2,3])
b = tf.constant(7)

# 构造最后一维的系数：第0、2位保持原值所以乘1，第1位乘b
coeff = tf.constant([1, b, 1], dtype=a.dtype)
# 自动广播到和a同形状相乘
c = a * coeff

方案2：切片更新（适合更复杂的修改场景）

如果需要修改的切片规则更复杂，可以通过Variable的切片赋值实现：

import tensorflow as tf

a = tf.ones([2,2,3])
b = tf.constant(7)

# 基于原张量初始化Variable
c = tf.Variable(a)
# 仅对指定切片执行乘法后赋值
c[:, :, 1].assign(c[:, :, 1] * b)
# 如需转为普通张量可调用 c = tf.convert_to_tensor(c)

两种方案输出的c都完全符合你要求的形状和数值规则。

内容的提问来源于stack exchange，提问作者user15110545