Question

如何在用户定义的Tensorflow Op中使用随机数？

我正在用cpp编写一个op，需要在Compute函数中使用随机数。

但是似乎我不应该直接使用cpp随机库，因为它不能由tf.set_random_seed控制。

我当前的代码如下所示，我应该在函数some_interesting_random_function中做什么？

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include <iostream>
#include <typeinfo>
#include <random>

using namespace tensorflow;

REGISTER_OP("MyRandom")
.Output("random: int32")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
        c->set_output(0, c->Scalar());
        return Status::OK();
});

int some_interesting_random_function(){
    return 10;
}

class MyRandomOp : public OpKernel {
    public:
        explicit MyRandomOp(OpKernelConstruction* context) : OpKernel(context) {}
        void Compute(OpKernelContext* context) override {
            Tensor* res;
            TensorShape shape;
            int dims[] = {};
            TensorShapeUtils::MakeShape(dims, 0, &shape);
            OP_REQUIRES_OK(context, context->allocate_output(0, shape,
                        &res));
            auto out1 = res->flat<int32>();
            out1(0) = some_interesting_random_function();
        }
};

REGISTER_KERNEL_BUILDER(Name("MyRandom").Device(DEVICE_CPU), MyRandomOp);

Answer 1

TensorFlow中所有随机数生成的核心是PhiloxRandom，通常通过其包装器GuardedPhiloxRandom访问。如tf.set_random_seed中所述，存在图级和操作级种子，这两种种子可能会或可能不会设置。如果您也想在操作中使用此功能，则需要做几件事。首先，应该使用两个可选属性seed和seed2来声明您的op；请参阅random_ops.cc中的现有操作。然后，在Python中，您有一些包装您的操作的用户API，使用tensorflow.python.framework.random_seed来生成这两个值，您必须将它们导入为tensorflow.python.framework import random_seed，然后执行seed1, seed2 = random_seed.get_seed(seed)；这将使用图形的种子和函数的可选seed参数正确地创建两个种子值（请参见random_ops.py）。显然，这些seed1和seed2值将作为seed和seed2属性传递给您的op。如果您做了所有这些操作，那么GuardedPhiloxRandom将使用正确的种子来正确初始化随机数生成器。

现在，转到内核实现。除了我上面提到的内容外，您还需要结合两件事：在core/kernels/random_op.h中声明的结构模板FillPhiloxRandom，这将帮助您用随机数据填充张量。还有一个Distribution，它只是一个对象，可以用随机数生成器调用该对象以生成一个值（请参见core/lib/random/random_distributions.h中的现有实现）。现在，主要要看在core/kernels/random_op.cc中是如何完成的，然后复制所需的位。其中的大多数内核都基于PhiloxRandomOp（尚未公开声明，但是您可以复制或改编）。这实质上是一个随机数生成器，在输出张量中分配空间（假定第一个输入是所需的形状），然后调用FillPhiloxRandom来完成工作。如果您要尝试创建这种操作（根据某种分布生成一些数据），那么您就准备就绪了！您的代码可能看起来像这样：

// Required for thread pool device
#define EIGEN_USE_THREADS

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/random_op.h"
#include "tensorflow/core/util/guarded_philox_random.h"

// Helper function to convert an 32-bit integer to a float between [0..1).
// Copied from core/lib/random/random_distributions.h
PHILOX_DEVICE_INLINE float Uint32ToFloat(uint32 x) {
  // IEEE754 floats are formatted as follows (MSB first):
  //    sign(1) exponent(8) mantissa(23)
  // Conceptually construct the following:
  //    sign == 0
  //    exponent == 127  -- an excess 127 representation of a zero exponent
  //    mantissa == 23 random bits
  const uint32 man = x & 0x7fffffu;  // 23 bit mantissa
  const uint32 exp = static_cast<uint32>(127);
  const uint32 val = (exp << 23) | man;

  // Assumes that endian-ness is same for float and uint32.
  float result;
  memcpy(&result, &val, sizeof(val));
  return result - 1.0f;
}

// Template class for your custom distribution
template <class Generator, typename RealType>
class MyDistribution;

// Implementation for tf.float32
template <class Generator>
class MyDistribution<Generator, float> {
 public:
  // The number of elements that will be returned (see below).
  static const int kResultElementCount = Generator::kResultElementCount;
  // Cost of generation of a single element (in cycles) (see below).
  static const int kElementCost = 3;
  // Indicate that this distribution may take variable number of samples
  // during the runtime (see below).
  static const bool kVariableSamplesPerOutput = false;
  typedef Array<float, kResultElementCount> ResultType;
  typedef float ResultElementType;

  PHILOX_DEVICE_INLINE
  ResultType operator()(Generator* gen) {
    typename Generator::ResultType sample = (*gen)();
    ResultType result;
    for (int i = 0; i < kResultElementCount; ++i) {
      float r = Uint32ToFloat(sample[i]);
      // Example distribution logic: produce 1 or 0 with 50% probability
      result[i] = 1.0f * (r < 0.5f);
    }
    return result;
  }
};

// Could add implementations for other data types...

// Base kernel
// Copied from core/kernels/random_op.cc
static Status AllocateOutputWithShape(OpKernelContext* ctx, const Tensor& shape,
                                      int index, Tensor** output) {
  TensorShape tensor_shape;
  TF_RETURN_IF_ERROR(ctx->op_kernel().MakeShape(shape, &tensor_shape));
  return ctx->allocate_output(index, tensor_shape, output);
}

template <typename Device, class Distribution>
class PhiloxRandomOp : public OpKernel {
 public:
  typedef typename Distribution::ResultElementType T;
  explicit PhiloxRandomOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, generator_.Init(ctx));
  }
  void Compute(OpKernelContext* ctx) override {
    const Tensor& shape = ctx->input(0);
    Tensor* output;
    OP_REQUIRES_OK(ctx, AllocateOutputWithShape(ctx, shape, 0, &output));
    auto output_flat = output->flat<T>();
    tensorflow::functor::FillPhiloxRandom<Device, Distribution>()(
        ctx, ctx->eigen_device<Device>(),
        // Multiplier 256 is the same as in FillPhiloxRandomTask; do not change
        // it just here.
        generator_.ReserveRandomOutputs(output_flat.size(), 256),
        output_flat.data(), output_flat.size(), Distribution());
  }
 private:
  GuardedPhiloxRandom generator_;
};

// Register kernel
typedef Eigen::ThreadPoolDevice CPUDevice;
template struct functor::FillPhiloxRandom<
    CPUDevice, MyDistribution<tensorflow::random::PhiloxRandom, float>>;
REGISTER_KERNEL_BUILDER(
    Name("MyDistribution")
        .Device(DEVICE_CPU)
        .HostMemory("shape")
        .TypeConstraint<float>("dtype"),
    PhiloxRandomOp<CPUDevice, MyDistribution<tensorflow::random::PhiloxRandom, float>>);

// Register kernels for more types, can use macros as in core/kernels/random_op.cc...

这里还有一些额外的细节。首先，您需要了解PhiloxRandom通常在每个步骤上生成四个无符号的32位整数，并且您必须从中获得随机值。 Uint32ToFloat是从该数字之一获得介于{0和1之间的float的助手。也有一些常量。 kResultElementCount是您的分布在每个步骤中产生的值的数量。如果您通过生成器为每个随机数生成一个值，则也可以Generator::kResultElementCount对其进行设置，例如此处（为4）。但是，例如，如果要生成double值（即tf.float64），则可能需要为每个值使用两个32位整数，因此您可能会在其中生成Generator::kResultElementCount / 2这种情况。 kElementCost应该用来指示您的分布产生一个元素需要多少个周期。我不知道TensorFlow团队是如何衡量的，但这只是在任务之间分配生成工作（由FillPhiloxRandom使用）的提示，因此您可以猜测一下，或者从类似的昂贵版本中复制分配。 kVariableSamplesPerOutput确定对您的发行版的每次调用是否会产生不同数量的输出；同样，当这是false时（这应该是常见的情况），FillPhiloxRandom将使值生成更加有效。 PHILOX_DEVICE_INLINE（在core/lib/random/philox_random.h中定义）是用于内联函数的编译器提示。然后可以为其他数据类型添加其他实现和内核注册，如果支持的话，还可以为DEVICE_GPU GPUDevice（带有typedef Eigen::GpuDevice GPUDevice）甚至DEVICE_SYCL（带有{ {1}}）。与此相关的是，typedef Eigen::SyclDevice SYCLDevice仅用于启用Eigen中的线程池执行设备，以使CPU实现成为多线程。

但是，如果您的用例不同（例如，您想生成一些随机数并进行其他计算），EIGEN_USE_THREADS可能对您没有用（或者可能是，但随后您还需要做其他事情）。看看FillPhiloxRandom和不同类的标题应该可以帮助您弄清楚如何使用它们解决问题。

如何在用户定义的Tensorflow Op中使用随机数？

1 个答案: