Method Compute

tensorflow/core/kernels/save_restore_v2_ops.cc:139–266

  }

  void Compute(OpKernelContext* context) override {
    const Tensor& prefix = context->input(0);
    const Tensor& tensor_names = context->input(1);
    const Tensor& shape_and_slices = context->input(2);
    const int kFixedInputs = 3;  // Prefix, tensor names, shape_and_slices.
    ValidateInputs(true /* is save op */, context, prefix, tensor_names,
                   shape_and_slices, kFixedInputs);
    if (!context->status().ok()) return;

    const int num_tensors = static_cast<int>(tensor_names.NumElements());
    const string& prefix_string = prefix.scalar<tstring>()();
    const auto& tensor_names_flat = tensor_names.flat<tstring>();
    const auto& shape_and_slices_flat = shape_and_slices.flat<tstring>();

    BundleWriter writer(Env::Default(), prefix_string);
    OP_REQUIRES_OK(context, writer.status());
    VLOG(1) << "BundleWriter, prefix_string: " << prefix_string;

    int start_index = 0;
    if (has_ev_) {
      start_index = 1;
    }

    int start_ev_key_index = 0;

    for (int i = start_index; i < num_tensors; ++i) {
      const string& tensor_name = tensor_names_flat(i);
      if (tensor_types_[i] == DT_RESOURCE) {
        auto& handle = HandleFromInput(context, i + kFixedInputs);
        if (IsHandle<EmbeddingVar<int64, float>>(handle)) {
          if (ev_key_types_[start_ev_key_index] == DT_INT32) {
            DumpEvWithGlobalStep<int32, float>(context,
                i + kFixedInputs, tensor_name, writer, tensor_types_[0]);
          } else if (ev_key_types_[start_ev_key_index] == DT_INT64) {
            DumpEvWithGlobalStep<int64, float>(context,
                i + kFixedInputs, tensor_name, writer, tensor_types_[0]);
          }
        } else if (IsHandle<HashTableResource>(handle)) {
          auto handles = context->input(i + kFixedInputs).flat<ResourceHandle>();
          int tensible_size = handles.size() - 1;
          std::vector<core::ScopedUnref> unrefs;
          HashTable* hashtable;
          std::vector<TensibleVariable*> tensibles;

          HashTableResource* htr;
          OP_REQUIRES_OK(context,
              LookupResource(context, handles(0), &htr));
          unrefs.emplace_back(htr);
          hashtable = htr->Internal();

          for (int j = 0; j < tensible_size; j++) {
            TensibleVariableResource* tvr;
            OP_REQUIRES_OK(context,
                LookupResource(context, handles(j + 1), &tvr));
            unrefs.emplace_back(tvr);
            tensibles.push_back(tvr->Internal());
          }