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 #include "tensorflow/core/framework/op.h"
     17 #include "tensorflow/core/framework/op_kernel.h"
     18 #include "tensorflow/core/lib/core/stringpiece.h"
     19 #include "tensorflow/core/lib/gtl/map_util.h"
     20 #include "tensorflow/core/lib/random/distribution_sampler.h"
     21 #include "tensorflow/core/lib/random/philox_random.h"
     22 #include "tensorflow/core/lib/random/simple_philox.h"
     23 #include "tensorflow/core/lib/strings/str_util.h"
     24 #include "tensorflow/core/platform/thread_annotations.h"
     25 #include "tensorflow/core/util/guarded_philox_random.h"
     26 
     27 namespace tensorflow {
     28 
     29 // Number of examples to precalculate.
     30 const int kPrecalc = 3000;
     31 // Number of words to read into a sentence before processing.
     32 const int kSentenceSize = 1000;
     33 
     34 namespace {
     35 
     36 bool ScanWord(StringPiece* input, string* word) {
     37   str_util::RemoveLeadingWhitespace(input);
     38   StringPiece tmp;
     39   if (str_util::ConsumeNonWhitespace(input, &tmp)) {
     40     word->assign(tmp.data(), tmp.size());
     41     return true;
     42   } else {
     43     return false;
     44   }
     45 }
     46 
     47 }  // end namespace
     48 
     49 class SkipgramOp : public OpKernel {
     50  public:
     51   explicit SkipgramOp(OpKernelConstruction* ctx)
     52       : OpKernel(ctx), rng_(&philox_) {
     53     string filename;
     54     OP_REQUIRES_OK(ctx, ctx->GetAttr("filename", &filename));
     55     OP_REQUIRES_OK(ctx, ctx->GetAttr("batch_size", &batch_size_));
     56     OP_REQUIRES_OK(ctx, ctx->GetAttr("window_size", &window_size_));
     57     OP_REQUIRES_OK(ctx, ctx->GetAttr("min_count", &min_count_));
     58     OP_REQUIRES_OK(ctx, ctx->GetAttr("subsample", &subsample_));
     59     OP_REQUIRES_OK(ctx, Init(ctx->env(), filename));
     60 
     61     mutex_lock l(mu_);
     62     example_pos_ = corpus_size_;
     63     label_pos_ = corpus_size_;
     64     label_limit_ = corpus_size_;
     65     sentence_index_ = kSentenceSize;
     66     for (int i = 0; i < kPrecalc; ++i) {
     67       NextExample(&precalc_examples_[i].input, &precalc_examples_[i].label);
     68     }
     69   }
     70 
     71   void Compute(OpKernelContext* ctx) override {
     72     Tensor words_per_epoch(DT_INT64, TensorShape({}));
     73     Tensor current_epoch(DT_INT32, TensorShape({}));
     74     Tensor total_words_processed(DT_INT64, TensorShape({}));
     75     Tensor examples(DT_INT32, TensorShape({batch_size_}));
     76     auto Texamples = examples.flat<int32>();
     77     Tensor labels(DT_INT32, TensorShape({batch_size_}));
     78     auto Tlabels = labels.flat<int32>();
     79     {
     80       mutex_lock l(mu_);
     81       for (int i = 0; i < batch_size_; ++i) {
     82         Texamples(i) = precalc_examples_[precalc_index_].input;
     83         Tlabels(i) = precalc_examples_[precalc_index_].label;
     84         precalc_index_++;
     85         if (precalc_index_ >= kPrecalc) {
     86           precalc_index_ = 0;
     87           for (int j = 0; j < kPrecalc; ++j) {
     88             NextExample(&precalc_examples_[j].input,
     89                         &precalc_examples_[j].label);
     90           }
     91         }
     92       }
     93       words_per_epoch.scalar<int64>()() = corpus_size_;
     94       current_epoch.scalar<int32>()() = current_epoch_;
     95       total_words_processed.scalar<int64>()() = total_words_processed_;
     96     }
     97     ctx->set_output(0, word_);
     98     ctx->set_output(1, freq_);
     99     ctx->set_output(2, words_per_epoch);
    100     ctx->set_output(3, current_epoch);
    101     ctx->set_output(4, total_words_processed);
    102     ctx->set_output(5, examples);
    103     ctx->set_output(6, labels);
    104   }
    105 
    106  private:
    107   struct Example {
    108     int32 input;
    109     int32 label;
    110   };
    111 
    112   int32 batch_size_ = 0;
    113   int32 window_size_ = 5;
    114   float subsample_ = 1e-3;
    115   int min_count_ = 5;
    116   int32 vocab_size_ = 0;
    117   Tensor word_;
    118   Tensor freq_;
    119   int64 corpus_size_ = 0;
    120   std::vector<int32> corpus_;
    121   std::vector<Example> precalc_examples_;
    122   int precalc_index_ = 0;
    123   std::vector<int32> sentence_;
    124   int sentence_index_ = 0;
    125 
    126   mutex mu_;
    127   random::PhiloxRandom philox_ GUARDED_BY(mu_);
    128   random::SimplePhilox rng_ GUARDED_BY(mu_);
    129   int32 current_epoch_ GUARDED_BY(mu_) = -1;
    130   int64 total_words_processed_ GUARDED_BY(mu_) = 0;
    131   int32 example_pos_ GUARDED_BY(mu_);
    132   int32 label_pos_ GUARDED_BY(mu_);
    133   int32 label_limit_ GUARDED_BY(mu_);
    134 
    135   // {example_pos_, label_pos_} is the cursor for the next example.
    136   // example_pos_ wraps around at the end of corpus_. For each
    137   // example, we randomly generate [label_pos_, label_limit) for
    138   // labels.
    139   void NextExample(int32* example, int32* label) EXCLUSIVE_LOCKS_REQUIRED(mu_) {
    140     while (true) {
    141       if (label_pos_ >= label_limit_) {
    142         ++total_words_processed_;
    143         ++sentence_index_;
    144         if (sentence_index_ >= kSentenceSize) {
    145           sentence_index_ = 0;
    146           for (int i = 0; i < kSentenceSize; ++i, ++example_pos_) {
    147             if (example_pos_ >= corpus_size_) {
    148               ++current_epoch_;
    149               example_pos_ = 0;
    150             }
    151             if (subsample_ > 0) {
    152               int32 word_freq = freq_.flat<int32>()(corpus_[example_pos_]);
    153               // See Eq. 5 in http://arxiv.org/abs/1310.4546
    154               float keep_prob =
    155                   (std::sqrt(word_freq / (subsample_ * corpus_size_)) + 1) *
    156                   (subsample_ * corpus_size_) / word_freq;
    157               if (rng_.RandFloat() > keep_prob) {
    158                 i--;
    159                 continue;
    160               }
    161             }
    162             sentence_[i] = corpus_[example_pos_];
    163           }
    164         }
    165         const int32 skip = 1 + rng_.Uniform(window_size_);
    166         label_pos_ = std::max<int32>(0, sentence_index_ - skip);
    167         label_limit_ =
    168             std::min<int32>(kSentenceSize, sentence_index_ + skip + 1);
    169       }
    170       if (sentence_index_ != label_pos_) {
    171         break;
    172       }
    173       ++label_pos_;
    174     }
    175     *example = sentence_[sentence_index_];
    176     *label = sentence_[label_pos_++];
    177   }
    178 
    179   Status Init(Env* env, const string& filename) {
    180     string data;
    181     TF_RETURN_IF_ERROR(ReadFileToString(env, filename, &data));
    182     StringPiece input = data;
    183     string w;
    184     corpus_size_ = 0;
    185     std::unordered_map<string, int32> word_freq;
    186     while (ScanWord(&input, &w)) {
    187       ++(word_freq[w]);
    188       ++corpus_size_;
    189     }
    190     if (corpus_size_ < window_size_ * 10) {
    191       return errors::InvalidArgument(
    192           "The text file ", filename,
    193           " contains too little data: ", corpus_size_, " words");
    194     }
    195     typedef std::pair<string, int32> WordFreq;
    196     std::vector<WordFreq> ordered;
    197     for (const auto& p : word_freq) {
    198       if (p.second >= min_count_) ordered.push_back(p);
    199     }
    200     LOG(INFO) << "Data file: " << filename << " contains " << data.size()
    201               << " bytes, " << corpus_size_ << " words, " << word_freq.size()
    202               << " unique words, " << ordered.size()
    203               << " unique frequent words.";
    204     word_freq.clear();
    205     std::sort(ordered.begin(), ordered.end(),
    206               [](const WordFreq& x, const WordFreq& y) {
    207                 return x.second > y.second;
    208               });
    209     vocab_size_ = static_cast<int32>(1 + ordered.size());
    210     Tensor word(DT_STRING, TensorShape({vocab_size_}));
    211     Tensor freq(DT_INT32, TensorShape({vocab_size_}));
    212     word.flat<string>()(0) = "UNK";
    213     static const int32 kUnkId = 0;
    214     std::unordered_map<string, int32> word_id;
    215     int64 total_counted = 0;
    216     for (std::size_t i = 0; i < ordered.size(); ++i) {
    217       const auto& w = ordered[i].first;
    218       auto id = i + 1;
    219       word.flat<string>()(id) = w;
    220       auto word_count = ordered[i].second;
    221       freq.flat<int32>()(id) = word_count;
    222       total_counted += word_count;
    223       word_id[w] = id;
    224     }
    225     freq.flat<int32>()(kUnkId) = corpus_size_ - total_counted;
    226     word_ = word;
    227     freq_ = freq;
    228     corpus_.reserve(corpus_size_);
    229     input = data;
    230     while (ScanWord(&input, &w)) {
    231       corpus_.push_back(gtl::FindWithDefault(word_id, w, kUnkId));
    232     }
    233     precalc_examples_.resize(kPrecalc);
    234     sentence_.resize(kSentenceSize);
    235     return Status::OK();
    236   }
    237 };
    238 
    239 REGISTER_KERNEL_BUILDER(Name("Skipgram").Device(DEVICE_CPU), SkipgramOp);
    240 
    241 class NegTrainOp : public OpKernel {
    242  public:
    243   explicit NegTrainOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    244     base_.Init(0, 0);
    245 
    246     OP_REQUIRES_OK(ctx, ctx->GetAttr("num_negative_samples", &num_samples_));
    247 
    248     std::vector<int32> vocab_count;
    249     OP_REQUIRES_OK(ctx, ctx->GetAttr("vocab_count", &vocab_count));
    250 
    251     std::vector<float> vocab_weights;
    252     vocab_weights.reserve(vocab_count.size());
    253     for (const auto& f : vocab_count) {
    254       float r = std::pow(static_cast<float>(f), 0.75f);
    255       vocab_weights.push_back(r);
    256     }
    257     sampler_ = new random::DistributionSampler(vocab_weights);
    258   }
    259 
    260   ~NegTrainOp() override { delete sampler_; }
    261 
    262   void Compute(OpKernelContext* ctx) override {
    263     Tensor w_in = ctx->mutable_input(0, false);
    264     OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(w_in.shape()),
    265                 errors::InvalidArgument("Must be a matrix"));
    266     Tensor w_out = ctx->mutable_input(1, false);
    267     OP_REQUIRES(ctx, w_in.shape() == w_out.shape(),
    268                 errors::InvalidArgument("w_in.shape == w_out.shape"));
    269     const Tensor& examples = ctx->input(2);
    270     OP_REQUIRES(ctx, TensorShapeUtils::IsVector(examples.shape()),
    271                 errors::InvalidArgument("Must be a vector"));
    272     const Tensor& labels = ctx->input(3);
    273     OP_REQUIRES(ctx, examples.shape() == labels.shape(),
    274                 errors::InvalidArgument("examples.shape == labels.shape"));
    275     const Tensor& learning_rate = ctx->input(4);
    276     OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(learning_rate.shape()),
    277                 errors::InvalidArgument("Must be a scalar"));
    278 
    279     auto Tw_in = w_in.matrix<float>();
    280     auto Tw_out = w_out.matrix<float>();
    281     auto Texamples = examples.flat<int32>();
    282     auto Tlabels = labels.flat<int32>();
    283     auto lr = learning_rate.scalar<float>()();
    284     const int64 vocab_size = w_in.dim_size(0);
    285     const int64 dims = w_in.dim_size(1);
    286     const int64 batch_size = examples.dim_size(0);
    287     OP_REQUIRES(ctx, vocab_size == sampler_->num(),
    288                 errors::InvalidArgument("vocab_size mismatches: ", vocab_size,
    289                                         " vs. ", sampler_->num()));
    290 
    291     // Gradient accumulator for v_in.
    292     Tensor buf(DT_FLOAT, TensorShape({dims}));
    293     auto Tbuf = buf.flat<float>();
    294 
    295     // Scalar buffer to hold sigmoid(+/- dot).
    296     Tensor g_buf(DT_FLOAT, TensorShape({}));
    297     auto g = g_buf.scalar<float>();
    298 
    299     // The following loop needs 2 random 32-bit values per negative
    300     // sample.  We reserve 8 values per sample just in case the
    301     // underlying implementation changes.
    302     auto rnd = base_.ReserveSamples32(batch_size * num_samples_ * 8);
    303     random::SimplePhilox srnd(&rnd);
    304 
    305     for (int64 i = 0; i < batch_size; ++i) {
    306       const int32 example = Texamples(i);
    307       DCHECK(0 <= example && example < vocab_size) << example;
    308       const int32 label = Tlabels(i);
    309       DCHECK(0 <= label && label < vocab_size) << label;
    310       auto v_in = Tw_in.chip<0>(example);
    311 
    312       // Positive: example predicts label.
    313       //   forward: x = v_in' * v_out
    314       //            l = log(sigmoid(x))
    315       //   backward: dl/dx = g = sigmoid(-x)
    316       //             dl/d(v_in) = g * v_out'
    317       //             dl/d(v_out) = v_in' * g
    318       {
    319         auto v_out = Tw_out.chip<0>(label);
    320         auto dot = (v_in * v_out).sum();
    321         g = (dot.exp() + 1.f).inverse();
    322         Tbuf = v_out * (g() * lr);
    323         v_out += v_in * (g() * lr);
    324       }
    325 
    326       // Negative samples:
    327       //   forward: x = v_in' * v_sample
    328       //            l = log(sigmoid(-x))
    329       //   backward: dl/dx = g = -sigmoid(x)
    330       //             dl/d(v_in) = g * v_out'
    331       //             dl/d(v_out) = v_in' * g
    332       for (int j = 0; j < num_samples_; ++j) {
    333         const int sample = sampler_->Sample(&srnd);
    334         if (sample == label) continue;  // Skip.
    335         auto v_sample = Tw_out.chip<0>(sample);
    336         auto dot = (v_in * v_sample).sum();
    337         g = -((-dot).exp() + 1.f).inverse();
    338         Tbuf += v_sample * (g() * lr);
    339         v_sample += v_in * (g() * lr);
    340       }
    341 
    342       // Applies the gradient on v_in.
    343       v_in += Tbuf;
    344     }
    345   }
    346 
    347  private:
    348   int32 num_samples_ = 0;
    349   random::DistributionSampler* sampler_ = nullptr;
    350   GuardedPhiloxRandom base_;
    351 };
    352 
    353 REGISTER_KERNEL_BUILDER(Name("NegTrain").Device(DEVICE_CPU), NegTrainOp);
    354 
    355 }  // end namespace tensorflow
    356