Search code examples
variablestensorflowslicequantization

Quantize the variable with slicing assignment in Tensorflow


I want to do quantization for only some area of variables. More specifically, I would like to perform only on any channel in the convolution layer.

This is binarization code:

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

I want to do quantization through the for-loop for a variable area that corresponds to an index. And this is the variable assignment which I found here(How to do slice assignment in Tensorflow):

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)

However, as far as I know, once variable is assigned, it is regarded as tensor so that it can not be assigned again.

This is the error message after first assignment:

ValueError: Sliced assignment is only supported for variables

Is there any other way to solve this problem?


For a more clear question, I've added some print functions.

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")

Result:

<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

Solution

  • I found myself the right answer, and I think I did something very silly.

    I solved this problem by creating a dummy tensor in the binarize function. The dummy tensor is a tensor of the same shape, with 1 at the desired index and 0 at the undesired index. By multiplying this I can quantize only the desired area of the variable.

    Here is the code:

    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)