Как использовать оператор TopK в vectorized_map в тензорном потоке - PullRequest
Новая
       0

Как использовать оператор TopK в vectorized_map в тензорном потоке

0 голосов
/ 19 июня 2020

Я пытался выполнить tf.vectorized_map на тензоре формы [размер партии, размер списка] и применить оператор tf.math.top_k поверх каждая строка в пакете и не удалась.

Например, данные могут быть: 

[ [1,2,4,5,6], [9,5,4,2,1] ]

, и я хотел бы применить topk к [1,2,4,5,6] и [9,5,4,2,1].

Однако мне удалось сделать то же самое с tf.map_fn, но vectorized_map должен работать быстрее. Я использую tensorflow 1.15 . 

import tensorflow as tf
import numpy as np

# create fake data
x = tf.convert_to_tensor([
    [1,2,4,5,6],
    [9,5,4,2,1],
], dtype=tf.float32)
x = tf.reshape(x, (2, -1))

B = x.shape[0] # batchsize
L = x.shape[1] # list size

print(f"B {B}, L {L}")

sess = tf.Session()

print(f"x tensor: {sess.run(x)}\n")

def fv(_x):
    #_tensor = tf.reshape(_x, (L,))  # doesnt work (1)
    _tensor = tf.reshape(tf.convert_to_tensor([9,5,4,2,1], dtype=tf.float32), (L,)) # work (2)
    #_tensor = tf.convert_to_tensor([9,5,4,2,1], dtype=tf.float32) # work (3)

    print(f"_tensor: {_tensor}")
    values, indices = tf.math.top_k(_tensor, k=3)
    # i just need the indices
    return indices

indices = tf.vectorized_map(
        fv,
        x,
)
print("\nindices ")
print(sess.run(indices))

Как мы видим, (2) и (3) выполняются, поэтому оператор topk должен быть пригоден для использования. Также, даже если (1) не работает, я могу использовать _x и, например, просто вернуть его как: 

def fv(_x):
    return _x * 10

Итак, _x можно использовать.

Итак, когда я запускаю код с (1) у меня есть ошибка: 

ValueError: No converter defined for TopKV2
name: "loop_body/TopKV2"
op: "TopKV2"
input: "loop_body/Reshape"
input: "loop_body/TopKV2/k"
attr {
  key: "T"
  value {
    type: DT_FLOAT
  }
}
attr {
  key: "sorted"
  value {
    b: true
  }
}

inputs: [WrappedTensor(t=<tf.Tensor 'loop_body/Reshape/pfor/Reshape:0' shape=(2, 5) dtype=float32>, is_stacked=True, is_sparse_stacked=False), WrappedTensor(t=<tf.Tensor 'loop_body/TopKV2/k:0' shape=() dtype=int32>, is_stacked=False, is_sparse_stacked=False)]. 
Either add a converter or set --op_conversion_fallback_to_while_loop=True, which may run slower

Process finished with exit code 1

Здесь я просто пытаюсь получить индексы, после того, как мне нужно будет обработать вектор, чтобы на выходе было как [[0,0,1,1,1], [1,1,1,0,0] ] для K=3 (1, если значения находятся в топе, иначе 0). А также чтобы задать другой тензор формы [размер пакета, 1], содержащий параметр K для каждой строки. (Мне уже удалось сделать это с помощью map_fn, поэтому я не думаю, что это будет проблемой позже).

Возможно, можно реализовать мой собственный оператор topk в векторизованной карте, но я бы предпочел не

1 Ответ

0 голосов
/ 21 июня 2020

Наконец-то я сделал что-то вроде этого: это не использует vectorized_map, но это то, что я хотел сделать. Но если кто-то сможет заставить его работать с vectorized_map, я бы посмотрел на решение. :) 

def topk(x, k):
    """
    x : shape [B, L]
    k : shape [B, 1]
    return : final_mask of shape [B,L] with final_mask[b,i] = 0 if x[b,i] is in  
    the k[b] biggest values of x[b,:], else final_mask[b,i] = 1
    """
    B = x.shape[0]  # batchsize
    L = x.shape[1]  # list size

    # the indices sorted in descending order
    indices_des = tf.argsort(x, axis=-1, direction='DESCENDING', stable=False, name='sorting_for_topk')

    
    mask = tf.reshape(tf.range(start=0, limit=L, dtype=tf.int32), [1, L])
    mask = tf.repeat(mask, [B], axis=0)
    mask = mask<k

    one_hot = tf.one_hot(indices_des, depth=L) * tf.cast(tf.reshape(mask, [B, L, 1]), tf.float32)
    final_mask = tf.reduce_sum(one_hot, axis=1)
    
    return final_mask
...