/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
you may not use this file except in compliance with the License.
|
You may obtain a copy of the License at
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
Unless required by applicable law or agreed to in writing, software
|
distributed under the License is distributed on an "AS IS" BASIS,
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
See the License for the specific language governing permissions and
|
limitations under the License.
|
==============================================================================*/
|
|
#include "tensorflow/core/util/tensor_slice_writer.h"
|
|
#include <array>
|
|
#include "tensorflow/core/framework/tensor_shape.pb.h"
|
#include "tensorflow/core/framework/versions.pb.h"
|
#include "tensorflow/core/lib/core/status_test_util.h"
|
#include "tensorflow/core/lib/core/stringpiece.h"
|
#include "tensorflow/core/lib/io/path.h"
|
#include "tensorflow/core/lib/strings/str_util.h"
|
#include "tensorflow/core/platform/logging.h"
|
#include "tensorflow/core/platform/protobuf.h"
|
#include "tensorflow/core/platform/test.h"
|
#include "tensorflow/core/public/version.h"
|
#include "tensorflow/core/util/saved_tensor_slice_util.h"
|
#include "tensorflow/core/util/tensor_slice_reader.h"
|
|
namespace tensorflow {
|
|
namespace checkpoint {
|
|
class TensorSliceWriteTestHelper {
|
public:
|
static void CheckEntries(const string& fname);
|
static void GetData(TensorSliceReader::Table* table, const string& name,
|
const TensorSlice& slice, SavedSlice* ss);
|
};
|
|
namespace {
|
|
// Testing that an array is what is expected
|
void ExpectIdenticalFloatArrays(const float* expected, int size,
|
const float* actual) {
|
// TODO(yangke): copy some of the Dump* functions over
|
// LOG(INFO) << "Expected = " << DumpFloatArray(expected, size);
|
// LOG(INFO) << "Actual = " << DumpFloatArray(actual, size);
|
for (int i = 0; i < size; ++i) {
|
EXPECT_NEAR(expected[i], actual[i], 1e-6);
|
}
|
}
|
|
template <typename T, typename U>
|
void ExpectIdenticalIntArrays(const T* expected, int size, const U* actual) {
|
for (int i = 0; i < size; ++i) {
|
EXPECT_EQ(expected[i], static_cast<T>(actual[i]));
|
}
|
}
|
|
// Nifty routine to get the size of an array
|
template <typename T, unsigned SIZE>
|
inline size_t ArraySize(const T (&v)[SIZE]) {
|
return SIZE;
|
}
|
|
// A simple test on writing a few tensor slices
|
// TODO(yangke): refactor into smaller tests: will do as we add more stuff to
|
// the writer.
|
TEST(TensorSliceWriteTest, SimpleWrite) {
|
const string filename = io::JoinPath(testing::TmpDir(), "checkpoint");
|
|
TensorSliceWriter writer(filename, CreateTableTensorSliceBuilder);
|
|
// Add some int32 tensor slices
|
{
|
TensorShape shape({5, 10});
|
TensorSlice slice = TensorSlice::ParseOrDie("-:0,1");
|
const int32 data[] = {0, 1, 2, 3, 4};
|
TF_CHECK_OK(writer.Add("test", shape, slice, data));
|
}
|
|
// Two slices share the same tensor name
|
{
|
TensorShape shape({5, 10});
|
TensorSlice slice = TensorSlice::ParseOrDie("-:3,1");
|
const int32 data[] = {10, 11, 12, 13, 14};
|
TF_CHECK_OK(writer.Add("test", shape, slice, data));
|
}
|
|
// Another slice from a different float tensor -- it has a different name and
|
// should be inserted in front of the previous tensor
|
{
|
TensorShape shape({3, 2});
|
TensorSlice slice = TensorSlice::ParseOrDie("-:-");
|
const float data[] = {1.2, 1.3, 1.4, 2.1, 2.2, 2.3};
|
TF_CHECK_OK(writer.Add("AA", shape, slice, data));
|
}
|
|
// A slice with int64 data
|
{
|
TensorShape shape({5, 10});
|
TensorSlice slice = TensorSlice::ParseOrDie("-:3,1");
|
const int64 data[] = {10, 11, 12, 13, 14};
|
TF_CHECK_OK(writer.Add("int64", shape, slice, data));
|
}
|
|
// A slice with int16 data
|
{
|
TensorShape shape({5, 10});
|
TensorSlice slice = TensorSlice::ParseOrDie("-:3,1");
|
const int16 data[] = {10, 11, 12, 13, 14};
|
TF_CHECK_OK(writer.Add("int16", shape, slice, data));
|
}
|
|
TF_CHECK_OK(writer.Finish());
|
|
// Now we examine the checkpoint file manually.
|
TensorSliceWriteTestHelper::CheckEntries(filename);
|
}
|
|
} // namespace
|
|
void TensorSliceWriteTestHelper::GetData(TensorSliceReader::Table* table,
|
const string& name,
|
const TensorSlice& slice,
|
SavedSlice* ss) {
|
string key = EncodeTensorNameSlice(name, slice);
|
string value;
|
EXPECT_TRUE(table->Get(key, &value));
|
SavedTensorSlices sts;
|
EXPECT_TRUE(ParseProtoUnlimited(&sts, value));
|
EXPECT_FALSE(sts.has_meta());
|
*ss = sts.data();
|
EXPECT_EQ(name, ss->name());
|
TensorSlice slice2(ss->slice());
|
EXPECT_EQ(slice.DebugString(), slice2.DebugString());
|
}
|
|
void TensorSliceWriteTestHelper::CheckEntries(const string& fname) {
|
TensorSliceReader::Table* tptr;
|
TF_CHECK_OK(OpenTableTensorSliceReader(fname, &tptr));
|
std::unique_ptr<TensorSliceReader::Table> table(tptr);
|
CHECK_NOTNULL(table.get());
|
|
// We expect a block of SavedTensorSlices
|
string value;
|
ASSERT_TRUE(table->Get(kSavedTensorSlicesKey, &value));
|
{
|
SavedTensorSlices sts;
|
EXPECT_TRUE(ParseProtoUnlimited(&sts, value));
|
// We also expect two entries for the tensors
|
EXPECT_TRUE(sts.has_meta());
|
EXPECT_EQ(4, sts.meta().tensor_size());
|
// We should have written nontrivial version information
|
EXPECT_LT(0, TF_CHECKPOINT_VERSION);
|
EXPECT_EQ(TF_CHECKPOINT_VERSION, sts.meta().versions().producer());
|
EXPECT_EQ(TF_CHECKPOINT_VERSION_MIN_CONSUMER,
|
sts.meta().versions().min_consumer());
|
// We don't expect any data in the first block.
|
EXPECT_FALSE(sts.has_data());
|
// The two tensors should be stored in the same order as they are first
|
// created.
|
{
|
// The two slices of the "test" tensor
|
const SavedSliceMeta& ssm = sts.meta().tensor(0);
|
EXPECT_EQ("test", ssm.name());
|
EXPECT_EQ(
|
"dim { size: 5 } "
|
"dim { size: 10 }",
|
ssm.shape().ShortDebugString());
|
EXPECT_EQ(DT_INT32, ssm.type());
|
EXPECT_EQ(2, ssm.slice_size());
|
TensorSlice s0(ssm.slice(0));
|
TensorSlice s1(ssm.slice(1));
|
EXPECT_EQ("-:0,1", s0.DebugString());
|
EXPECT_EQ("-:3,1", s1.DebugString());
|
}
|
{
|
// The "AA" tensor
|
const SavedSliceMeta& ssm = sts.meta().tensor(1);
|
EXPECT_EQ("AA", ssm.name());
|
EXPECT_EQ(
|
"dim { size: 3 } "
|
"dim { size: 2 }",
|
ssm.shape().ShortDebugString());
|
EXPECT_EQ(DT_FLOAT, ssm.type());
|
EXPECT_EQ(1, ssm.slice_size());
|
TensorSlice s0(ssm.slice(0));
|
EXPECT_EQ("-:-", s0.DebugString());
|
}
|
{
|
// The "int64" tensor
|
const SavedSliceMeta& ssm = sts.meta().tensor(2);
|
EXPECT_EQ("int64", ssm.name());
|
EXPECT_EQ(
|
"dim { size: 5 } "
|
"dim { size: 10 }",
|
ssm.shape().ShortDebugString());
|
EXPECT_EQ(DT_INT64, ssm.type());
|
EXPECT_EQ(1, ssm.slice_size());
|
TensorSlice s0(ssm.slice(0));
|
EXPECT_EQ("-:3,1", s0.DebugString());
|
}
|
{
|
// The "int16" tensor
|
const SavedSliceMeta& ssm = sts.meta().tensor(3);
|
EXPECT_EQ("int16", ssm.name());
|
EXPECT_EQ(
|
"dim { size: 5 } "
|
"dim { size: 10 }",
|
ssm.shape().ShortDebugString());
|
EXPECT_EQ(DT_INT16, ssm.type());
|
EXPECT_EQ(1, ssm.slice_size());
|
TensorSlice s0(ssm.slice(0));
|
EXPECT_EQ("-:3,1", s0.DebugString());
|
}
|
}
|
|
// We expect 5 blocks of tensor data
|
{
|
// Block 1: we expect it to be the full slice of the "AA" tensor
|
SavedSlice ss;
|
GetData(table.get(), "AA", TensorSlice(2), &ss);
|
const float data[] = {1.2, 1.3, 1.4, 2.1, 2.2, 2.3};
|
EXPECT_EQ(ArraySize(data), ss.data().float_val_size());
|
ExpectIdenticalFloatArrays(data, ArraySize(data),
|
ss.data().float_val().data());
|
}
|
|
{
|
// Block 2: we expect it to be the first slice of the "test" tensor
|
SavedSlice ss;
|
GetData(table.get(), "test", TensorSlice({{0, -1}, {0, 1}}), &ss);
|
const int32 data[] = {0, 1, 2, 3, 4};
|
EXPECT_EQ(ArraySize(data), ss.data().int_val_size());
|
ExpectIdenticalIntArrays(data, ArraySize(data), ss.data().int_val().data());
|
}
|
|
{
|
// Block 3: we expect it to be the second slice of the "test" tensor
|
SavedSlice ss;
|
GetData(table.get(), "test", TensorSlice({{0, -1}, {3, 1}}), &ss);
|
const int32 data[] = {10, 11, 12, 13, 14};
|
EXPECT_EQ(ArraySize(data), ss.data().int_val_size());
|
ExpectIdenticalIntArrays(data, ArraySize(data), ss.data().int_val().data());
|
}
|
|
{
|
// Block 4: we expect it to be the slice of the "int64" tensor
|
SavedSlice ss;
|
GetData(table.get(), "int64", TensorSlice({{0, -1}, {3, 1}}), &ss);
|
const int64 data[] = {10, 11, 12, 13, 14};
|
EXPECT_EQ(ArraySize(data), ss.data().int64_val_size());
|
ExpectIdenticalIntArrays(data, ArraySize(data),
|
ss.data().int64_val().data());
|
}
|
|
{
|
// Block 5: we expect it to be the slice of the "int16" tensor
|
SavedSlice ss;
|
GetData(table.get(), "int16", TensorSlice({{0, -1}, {3, 1}}), &ss);
|
const int16 data[] = {10, 11, 12, 13, 14};
|
EXPECT_EQ(ArraySize(data), ss.data().int_val_size());
|
ExpectIdenticalIntArrays(data, ArraySize(data), ss.data().int_val().data());
|
}
|
}
|
|
template <typename DT>
|
size_t BytesPerElementHelper(DT value) {
|
SavedSlice ss;
|
std::array<DT, 1> lo_data;
|
std::fill(lo_data.begin(), lo_data.end(), value);
|
TF_EXPECT_OK(
|
TensorSliceWriter::SaveData(lo_data.data(), lo_data.size(), &ss));
|
size_t lo_byte_size = ss.ByteSizeLong();
|
|
std::array<DT, 1001> hi_data;
|
std::fill(hi_data.begin(), hi_data.end(), value);
|
TF_EXPECT_OK(
|
TensorSliceWriter::SaveData(hi_data.data(), hi_data.size(), &ss));
|
size_t hi_byte_size = ss.ByteSizeLong();
|
|
return (hi_byte_size - lo_byte_size) / (hi_data.size() - lo_data.size());
|
}
|
|
TEST(TensorSliceWriteTest, CheckpointSize) {
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_BOOL),
|
BytesPerElementHelper<bool>(false));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_BOOL),
|
BytesPerElementHelper<bool>(true));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_FLOAT),
|
BytesPerElementHelper<float>(-1.0));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_DOUBLE),
|
BytesPerElementHelper<double>(-1.0));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_COMPLEX64),
|
BytesPerElementHelper<complex64>(-1.0));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_COMPLEX128),
|
BytesPerElementHelper<complex128>(-1.0));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_INT32),
|
BytesPerElementHelper<int32>(-1));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_INT64),
|
BytesPerElementHelper<int64>(-1));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_UINT16),
|
BytesPerElementHelper<uint16>(std::numeric_limits<uint16>::max()));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_UINT8),
|
BytesPerElementHelper<uint8>(std::numeric_limits<uint8>::max()));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_INT8),
|
BytesPerElementHelper<int8>(-1));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_INT16),
|
BytesPerElementHelper<int16>(-1));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_QINT8),
|
BytesPerElementHelper<qint8>(-1));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_QUINT8),
|
BytesPerElementHelper<quint8>(std::numeric_limits<uint8>::max()));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_QINT32),
|
BytesPerElementHelper<qint32>(-1));
|
EXPECT_EQ(TensorSliceWriter::MaxBytesPerElement(DT_HALF),
|
BytesPerElementHelper<Eigen::half>(Eigen::half(-1.0)));
|
}
|
|
TEST(TensorSliceWriteTest, SizeErrors) {
|
const string filename = io::JoinPath(testing::TmpDir(), "checkpoint");
|
|
TensorSliceWriter writer(filename, CreateTableTensorSliceBuilder);
|
|
// Add a 300MB int8 tensor slice, which will fail because it expands to 3GB.
|
{
|
TensorShape shape({300, 1000000});
|
TensorSlice slice = TensorSlice::ParseOrDie("-:-");
|
const std::vector<int8> data(300000000, -1);
|
Status s = writer.Add("test1", shape, slice, data.data());
|
EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
|
EXPECT_TRUE(str_util::StrContains(
|
s.error_message(), "Tensor slice is too large to serialize"));
|
}
|
|
// Add a large string tensor slice, which will fail.
|
{
|
TensorShape shape({256, 1024});
|
TensorSlice slice = TensorSlice::ParseOrDie("-:-");
|
const std::vector<string> data(256 * 1024, std::string(8192, 'f'));
|
Status s = writer.Add("test2", shape, slice, data.data());
|
EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
|
EXPECT_TRUE(str_util::StrContains(
|
s.error_message(), "Tensor slice is too large to serialize"));
|
}
|
}
|
|
} // namespace checkpoint
|
|
} // namespace tensorflow
|
