Tensorflow高效重疊添加

Question

我試圖在Tensorflow中實現overlap-add，但我努力將numpy output_seq[start:end] += chunk轉換爲Tensorflow。現在我是output_seq = output_seq + tf.pad(chunk, [[start, length - end]])，但是這對於長序列來說確實很慢。

我也有一種預感，可能有一些技巧可以用來收集/分散，但我無法弄清楚。下面是我的蠻力嘗試：

import tensorflow as tf 

input = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]] 

def overlap_add(overlap): 
    with tf.Graph().as_default(), tf.Session() as sess: 

     x = tf.constant(input) 

     num_chunks = tf.shape(x)[0] 
     chunk_size = tf.shape(x)[1] 
     hop_length = chunk_size - overlap 
     out_len = chunk_size + hop_length * (num_chunks - 1) 

     y = tf.zeros((out_len,), dtype=tf.int32) 

     def body(i, y): 
      j = i * hop_length 
      padding = [[j, out_len - (j + chunk_size)]] 
      chunk = x[i] 
      y = y + tf.pad(chunk, padding) 
      return (i + 1, y) 

     i = tf.constant(0) 
     i, y = tf.while_loop(
      cond=lambda i, _: tf.less(i, num_chunks), 
      body=body, 
      loop_vars=[i, y]) 

     return sess.run(y) 


for i in range(4): 
    print 'overlap_add(%d): %s' % (i, overlap_add(i)) 

# overlap_add(0): [ 1 2 3 4 5 6 7 8 9 10 11 12] 
# overlap_add(1): [ 1 2 3 9 6 7 17 10 11 12] 
# overlap_add(2): [ 1 2 8 10 16 18 11 12] 
# overlap_add(3): [ 1 7 18 21 19 12] 

Answer 1

UPDATE：現在有在Tensorflow本身overlap_and_add功能。

OLD答： 通過文檔拖網，發現unsorted_segment_sum：

import tensorflow as tf 

input = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]] 

def tf_repeat(a, repeats): 
    return tf.reshape(tf.tile(tf.reshape(a, [-1, 1]), 
           [1, repeats]), [-1]) 

def overlap_add(overlap): 
    with tf.Graph().as_default(), tf.Session() as sess: 

     x = tf.constant(input) 
     x_flat = tf.reshape(x, [-1]) 

     num_chunks = tf.shape(x)[0] 
     chunk_size = tf.shape(x)[1] 
     hop_len = chunk_size - overlap 
     flat_len = num_chunks * chunk_size 
     out_len = chunk_size + hop_len * (num_chunks - 1) 

     # e.g. [0,1,2,3, 2,3,4,5, 4,5,6,7] for overlap == 2 
     indexes = tf.range(flat_len) - tf_repeat(tf.range(num_chunks), chunk_size) * overlap 

     return sess.run(tf.unsorted_segment_sum(x_flat, indexes, out_len)) 


for i in range(4): 
    print 'overlap_add(%d): %s' % (i, overlap_add(i)) 

# overlap_add(0): [ 1 2 3 4 5 6 7 8 9 10 11 12] 
# overlap_add(1): [ 1 2 3 9 6 7 17 10 11 12] 
# overlap_add(2): [ 1 2 8 10 16 18 11 12] 
# overlap_add(3): [ 1 7 18 21 19 12] 

Answer 2

可以在Tensorflow使用切片以及：

a[1:3].assign(a[1:3] + b[1:3]).eval() 

出於某種原因assign_add未實現。這對我來說似乎是一個錯誤。

a[1:3].assign_add(b[1:3]).eval() # Doesn't work