pdf_sample_layer.py

import tensorflow as tf
import numpy as np


def pdf_sample(pdf, uniform_noise):
  pdf = pdf / (tf.reduce_sum(pdf, axis=1, keep_dims=True) + 1e-36)
  cdf = tf.cumsum(pdf, axis=1, exclusive=True)
  indices = tf.reduce_sum(
      tf.cast(tf.less(cdf, uniform_noise), tf.int32), axis=1) - 1
  return indices


def pdf_sample_2d(pdf, uniform_noise):
  height, width = pdf.get_shape()[1], pdf.get_shape()[2]
  pdf = tf.reshape(pdf, (int(pdf.get_shape()[0]), -1))
  indices_1d = pdf_sample(pdf, uniform_noise)
  indices = tf.stack(
      [tf.clip_by_value(indices_1d / width, 0, height - 1), indices_1d % width],
      axis=1)
  return indices


def test1():
  import cv2
  batch_size = 1024
  img = cv2.imread('data/doggy.jpg').mean(axis=2)

  pdf_batch = np.empty(
      shape=(batch_size, img.shape[0], img.shape[1]), dtype=np.float32)

  for i in range(batch_size):
    pdf_batch[i] = img

  pdf = tf.placeholder(tf.float32, (batch_size, img.shape[0], img.shape[1]))
  noise = tf.placeholder(tf.float32, (batch_size, 1))

  with tf.Session() as sess:
    indices = pdf_sample_2d(pdf, noise)
    image_buffer = np.zeros(
        shape=(img.shape[0], img.shape[1]), dtype=np.float32)

    while True:
      indices_out = sess.run(
          indices,
          feed_dict={pdf: pdf_batch,
                     noise: np.random.rand(batch_size, 1)})

      for ind in indices_out:
        image_buffer[ind[0]][ind[1]] += 1

      cv2.imshow('img', image_buffer / np.max(image_buffer))
      cv2.waitKey(30)


def test2():
  batch_size = 1024
  n = 3

  pdf_batch = [[2.0**i for i in range(1, n + 1)] for _ in range(batch_size)]

  pdf = tf.placeholder(tf.float32, (batch_size, n))
  noise = tf.placeholder(tf.float32, (batch_size, 1))

  counter = [0 for _ in range(n)]

  with tf.Session() as sess:
    indices = pdf_sample(pdf_batch, noise)

    for i in range(1000):
      indices_out = sess.run(
          indices,
          feed_dict={pdf: pdf_batch,
                     noise: np.random.rand(batch_size, 1)})
      for i in indices_out:
        counter[indices_out[i]] += 1

    for i in range(n):
      print(counter[i] * 1.0 / 100 / batch_size)


if __name__ == '__main__':
  test2()
  # test()