Method RestoreStateFromSnapshot

dali/operators/reader/loader/loader.h:173–205

* @brief Restores the loader's state from a snapshot. */

   * @brief Restores the loader's state from a snapshot.
   */
  void RestoreStateFromSnapshot(const LoaderStateSnapshot& state) {
    DALI_ENFORCE(IsCheckpointingEnabled(),
                 "Checkpointing was not enabled. Please make sure you set"
                 " enable_checkpointing to True when creating the pipeline.");

    RestoreEpochState(state);
    SaveStateSnapshot(current_snapshot_);

    // Now fast forward the loader by `state.age` steps:
    //
    // 1. Run in dry mode to see which samples actually need to be read.
    //    We don't need to read samples that will leave the buffer during fast-forward.
    for (Index i = 0; i < state.age; i++) {
      Index pos_in_batch = (returned_sample_counter_ + i) % max_batch_size_;
      ReadOne(pos_in_batch == 0, pos_in_batch == max_batch_size_ - 1, [](Index i){ return false; });
    }
    auto missing = GetMissingSamples();

    // 2. Restore the state and run reading again, this time reading the needed samples.
    RestoreEpochState(state);
    for (Index i = 0; i < state.age; i++) {
      Index pos_in_batch = (returned_sample_counter_ + i) % max_batch_size_;
      auto filter = [&missing](Index idx) {
        return missing.find(idx) != missing.end();
      };
      ReadOne(pos_in_batch == 0, pos_in_batch == max_batch_size_ - 1, filter);
    }

    DALI_ENFORCE(GetMissingSamples().empty(), "Internal error: reading missing samples failed");

    // current_snapshot_.age was increased by `ReadOne` calls, reset it to correct value
    current_snapshot_.age = state.age;
  }

  bool ShouldPadBatch(bool is_new_batch) {
    // If the reader has depleted samples from the given shard, but shards are not equal