tensorflow index问题-1

记录一下tensorflow中的tensor index的解决方法。

问题描述：对一个矩阵的不同行做max-pooling/softmax。
比如对下面的矩阵按照编号0，1,2,3

0
0
1
1
1
2
3
3

max/avg/sum/product op
- 使用tf.segment分段
- 仅仅支持简单操作，不支持比如分组softmax的操作

tf.segment_max(data, segment_ids, name=None)
data: A Tensor. Must be one of the following types: float32, float64, int32, int64, uint8, int16, int8, uint16, half.
segment_ids: A Tensor. Must be one of the following types: int32, int64. A 1-D tensor whose rank is equal to the rank of data’s first dimension. Values should be sorted and can be repeated.

self.pooled = tf.nn.dropout(poolre, dropout)
self.bag_pool = tf.segment_mean(self.pooled, self.bag_seg)

分组操作前后矩阵维度变化
data:[8,11] + segment_ids([0,0,1,1,1,2,3,3];0/1/2/3是组编号)
-> data:[4,11]

对分组进行softmax

self.select_weigh = tf.exp(tf.gather_nd(weigh, self.bag_param))
self.select_sum = tf.gather(tf.segment_sum(self.select_weigh, self.bag_seg), self.bag_seg)
self.select_softmax = self.select_weigh / self.select_sum
self.poolne = tf.segment_sum(tf.expand_dims(self.select_softmax, 1) * poolre, self.bag_seg)
self.pooled = tf.nn.dropout(self.poolne, dropout)

这里获得和原始向量大小相同的向量。
segment_sum和gather配合使用

self.select_sum = tf.gather(tf.segment_sum(self.select_weigh, self.bag_seg), self.bag_seg)

相除获得softmax

self.select_softmax = self.select_weigh / self.select_sum

通过broadcast乘法得到softmax后的矩阵

self.poolne = tf.segment_sum(tf.expand_dims(self.select_softmax, 1) * poolre, self.bag_seg)

对矩阵先进行列采样再行采样
取元素方式如下

0	×
0	×
1		o
1		o
1		o
2			*
3				@
3				@

下面语句的
bag_param:[[0,0],[1,0],[2,1],[3,1],[4,1],[5,2],[6,3],[7,3]]
就可以得到以上index操作

self.select_weigh = tf.exp(tf.gather_nd(weigh, self.bag_param))