bool AVInputOp::Prefetch()

in ops/av_input_op.h [657:739]
62 lines of code
9 McCabe index (conditional complexity)

bool AVInputOp<Context>::Prefetch() {
  try {
    // We will get the reader pointer from input.
    // If we use local clips, db will store the list
    reader_ = &OperatorBase::Input<db::DBReader>(0);

    // Call mutable_data() once to allocate the underlying memory.
    prefetched_clip_rgb_.mutable_data<float>();
    prefetched_clip_of_.mutable_data<float>();
    prefetched_clip_logmels_.mutable_data<float>();
    prefetched_label_.mutable_data<int>();
    prefetched_video_id_.mutable_data<int64_t>();

    // Prefetching handled with a thread pool of "decode_threads" threads.
    std::mt19937 meta_randgen(time(nullptr));
    std::vector<std::mt19937> randgen_per_thread;
    for (int i = 0; i < num_decode_threads_; ++i) {
      randgen_per_thread.emplace_back(meta_randgen());
    }

    std::bernoulli_distribution mirror_this_clip(0.5);
    for (int item_id = 0; item_id < batch_size_; ++item_id) {
      std::mt19937* randgen =
          &randgen_per_thread[item_id % num_decode_threads_];

      int frame_size = crop_size_ * crop_size_;
      // get the clip data pointer for the item_id -th example
      float* clip_rgb_data = prefetched_clip_rgb_.mutable_data<float>() +
          frame_size * length_rgb_ * channels_rgb_ * item_id * clip_per_video_;

      // get the logmels data for the current clip
      float* clip_logmels_data =
          prefetched_clip_logmels_.mutable_data<float>() +
              item_id * logMelFilters_ * logMelFrames_ * clip_per_video_;

      // get the label data pointer for the item_id -th example
      int* label_data = prefetched_label_.mutable_data<int>() +
          (do_multi_label_ ? num_of_class_ : 1) * item_id * clip_per_video_;

      // get the video id data pointer for the item_id -th example
      int64_t* video_id_data =
          prefetched_video_id_.mutable_data<int64_t>() + item_id * clip_per_video_;

      std::string key, value;
      // read data
      reader_->Read(&key, &value);

      thread_pool_->run(std::bind(
          &AVInputOp<Context>::DecodeAndTransform,
          this,
          std::string(value),
          clip_rgb_data,
          clip_logmels_data,
          label_data,
          video_id_data,
          randgen,
          &mirror_this_clip));
    } // for over the batch
    thread_pool_->waitWorkComplete();

    // If the context is not CPUContext, we will need to do a copy in the
    // prefetch function as well.
    if (!std::is_same<Context, CPUContext>::value) {
      if (get_rgb_) {
        prefetched_clip_rgb_on_device_.CopyFrom(
            prefetched_clip_rgb_, &context_);
      }
      if (get_logmels_) {
        prefetched_clip_logmels_on_device_.CopyFrom(
            prefetched_clip_logmels_, &context_);
      }
      prefetched_label_on_device_.CopyFrom(prefetched_label_, &context_);
      if (get_video_id_) {
        prefetched_video_id_on_device_.CopyFrom(
            prefetched_video_id_, &context_);
      }
    }
  } catch (const std::exception& exc) {
    std::cerr << "While calling Prefetch()\n";
    std::cerr << exc.what();
  }
  return true;
}