Tensorflow: эффективная многочленная выборка (Theano x50 быстрее?)

Question

Я хочу иметь возможность сэмплировать из полиномиального дистрибутива очень эффективно и, очевидно, мой код TensorFlow очень ... очень медленный ...

Идея в том, что у меня есть:

• Вектор: counts = [40, 50, 26, ..., 19] например
• Матрица вероятностей: probs = [[0.1, ..., 0.5], ... [0.3, ..., 0.02]] такая, что np.sum(probs, axis=1) = 1

Допустим, len(counts) = N и len(probs) = (N, 50).Что я хочу сделать, это (в нашем примере):

• выборка 40 раз из первого вектора вероятности матрицы probs
• выборка 50 раз из второго вектора вероятностиматрица probs
• ...
• выборка 19 раз из N-го вектора вероятности матрицы probs

, так что моя окончательная матрица выглядит так(например): A = [[22, ... 13], ..., [12, ..., 3]], где np.sum(A, axis=1) == counts (то есть сумма по каждой строке = число в соответствующей строке counts вектора)

Вот мой пример кода TensorFlow:

import numpy as np
import tensorflow as tf
import tensorflow.contrib.distributions as ds
import time

nb_distribution = 100 # number of probability distributions

counts = np.random.randint(2000, 3500, size=nb_distribution) # define number of counts (vector of size 100 with int in 2000, 3500)
# print(u[:40]) # should be the same as the output of print(np.sum(res, 1)[:40]) in the tf.Session()

# probsn is a matrix of probability:
# each row of probsn contains a vector of size 30 that sums to 1
probsn = np.random.uniform(size=(nb_distribution, 30))
probsn /= np.sum(probsn, axis=1)[:, None]

counts = tf.Variable(counts, dtype=tf.float32)
probs = tf.Variable(tf.convert_to_tensor(probsn.astype(np.float32)))

# sample from the multinomial
dist = ds.Multinomial(total_count=counts, probs=probs)
out = dist.sample()

start = time.time()
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    res = sess.run(out)
    # print(np.sum(res, 1)[:40])
print(time.time() - start)

истекшее время : 0,12 секунды

Мой эквивалентный код в Theano :

import numpy as np
import theano
from theano.tensor import _shared

nb_distribution = 100 # number of probability distributions

counts = np.random.randint(2000, 3500, size=nb_distribution)
#print(u[:40]) # should be the same as the output of print(np.sum(v_sample(), 1)[:40])

counts = _shared(counts) # define number of counts (vector of size 100 with int in 2000, 3500)

# probsn is a matrix of probability:
# each row of probsn contains a vector that sums to 1
probsn = np.random.uniform(size=(nb_distribution, 30)) 
probsn /= np.sum(probsn, axis=1)[:, None]
probsn = _shared(probsn)

from theano.tensor.shared_randomstreams import RandomStreams

np_rng = np.random.RandomState(12345)
theano_rng = RandomStreams(np_rng.randint(2 ** 30))

v_sample = theano.function(inputs=[], outputs=theano_rng.multinomial(n=counts, pvals=probsn))

start_t = time.time()
out = np.sum(v_sample(), 1)[:40]
# print(out)
print(time.time() - start_t)

истекшее время : 0.0025 секунд

Theano примерно в 100 раз быстрее ... Что-то не так с моим кодом TensorFlow?Как мне эффективно произвести выборку из полиномиального распределения в TensorFlow?

Answer 1

Проблема в том, что полиномиальный метод TensorFlow sample() фактически использует вызовы метода _sample_n().Этот метод определен здесь .Как мы видим из кода для выборки из многочлена, код создает матрицу one_hot для каждой строки , а затем сокращает матрицу до вектора путем суммирования по строкам:

math_ops.reduce_sum(array_ops.one_hot(x, depth=k), axis=-2)

Это неэффективно, потому что он использует дополнительную память.Чтобы избежать этого, я использовал функцию tf.scatter_nd.Вот полностью работоспособный пример:

import tensorflow as tf
import numpy as np
import tensorflow.contrib.distributions as ds
import time

tf.reset_default_graph()

nb_distribution = 100 # number of probabilities distribution

u = np.random.randint(2000, 3500, size=nb_distribution) # define number of counts (vector of size 100 with int in 2000, 3500)

# probsn is a matrix of probability:
# each row of probsn contains a vector of size 30 that sums to 1
probsn = np.random.uniform(size=(nb_distribution, 30))
probsn /= np.sum(probsn, axis=1)[:, None]

counts = tf.Variable(u, dtype=tf.float32)
probs = tf.Variable(tf.convert_to_tensor(probsn.astype(np.float32)))

# sample from the multinomial
dist = ds.Multinomial(total_count=counts, probs=probs)
out = dist.sample()


with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    res = sess.run(out) # if remove this line the code is slower...
    start = time.time()
    res = sess.run(out)
    print(time.time() - start)
    print(np.all(u == np.sum(res, axis=1)))

Этот код занял 0,05 секунды для вычисления

def vmultinomial_sampling(counts, pvals, seed=None):
    k = tf.shape(pvals)[1]
    logits = tf.expand_dims(tf.log(pvals), 1)

    def sample_single(args):
        logits_, n_draw_ = args[0], args[1]
        x = tf.multinomial(logits_, n_draw_, seed)
        indices = tf.cast(tf.reshape(x, [-1,1]), tf.int32)
        updates = tf.ones(n_draw_) # tf.shape(indices)[0]
        return tf.scatter_nd(indices, updates, [k])

    x = tf.map_fn(sample_single, [logits, counts], dtype=tf.float32)

    return x

xx = vmultinomial_sampling(u, probsn)
# check = tf.expand_dims(counts, 1) * probs

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    res = sess.run(xx) # if remove this line the code is slower...
    start_t = time.time()
    res = sess.run(xx)
    print(time.time() -start_t)
    #print(np.sum(res, axis=1))
    print(np.all(u == np.sum(res, axis=1)))

Этот код занял 0,016 секунды

Недостаток в том, что мой код нена самом деле не распараллеливать вычисления (хотя параметр parallel_iterations по умолчанию равен 10 в map_fn, установка его в 1 ничего не меняет ...)

Может быть, кто-то найдет что-то лучше, потому чтоон все еще очень медленный по сравнению с реализацией Theano (из-за того, что он не использует преимущества распараллеливания ... и, тем не менее, здесь распараллеливание имеет смысл, поскольку выборка одной строки не зависит от выборки другой ...)