Status SetInputAndInvokeMiniBatch()

in tensorflow_serving/servables/tensorflow/tflite_session.cc [183:250]
65 lines of code
14 McCabe index (conditional complexity)

Status SetInputAndInvokeMiniBatch(
    std::unique_ptr<internal::TfLiteInterpreterWrapper>& interpreter_wrapper,
    const std::vector<int>& tflite_input_indices,
    const std::vector<std::vector<const Tensor*>>& inputs, int batch_size,
    int* fixed_batch_size) {
  auto* interpreter = interpreter_wrapper->Get();
  // Load input data from Tensorflow tensors.
  for (int i = 0; i < tflite_input_indices.size(); ++i) {
    int tflite_input_idx = tflite_input_indices[i];
    auto tflite_input_tensor = interpreter->tensor(tflite_input_idx);
    const auto& tf_input_tensors = inputs[i];
    if (tflite_input_tensor->type != kTfLiteString) {
      const Tensor* tf_input_tensor = tf_input_tensors[0];
      if (tf_input_tensors.size() > 1) {
        Tensor concated;
        std::vector<Tensor> to_concatenate;
        to_concatenate.reserve(tf_input_tensors.size());
        for (const auto* t : tf_input_tensors) {
          to_concatenate.push_back(std::move(*t));
        }
        TF_RETURN_IF_ERROR(tensor::Concat(to_concatenate, &concated));
        tf_input_tensor = &concated;
      }
      auto tensor_bytes = tf_input_tensor->tensor_data();
      std::vector<int> tf_dims = TensorDims(*tf_input_tensor);
      std::vector<int> tflite_dims(
          tflite_input_tensor->dims->data,
          tflite_input_tensor->dims->data + tflite_input_tensor->dims->size);
      if (tensor_bytes.size() != tflite_input_tensor->bytes ||
          tf_dims != tflite_dims) {
        if (interpreter->ResizeInputTensor(tflite_input_idx, tf_dims) !=
            kTfLiteOk) {
          return errors::Internal(
              "Failed to resize input tensor: ", tflite_input_tensor->name,
              " from ", tflite_input_tensor->bytes, " to ", tensor_bytes.size(),
              " bytes.");
        }
        if (interpreter->AllocateTensors() != kTfLiteOk) {
          return errors::Internal("Failed to allocate tensors");
        }
      }
      std::memcpy(tflite_input_tensor->data.raw, tensor_bytes.data(),
                  tensor_bytes.size());
    } else {
      // Copy the string tensor data to the input tflite tensor.
      const bool needs_resize =
          fixed_batch_size ? batch_size > interpreter_wrapper->GetBatchSize()
                           : batch_size != interpreter_wrapper->GetBatchSize();
      if (needs_resize) {
        // std::cout << "resizing to: " << batch_size << std::endl;
        interpreter->ResizeInputTensor(tflite_input_idx, {batch_size});
        interpreter_wrapper->SetBatchSize(batch_size);
        if (interpreter->AllocateTensors() != kTfLiteOk) {
          return errors::Internal("Failed to allocate tensors");
        }
      }
      if (fixed_batch_size) {
        *fixed_batch_size = interpreter_wrapper->GetBatchSize();
      }
      TF_RETURN_IF_ERROR(interpreter_wrapper->SetStringData(
          tf_input_tensors, tflite_input_tensor, tflite_input_idx, batch_size));
    }
  }
  if (interpreter_wrapper->Invoke() != kTfLiteOk) {
    return errors::Internal("Failed to invoke TfLite interpreter");
  }
  return Status::OK();
}