Esta página foi traduzida pela API Cloud Translation.
Switch to English

tf.vectorized_map

Versão TensorFlow 1 Ver fonte no GitHub

Mapa paralelo na lista de tensores descompactados dos elems na dimensão 0.

Usado nos cadernos

Usado no guia

Esse método funciona de forma semelhante ao tf.map_fn, mas é otimizado para executar muito mais rápido, possivelmente com um espaço de memória muito maior. Os speedups são obtidos por vetorização (ver https://arxiv.org/pdf/1903.04243.pdf ). A idéia por trás da vetorização é lançar semanticamente todas as invocações de fn em paralelo e fundir as operações correspondentes em todas essas invocações. Essa fusão é feita estaticamente no momento da geração do gráfico e o código gerado geralmente é semelhante em desempenho a uma versão fundida manualmente.

Como tf.vectorized_map paralela totalmente o lote, esse método geralmente será significativamente mais rápido do que usar tf.map_fn , especialmente no modo tf.map_fn . No entanto, este é um recurso experimental e atualmente tem várias limitações:

  • Não deve haver dependência de dados entre as diferentes invocações semânticas de fn , ou seja, deve ser seguro mapear os elementos das entradas em qualquer ordem.
  • Os kernels com estado podem geralmente não ser suportados, pois geralmente implicam uma dependência de dados. No entanto, oferecemos suporte a um conjunto limitado de kernels com estado (como RandomFoo, operações variáveis ​​como leituras, etc).
  • fn possui suporte limitado para operações de controle de fluxo.
  • fn deve retornar a estrutura aninhada de tensores ou operações. No entanto, se uma operação for retornada, ela deve ter zero saídas.
  • A forma e o tipo de qualquer tensor intermediário ou de saída no cálculo de fn não deve depender da entrada em fn .

Exemplos:

 def outer_product(a):
  return tf.tensordot(a, a, 0)

batch_size = 100
a = tf.ones((batch_size, 32, 32))
c = tf.vectorized_map(outer_product, a)
assert c.shape == (batch_size, 32, 32, 32, 32)
 
 # Computing per-example gradients

batch_size = 10
num_features = 32
layer = tf.keras.layers.Dense(1)

def model_fn(arg):
  with tf.GradientTape() as g:
    inp, label = arg
    inp = tf.expand_dims(inp, 0)
    label = tf.expand_dims(label, 0)
    prediction = layer(inp)
    loss = tf.nn.l2_loss(label - prediction)
  return g.gradient(loss, (layer.kernel, layer.bias))

inputs = tf.random.uniform([batch_size, num_features])
labels = tf.random.uniform([batch_size, 1])
per_example_gradients = tf.vectorized_map(model_fn, (inputs, labels))
assert per_example_gradients[0].shape == (batch_size, num_features, 1)
assert per_example_gradients[1].shape == (batch_size, 1)
 

fn A chamada a ser executada. Ele aceita um argumento, que terá a mesma estrutura (possivelmente aninhada) dos elems e retorna uma estrutura possivelmente aninhada de Tensores e Operações, que podem ser diferentes da estrutura dos elems .
elems Um tensor ou sequência (possivelmente aninhada) de tensores, cada um dos quais será desempacotado ao longo de sua primeira dimensão. A sequência aninhada das fatias resultantes será mapeada por fn .
fallback_to_while_loop Se true, ao falhar em vetorizar uma operação, a operação não suportada é agrupada em um tf. while_loop para executar as iterações do mapa. Observe que esse fallback só acontece para operações não suportadas e outras partes de fn ainda são vetorizadas. Se for falso, ao encontrar um op não suportado, um ValueError é lançado. Observe que os fallbacks podem resultar em lentidão, uma vez que a vetorização geralmente produz aceleração de uma a duas ordens de magnitude.

Um tensor ou (possivelmente aninhado) sequência de tensores. Cada tensor compacta os resultados da aplicação de fn aos tensores desembalados dos elementos ao longo da primeira dimensão, do primeiro ao último.

ValueError Se a vetorização falhar e fallback_to_ while_loop for False.