如何在卷积中使用tf.where？

Question

I want to create a graph that splits into several other graphs at after certain point according to classification results. I thiought tf.cond or tf.where might be right to use but im not sure how.

It is impossible to copy here all of my code but I created a small segment that illustrates the issue.

import os
import sys
import tensorflow as tf
GPU_INDEX = 2

net_class = tf.constant([[0.1, 0.2, 0.3, 0.4], [0.4, 0.3, 0.2, 0.1],[0.2, 0.4, 0.3, 0.1], [0.3, 0.2, 0.4, 0.1],[0.1, 0.3, 0.3, 0.4]]) # 3,0,1,2,3
classes = tf.argmax(net_class, axis=1)
cls_0_idx = tf.squeeze(tf.where(tf.equal(classes, 0)))
cls_3_idx = tf.squeeze(tf.where(tf.equal(classes, 3)))

cls_0 = tf.gather(params=net_class, indices=cls_0_idx)
cls_3 = tf.gather(params=net_class, indices=cls_3_idx)

params_0 = tf.constant([1.0,1,1,1])
params_3 = tf.constant([3.0,3,3,3])


output = tf.stack([tf.nn.conv1d(cls_0, params_0, 1,  padding='VALID'), tf.nn.conv1d(cls_3, params_3, 1,  padding='VALID')])

sess = tf.Session()
cls_0_idx_val = sess.run(output)

print(output)

Here I tried to extract the indexes of the input that classifies as 0 or 3 and use different variables to multiply them for the output (shared weights for each class, this is why I use the convolution).

I get the following error :

ValueError: Shape must be rank 4 but is rank 2 for 'conv1d/Conv2D' (op: 'Conv2D') with input shapes: ?, [1,4].

I understand why I get the error ( because tf.where doesn't "know" its size) but the question is how do I fix it? (the classes are not equal and may even be empty in my "real" problem as well)

Answer 1

I think you should

1. setup axis to 1 in tf.squeeze

2. change tf.nn.conv1d to simple multiply

3. change tf.stack to tf.concat

then you will have something like this:

net_class = tf.constant([[0.1, 0.2, 0.3, 0.4], [0.4, 0.3, 0.2, 0.1],[0.2, 0.4, 0.3, 0.1], [0.3, 0.2, 0.4, 0.1],[0.1, 0.3, 0.3, 0.4]]) # 3,0,1,2,3
classes = tf.argmax(net_class, axis=1)
cls_0_idx = tf.squeeze(tf.where(tf.equal(classes, 0)), -1)
cls_3_idx = tf.squeeze(tf.where(tf.equal(classes, 3)), -1)

cls_0 = tf.gather(params=net_class, indices=cls_0_idx)
cls_3 = tf.gather(params=net_class, indices=cls_3_idx)

params_0 = tf.constant([1.0,1,1,1])
params_3 = tf.constant([3.0,3,3,3])
output = tf.concat([cls_0 * params_0, cls_3 * params_3], axis = 0)