summaryrefslogtreecommitdiff
path: root/tests/tools/tflite_run/src/tensor_loader.cc
blob: de605bacfe522525bdd93d2bd7e26c6bf38caae7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#include "tensor_loader.h"

#include <assert.h>

#include <fstream>

#include "misc/tensor/Shape.h"

namespace TFLiteRun
{

TensorLoader::TensorLoader(tflite::Interpreter &interpreter)
    : _interpreter(interpreter), _raw_data(nullptr)
{
}

void TensorLoader::loadDumpedTensors(const std::string &filename)
{
  // TODO Handle file open/read error
  std::ifstream file(filename, std::ios::ate | std::ios::binary);
  size_t file_size = file.tellg();
  file.seekg(0, std::ios::beg);

  uint32_t num_tensors = 0;
  file.read(reinterpret_cast<char *>(&num_tensors), sizeof(num_tensors));

  int tensor_indices_raw[num_tensors];
  file.read(reinterpret_cast<char *>(tensor_indices_raw), sizeof(tensor_indices_raw));
  std::vector<int> tensor_indices(tensor_indices_raw, tensor_indices_raw + num_tensors);

  _raw_data = std::unique_ptr<float[]>(new float[file_size]);
  file.read(reinterpret_cast<char *>(_raw_data.get()), file_size);
  file.close();

  size_t read_bytes = loadTensorsFromRawData(tensor_indices);

  // The file size and total output tensor size must match
  assert(file_size ==
         sizeof(num_tensors) + sizeof(tensor_indices_raw) + read_bytes * sizeof(float));
}

void TensorLoader::loadRawTensors(const std::string &filename,
                                  const std::vector<int> &tensor_indices)
{
  // TODO Handle file open/read error
  std::ifstream file(filename, std::ios::ate | std::ios::binary);
  size_t file_size = file.tellg();
  file.seekg(0, std::ios::beg);

  _raw_data = std::unique_ptr<float[]>(new float[file_size]);
  file.read(reinterpret_cast<char *>(_raw_data.get()), file_size);
  file.close();

  size_t read_bytes = loadTensorsFromRawData(tensor_indices);

  // The file size and total output tensor size must match
  assert(file_size == read_bytes * sizeof(float));
}

size_t TensorLoader::loadTensorsFromRawData(const std::vector<int> &tensor_indices)
{
  size_t offset = 0;
  for (const auto &o : tensor_indices)
  {
    const TfLiteTensor *tensor = _interpreter.tensor(o);

    // Convert tensor shape to `Shape` from `tensor->dims`
    nnfw::misc::tensor::Shape shape(static_cast<size_t>(tensor->dims->size));
    for (int d = 0; d < tensor->dims->size; d++)
    {
      shape.dim(d) = tensor->dims->data[d];
    }

    float *base = _raw_data.get() + offset;

    assert(tensor->bytes % sizeof(float) == 0);
    offset += (tensor->bytes / sizeof(float));

    _tensor_map.insert(std::make_pair(o, nnfw::tflite::TensorView<float>(shape, base)));
  }

  return offset;
}

const nnfw::tflite::TensorView<float> &TensorLoader::get(int tensor_idx) const
{
  auto found = _tensor_map.find(tensor_idx);
  assert(found != _tensor_map.end());
  return found->second;
}

} // end of namespace TFLiteRun