Function TEST

tensorflow/core/kernels/embedding_variable_memory_test.cc:45–80

}

TEST(EmbeddingVariabelMemoryTest, TestMemory) {
  int value_size = 32;
  int64 default_value_dim = 4096;
  int filter_freq = 2;
  Tensor default_value(
      DT_FLOAT, TensorShape({default_value_dim, value_size}));
  auto default_value_matrix = default_value.matrix<float>();
	for (int i = 0; i < default_value_dim; i++) {
		for (int j = 0 ; j < value_size; j++) {
			default_value_matrix(i, j) = i * value_size + j;
		}
	}

  int num_of_ids = 1000000;
  std::vector<int64> id_list(num_of_ids);
  for (int i = 0; i < num_of_ids; i++) {
    id_list[i] = i;
  }
  float used_mb = PerfMemory(default_value, id_list,
                             value_size, default_value_dim);
  float theoritical_mb =
      50 + num_of_ids * (value_size * sizeof(float)) / 1000000;
  EXPECT_TRUE((used_mb > theoritical_mb * 0.99) &&
              (used_mb < theoritical_mb * 1.01));

  for (int i = 0; i < num_of_ids / 2; i++) {
    id_list.emplace_back(i);
  }
  used_mb = PerfMemory(default_value, id_list, value_size,
                       default_value_dim, filter_freq);
  theoritical_mb =
      50 + num_of_ids * (8 + value_size * sizeof(float) / 2
                         + 4/*memory for ids_list*/) / 1000000;
  EXPECT_TRUE((used_mb > theoritical_mb * 0.99) &&
              (used_mb < theoritical_mb * 1.02));
}
} //namespace embedding
} //namespace tensorflow