在 Tensorflow 中使用切片赋值量化变量

Quantize the variable with slicing assignment in Tensorflow

我只想对某些变量区域进行量化。 更具体地说,我只想在卷积层的任何通道上执行。

这是二值化代码:

G = tf.get_default_graph()
def binarize(x):
    with G.gradient_override_map({"Sign": "Identity"}):
        E = tf.stop_gradient(tf.reduce_mean(tf.abs(x)))
        return tf.sign(x) * E

我想通过for循环对索引对应的变量区域进行量化。这是我在这里找到的变量赋值():

v = tf.get_variable(name, shape, **kwargs)

for i in range(start_index, end_index):
    if i in index:
        with tf.control_dependencies([v[:,:,i,:].assign(binarize(v[:,:,i,:]))]):
            v = tf.identity(v)

但是,据我所知,变量一旦被赋值,就被视为张量,无法再次赋值。

这是第一次赋值后的错误信息:

ValueError: Sliced assignment is only supported for variables

还有其他方法可以解决这个问题吗?


为了更清楚的问题,我添加了一些打印功能。

v = tf.get_variable(name, shape, **kwargs)

for i in range(start_index, end_index):
    if i in index:
        print(v)
        print(type(v))
        with tf.control_dependencies([v[:,:,i,:].assign(binarize(v[:,:,i,:]))]):
            v = tf.identity(v)
            print(v)
            print(type(v))
            print("Done")

结果:

<tf.Variable 'conv2/weights:0' shape=(5, 5, 96, 256) dtype=float32_ref>
<class 'tensorflow.python.ops.variables.Variable'>
<Tensor("tower0/conv2/Identity:0", shape=(5, 5, 96, 256), dtype=float32, device=/device:CPU:0)
<class 'tensorflow.python.framework.ops.Tensor'>
Done
<Tensor("tower0/conv2/Identity:0", shape=(5, 5, 96, 256), dtype=float32, device=/device:CPU:0)
<class 'tensorflow.python.framework.ops.Tensor'>    

ValueError: Sliced assignment is only supported for variables

我找到了正确的答案,我认为我做了一件很愚蠢的事情。

我通过在二值化函数中创建一个虚拟张量解决了这个问题。虚拟张量是相同形状的张量,在所需索引处为 1,在不需要索引处为 0。通过乘以它,我可以只量化变量的所需区域。

代码如下:

def binarize(x):
    positive_dummy_tensor, negative_dummy_tensor = dummy_tensor(x)
    pos_x = tf.multiply(x, positive_dummy_tensor)
    neg_x = tf.multiply(x, negative_dummy_tensor)
    with G.gradient_override_map({"Sign": "Identity"}):
        E = tf.stop_gradient(tf.reduce_mean(tf.abs(x)))
        return tf.add(tf.multiply(tf.sign(pos_x), E), neg_x)