Home | History | Annotate | Download | only in kernels 
      1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 // See docs in ../ops/random_ops.cc.
     17 
     18 #include <vector>
     19 #include "tensorflow/core/framework/op_kernel.h"
     20 #include "tensorflow/core/framework/register_types.h"
     21 #include "tensorflow/core/framework/tensor.h"
     22 #include "tensorflow/core/framework/tensor_shape.h"
     23 #include "tensorflow/core/framework/tensor_util.h"
     24 #include "tensorflow/core/lib/random/random_distributions.h"
     25 #include "tensorflow/core/platform/logging.h"
     26 #include "tensorflow/core/util/guarded_philox_random.h"
     27 
     28 namespace tensorflow {
     29 
     30 // TODO(irving): If performance is critical, generate output directly instead
     31 // of an in-place shuffle using a pseudorandom permutation like
     32 //
     33 //   https://github.com/otherlab/geode/blob/master/geode/random/permute.cpp
     34 //
     35 // This is probably also the right thing if we want a GPU version of shuffling.
     36 
     37 // We use our own version of std::random_shuffle to guarantee that exactly
     38 // size - 1 samples are used.
     39 template <class Iter, class Random>
     40 static inline void RandomShuffle(Iter first, Iter last, Random& uniform) {
     41   if (first == last) return;
     42   const auto stop = last - 1;
     43   for (auto i = first; i != stop; ++i) {
     44     using std::iter_swap;
     45     iter_swap(i, i + uniform(last - i));
     46   }
     47 }
     48 
     49 template <class IntT, class InT, class OutT, class Random>
     50 static void IndexedShuffle(const int64 size, const InT& input_mat,
     51                            OutT output_mat, Random& uniform) {
     52   std::vector<IntT> permutation(size);
     53   for (IntT i = 0; i < size; i++) {
     54     permutation[i] = i;
     55   }
     56   RandomShuffle(permutation.begin(), permutation.end(), uniform);
     57   for (IntT i = 0; i < size; i++) {
     58     output_mat.template chip<0>(i) = input_mat.template chip<0>(permutation[i]);
     59   }
     60 }
     61 
     62 template <typename T>
     63 class RandomShuffleOp : public OpKernel {
     64  public:
     65   explicit RandomShuffleOp(OpKernelConstruction* context) : OpKernel(context) {
     66     OP_REQUIRES_OK(context, generator_.Init(context));
     67   }
     68 
     69   void Compute(OpKernelContext* context) override {
     70     const Tensor& input = context->input(0);
     71 
     72     if (input.NumElements() <= 1 || input.dim_size(0) <= 1) {
     73       // No shuffling is required, so copy input directly to output
     74       context->set_output(0, input);
     75     } else {
     76       // Reserve enough random samples for shuffling
     77       const int64 size = input.dim_size(0);
     78       const int64 samples = size - 1;
     79       auto local_gen = generator_.ReserveSamples32(samples);
     80       random::SingleSampleAdapter<random::PhiloxRandom> single(&local_gen);
     81       const auto uniform = [&single](uint32 n) { return single() % n; };
     82 
     83       if (input.dims() == 1) {
     84         // For 1D data, copy and then shuffle in place
     85         context->set_output(0, tensor::DeepCopy(input));
     86         auto vec = context->mutable_output(0)->vec<T>();
     87         RandomShuffle(vec.data(), vec.data() + size, uniform);
     88       } else {
     89         // For >= 2D, shuffle indices and then copy across
     90         Tensor* output = nullptr;
     91         OP_REQUIRES_OK(context,
     92                        context->allocate_output(0, input.shape(), &output));
     93         const auto input_mat = input.flat_outer_dims<T>();
     94         auto output_mat = output->flat_outer_dims<T>();
     95         if (size < kint32max) {
     96           IndexedShuffle<int32>(size, input_mat, output_mat, uniform);
     97         } else {
     98           IndexedShuffle<int64>(size, input_mat, output_mat, uniform);
     99         }
    100       }
    101     }
    102   }
    103 
    104  private:
    105   GuardedPhiloxRandom generator_;
    106 };
    107 
    108 #define REGISTER(T)                                                    \
    109   REGISTER_KERNEL_BUILDER(                                             \
    110       Name("RandomShuffle").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
    111       RandomShuffleOp<T>);
    112 TF_CALL_ALL_TYPES(REGISTER)
    113 
    114 }  // namespace tensorflow
    115