| # Copyright 2019 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. |
| # ============================================================================== |
| """Tests for the `SnapshotDataset` transformation.""" |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| |
| import multiprocessing |
| import os |
| import shutil |
| import time |
| |
| from absl.testing import parameterized |
| import numpy as np |
| |
| from tensorflow.python.data.experimental.ops import snapshot |
| from tensorflow.python.data.kernel_tests import checkpoint_test_base |
| from tensorflow.python.data.kernel_tests import test_base |
| from tensorflow.python.data.kernel_tests import tf_record_test_base |
| from tensorflow.python.data.ops import dataset_ops |
| from tensorflow.python.data.ops import readers as core_readers |
| from tensorflow.python.framework import combinations |
| from tensorflow.python.framework import errors |
| from tensorflow.python.framework import ops |
| from tensorflow.python.ops import gen_array_ops |
| from tensorflow.python.ops import string_ops |
| from tensorflow.python.platform import test |
| |
| |
| def is_graphdef_file(filename): |
| return filename.endswith("-graph.pbtxt") |
| |
| |
| def is_temp_file(filename): |
| return "-tmp-" in filename |
| |
| |
| def listdir_and_filter(dirname, filter_fn): |
| return [path for path in sorted(os.listdir(dirname)) if filter_fn(path)] |
| |
| |
| class SnapshotTest(tf_record_test_base.TFRecordTestBase, |
| parameterized.TestCase): |
| |
| def setUp(self): |
| super(SnapshotTest, self).setUp() |
| tmpdir = self.get_temp_dir() |
| tmpdir = os.path.join(tmpdir, "snapshot") |
| os.mkdir(tmpdir) |
| self._snapshot_dir = tmpdir |
| |
| def tearDown(self): |
| super(SnapshotTest, self).tearDown() |
| shutil.rmtree(self._snapshot_dir) |
| |
| def createTFRecords(self, num_files=10, num_records=100): |
| self._num_files = num_files |
| self._num_records = num_records |
| self._filenames = self._createFiles() |
| |
| def removeTFRecords(self): |
| for filename in self._filenames: |
| os.remove(filename) |
| self._filenames = [] |
| self._num_files = None |
| self._num_records = None |
| |
| def assertDatasetProducesSet(self, dataset, expected): |
| actual = [] |
| next_fn = self.getNext(dataset) |
| for _ in range(len(expected)): |
| elem = self.evaluate(next_fn()) |
| actual.append(elem) |
| self.assertCountEqual(actual, expected) |
| with self.assertRaises(errors.OutOfRangeError): |
| self.evaluate(next_fn()) |
| |
| def assertSnapshotDirectoryContains(self, directory, num_fingerprints, |
| num_runs_per_fingerprint, |
| num_snapshot_shards_per_run): |
| |
| # Ignore the graphdef pbtxts we write for debugging purposes and temporary |
| # files that are an artifact of how TF writes files. |
| dirlist = listdir_and_filter( |
| directory, |
| lambda p: not (is_graphdef_file(p) or is_temp_file(p))) |
| self.assertLen(dirlist, num_fingerprints) |
| |
| for i in range(num_fingerprints): |
| fingerprint_dir = os.path.join(directory, dirlist[i]) |
| fingerprint_dir_list = listdir_and_filter( |
| fingerprint_dir, lambda p: not is_temp_file(p)) |
| self.assertLen(fingerprint_dir_list, num_runs_per_fingerprint + 1) |
| self.assertEqual(fingerprint_dir_list[num_runs_per_fingerprint], |
| "snapshot.metadata") |
| |
| for j in range(num_runs_per_fingerprint): |
| run_dir = os.path.join(fingerprint_dir, fingerprint_dir_list[j]) |
| run_dirlist = sorted(os.listdir(run_dir)) |
| self.assertLen(run_dirlist, num_snapshot_shards_per_run) |
| |
| file_counter = 0 |
| for filename in run_dirlist: |
| self.assertEqual(filename, "%08d.shard" % file_counter) |
| file_counter += 1 |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testCreateSnapshotDataset(self): |
| dataset = dataset_ops.Dataset.from_tensors([1, 2, 3]) |
| dataset.snapshot(path=self._snapshot_dir) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadSnapshotDatasetDefault(self): |
| self.createTFRecords() |
| filenames = self._filenames |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 100) |
| ] |
| |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset, expected) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| self.removeTFRecords() |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset2, expected) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadSnapshotDatasetAutoWriteSnappyRead(self): |
| self.createTFRecords() |
| filenames = self._filenames |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 100) |
| ] |
| |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.snapshot(path=self._snapshot_dir, compression="AUTO") |
| self.assertDatasetProduces(dataset, expected) |
| |
| self.removeTFRecords() |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.snapshot( |
| path=self._snapshot_dir, |
| compression="SNAPPY" |
| ) |
| self.assertDatasetProduces(dataset2, expected) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadSnapshotDatasetCustomShardFn(self): |
| self.createTFRecords() |
| filenames = self._filenames |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 100) |
| ] |
| |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.snapshot( |
| path=self._snapshot_dir, |
| shard_func=lambda _: np.int64(0) |
| ) |
| self.assertDatasetProduces(dataset, expected) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=1) |
| |
| self.removeTFRecords() |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.snapshot( |
| path=self._snapshot_dir, |
| shard_func=lambda _: 0 |
| ) |
| self.assertDatasetProduces(dataset2, expected) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadSnapshotDatasetCustomReaderFn(self): |
| self.createTFRecords() |
| filenames = self._filenames |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 100) |
| ] |
| |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.snapshot( |
| path=self._snapshot_dir, |
| reader_func=( |
| lambda ds: ds.interleave( # pylint:disable=g-long-lambda |
| lambda x: x, |
| cycle_length=4, |
| num_parallel_calls=4))) |
| self.assertDatasetProduces(dataset, expected) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| self.removeTFRecords() |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.snapshot( |
| self._snapshot_dir, |
| reader_func=( |
| lambda ds: ds.interleave( # pylint:disable=g-long-lambda |
| lambda x: x, |
| cycle_length=4, |
| num_parallel_calls=4))) |
| self.assertDatasetProducesSet(dataset2, expected) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testSnapshotDatasetInvalidShardFn(self): |
| dataset = dataset_ops.Dataset.range(1000) |
| with self.assertRaises(TypeError): |
| dataset = dataset.snapshot( |
| path=self._snapshot_dir, |
| shard_func=lambda _: "invalid_fn") |
| next_fn = self.getNext(dataset) |
| self.evaluate(next_fn()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testSnapshotDatasetInvalidReaderFn(self): |
| dataset = dataset_ops.Dataset.range(1000) |
| with self.assertRaises(TypeError): |
| dataset = dataset.snapshot( |
| path=self._snapshot_dir, |
| reader_func=lambda x: x + 1 |
| ) |
| next_fn = self.getNext(dataset) |
| self.evaluate(next_fn()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testRoundtripEmptySnapshot(self): |
| dataset = dataset_ops.Dataset.range(0) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset, []) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=0) |
| |
| dataset2 = dataset_ops.Dataset.range(0) |
| dataset2 = dataset.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset2, []) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotDatasetSimple(self): |
| dataset = dataset_ops.Dataset.range(1000) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset, list(range(1000))) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotDatasetMultipleFingerprints(self): |
| dataset1 = dataset_ops.Dataset.range(1000) |
| dataset1 = dataset1.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset1, list(range(1000))) |
| |
| dataset2 = dataset_ops.Dataset.range(2000) |
| dataset2 = dataset2.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset2, list(range(2000))) |
| |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=2, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotDatasetSameFingerprintMultipleCompleteRuns(self): |
| dataset1 = dataset_ops.Dataset.range(1000) |
| dataset1 = dataset1.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset1, list(range(1000))) |
| dataset2 = dataset_ops.Dataset.range(1000) |
| dataset2 = dataset2.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset2, list(range(1000))) |
| |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotDatasetSameFingerprintIncompleteRunRestart(self): |
| dataset1 = dataset_ops.Dataset.range(1000) |
| dataset1 = dataset1.snapshot(path=self._snapshot_dir) |
| next1 = self.getNext(dataset1) |
| for i in range(500): |
| self.assertEqual(i, self.evaluate(next1())) |
| |
| dataset2 = dataset_ops.Dataset.range(1000) |
| dataset2 = dataset2.snapshot(path=self._snapshot_dir) |
| next2 = self.getNext(dataset2) |
| for i in range(500): |
| self.assertEqual(i, self.evaluate(next2())) |
| |
| for i in range(500, 1000): |
| self.assertEqual(i, self.evaluate(next1())) |
| self.assertEqual(i, self.evaluate(next2())) |
| |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=2, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotCustomShardFunction(self): |
| dataset = dataset_ops.Dataset.range(1000) |
| dataset = dataset.enumerate() |
| dataset = dataset.snapshot( |
| path=self._snapshot_dir, |
| shard_func=lambda i, _: i % 2 |
| ) |
| dataset = dataset.map(lambda _, elem: elem) |
| self.assertDatasetProduces(dataset, list(range(1000))) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=2) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotDatasetWithTuples(self): |
| dataset1 = dataset_ops.Dataset.range(0, 1000) |
| dataset2 = dataset_ops.Dataset.range(1000, 2000) |
| dataset3 = dataset_ops.Dataset.range(2000, 3000) |
| dataset4 = dataset_ops.Dataset.range(3000, 4000) |
| |
| dataset = dataset_ops.Dataset.zip((dataset1, dataset2, dataset3, dataset4)) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| |
| expected = list( |
| zip( |
| range(0, 1000), range(1000, 2000), range(2000, 3000), |
| range(3000, 4000))) |
| self.assertDatasetProduces(dataset, expected) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotShuffleSameFingerprint(self): |
| |
| def make_dataset(): |
| dataset = dataset_ops.Dataset.range(1000) |
| dataset = dataset.shuffle(1000) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| return dataset |
| |
| dataset1 = make_dataset() |
| self.assertDatasetProducesSet(dataset1, list(range(1000))) |
| dataset2 = make_dataset() |
| self.assertDatasetProducesSet(dataset2, list(range(1000))) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadUsingFlatMap(self): |
| dataset = dataset_ops.Dataset.range(1000) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProduces(dataset, list(range(1000))) |
| flat_map = dataset_ops.Dataset.from_tensors(dataset).flat_map(lambda x: x) |
| self.assertDatasetProduces(flat_map, list(range(1000))) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadOptimizableUsingFlatMap(self): |
| dataset = dataset_ops.Dataset.range(1000) |
| # Will be optimized into ShuffleAndRepeat. |
| dataset = dataset.shuffle(10) |
| dataset = dataset.repeat(2) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| self.assertDatasetProducesSet(dataset, 2 * list(range(1000))) |
| flat_map = dataset_ops.Dataset.from_tensors(dataset).flat_map(lambda x: x) |
| self.assertDatasetProducesSet(flat_map, 2 * list(range(1000))) |
| self.assertSnapshotDirectoryContains( |
| self._snapshot_dir, |
| num_fingerprints=1, |
| num_runs_per_fingerprint=1, |
| num_snapshot_shards_per_run=multiprocessing.cpu_count()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testRepeatAndPrefetch(self): |
| """This test reproduces github.com/tensorflow/tensorflow/issues/48903.""" |
| dataset = dataset_ops.Dataset.from_tensor_slices(np.random.rand(16, 32)) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| dataset = dataset.shuffle(buffer_size=16) |
| dataset = dataset.batch(16) |
| dataset = dataset.repeat() |
| dataset = dataset.prefetch(1) |
| next_element = self.getNext(dataset) |
| for _ in range(30): |
| self.evaluate(next_element()) |
| |
| |
| class LegacySnapshotTest(tf_record_test_base.TFRecordTestBase, |
| parameterized.TestCase): |
| |
| def setUp(self): |
| super(LegacySnapshotTest, self).setUp() |
| self.removeTFRecords() |
| tmpdir = self.get_temp_dir() |
| tmpdir = os.path.join(tmpdir, "snapshot") |
| os.mkdir(tmpdir) |
| self.snapshot_dir = tmpdir |
| |
| def tearDown(self): |
| super(LegacySnapshotTest, self).tearDown() |
| shutil.rmtree(self.snapshot_dir) |
| |
| def removeTFRecords(self): |
| for filename in self._filenames: |
| os.remove(filename) |
| self._filenames = [] |
| |
| def setUpTFRecord(self, num_files=10, num_records=10): |
| self._num_files = num_files |
| self._num_records = num_records |
| self._filenames = self._createFiles() |
| |
| def makeSnapshotDirectory(self): |
| return self.snapshot_dir |
| |
| def assertSnapshotDirectoryContains(self, directory, num_fingerprints, |
| num_runs_per_fp, num_snapshot_files): |
| # Ignore the graphdef pbtxts we write for debugging purposes and temporary |
| # files that are an artifact of how TF writes files. |
| dirlist = listdir_and_filter( |
| directory, |
| lambda p: not (is_graphdef_file(p) or is_temp_file(p))) |
| self.assertLen(dirlist, num_fingerprints) |
| |
| for i in range(num_fingerprints): |
| fingerprint_dir = os.path.join(directory, dirlist[i]) |
| fingerprint_dir_list = listdir_and_filter( |
| fingerprint_dir, lambda p: not is_temp_file(p)) |
| self.assertLen(fingerprint_dir_list, num_runs_per_fp + 1) |
| self.assertEqual(fingerprint_dir_list[num_runs_per_fp], |
| "snapshot.metadata") |
| |
| for j in range(num_runs_per_fp): |
| run_dir = os.path.join(fingerprint_dir, fingerprint_dir_list[j]) |
| run_dirlist = sorted(os.listdir(run_dir)) |
| self.assertLen(run_dirlist, num_snapshot_files) |
| |
| file_counter = 0 |
| for filename in run_dirlist: |
| self.assertEqual(filename, "%08d.snapshot" % file_counter) |
| file_counter += 1 |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteDifferentPipelinesInOneDirectory(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.range(1000) |
| dataset = dataset.apply(snapshot.legacy_snapshot(tmpdir)) |
| self.assertDatasetProduces(dataset, list(range(1000))) |
| |
| dataset = dataset_ops.Dataset.range(1001) |
| dataset = dataset.apply(snapshot.legacy_snapshot(tmpdir)) |
| self.assertDatasetProduces(dataset, list(range(1001))) |
| |
| self.assertSnapshotDirectoryContains(tmpdir, 2, 1, 1) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testWriteSnapshotMultipleSimultaneous(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset1 = dataset_ops.Dataset.range(1000) |
| dataset1 = dataset1.apply(snapshot.legacy_snapshot(tmpdir)) |
| next1 = self.getNext(dataset1) |
| |
| dataset2 = dataset_ops.Dataset.range(1000) |
| dataset2 = dataset2.apply(snapshot.legacy_snapshot(tmpdir)) |
| next2 = self.getNext(dataset2) |
| |
| for i in range(0, 1000): |
| self.assertEqual(i, self.evaluate(next1())) |
| self.assertEqual(i, self.evaluate(next2())) |
| |
| # we check that only one copy of the metadata has been written, and the |
| # one that lost the race would be in passthrough mode. |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testGetNextCreatesDir(self): |
| tmpdir = self.snapshot_dir |
| |
| # We create two iterators but call getNext on only one. |
| dataset1 = dataset_ops.Dataset.range(1000) |
| dataset1 = dataset1.apply(snapshot.legacy_snapshot(tmpdir)) |
| next1 = self.getNext(dataset1) |
| |
| dataset2 = dataset_ops.Dataset.range(1001) |
| dataset2 = dataset2.apply(snapshot.legacy_snapshot(tmpdir)) |
| _ = self.getNext(dataset2) |
| |
| for _ in range(1000): |
| self.evaluate(next1()) |
| |
| # We check that only one directory is created. |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(compression=[ |
| snapshot.COMPRESSION_NONE, snapshot.COMPRESSION_GZIP, |
| snapshot.COMPRESSION_SNAPPY |
| ]))) |
| def testWriteSnapshotSimpleSuccessful(self, compression): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.range(1000) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, compression=compression)) |
| self.assertDatasetProduces(dataset, list(range(1000))) |
| |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(compression=[ |
| snapshot.COMPRESSION_NONE, snapshot.COMPRESSION_GZIP, |
| snapshot.COMPRESSION_SNAPPY |
| ]))) |
| def testWriteSnapshotRepeatAfterwards(self, compression): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.range(10) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, compression=compression)) |
| dataset = dataset.repeat(10) |
| self.assertDatasetProduces(dataset, list(range(10)) * 10) |
| |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(compression=[ |
| snapshot.COMPRESSION_NONE, snapshot.COMPRESSION_GZIP, |
| snapshot.COMPRESSION_SNAPPY |
| ]))) |
| def testWriteSnapshotMixTypes(self, compression): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.range(10) |
| |
| def map_fn(x): |
| return (x, string_ops.as_string(x), string_ops.as_string(2 * x), 2 * x) |
| |
| dataset = dataset.map(map_fn) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, compression=compression)) |
| dataset = dataset.repeat(10) |
| |
| expected = [] |
| for i in range(10): |
| expected.append((i, str(i), str(2 * i), 2 * i)) |
| self.assertDatasetProduces(dataset, expected * 10) |
| |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testSpecifySnapshotNameWriteAndRead(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.range(10) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, snapshot_name="my_custom_snapshot")) |
| dataset = dataset.repeat(10) |
| self.assertDatasetProduces(dataset, list(range(10)) * 10) |
| |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| self.assertTrue( |
| os.path.exists(os.path.join(tmpdir, "custom-my_custom_snapshot"))) |
| self.assertTrue( |
| os.path.exists( |
| os.path.join(tmpdir, "custom-my_custom_snapshot", "custom"))) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testForcePassthroughMode(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.range(10) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, mode="passthrough")) |
| dataset = dataset.repeat(10) |
| self.assertDatasetProduces(dataset, list(range(10)) * 10) |
| |
| self.assertSnapshotDirectoryContains(tmpdir, 0, 0, 0) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testForceWriteMode(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.range(10) |
| dataset = dataset.apply(snapshot.legacy_snapshot(tmpdir, mode="write")) |
| dataset = dataset.repeat(10) |
| self.assertDatasetProduces(dataset, list(range(10)) * 10) |
| |
| # We will end up writing 10 different runs. |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 10, 1) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testForceReadMode(self): |
| tmpdir = self.snapshot_dir |
| |
| # We write a copy of the snapshot first. |
| dataset = dataset_ops.Dataset.range(10) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, mode="write", snapshot_name="my_custom_snapshot")) |
| self.assertDatasetProduces(dataset, list(range(10))) |
| |
| # We move the run to a new name. |
| shutil.move( |
| os.path.join(tmpdir, "custom-my_custom_snapshot"), |
| os.path.join(tmpdir, "custom-my_custom_snapshot_2")) |
| |
| # Even though the snapshot.metadata is pointing to the old run that no |
| # longer exists after we moved, we force it to read from the run we specify. |
| dataset = dataset_ops.Dataset.range(10) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, mode="read", snapshot_name="my_custom_snapshot_2")) |
| self.assertDatasetProduces(dataset, list(range(10))) |
| |
| # We should still have one snapshot and one run. |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testForceReadNonexistentSnapshot(self): |
| tmpdir = self.snapshot_dir |
| dataset = dataset_ops.Dataset.range(10) |
| with self.assertRaises(errors.NotFoundError): |
| dataset = dataset.apply(snapshot.legacy_snapshot(tmpdir, mode="read")) |
| get_next = self.getNext(dataset) |
| self.evaluate(get_next()) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testForceReadNonexistentNamedSnapshot(self): |
| tmpdir = self.snapshot_dir |
| dataset = dataset_ops.Dataset.range(10) |
| with self.assertRaises(errors.NotFoundError): |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, mode="read", snapshot_name="my_nonexistent_snapshot")) |
| get_next = self.getNext(dataset) |
| self.evaluate(get_next()) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(compression=[ |
| snapshot.COMPRESSION_NONE, snapshot.COMPRESSION_GZIP, |
| snapshot.COMPRESSION_SNAPPY |
| ]))) |
| def testReadSnapshotBackAfterWrite(self, compression): |
| self.setUpTFRecord() |
| filenames = self._filenames |
| |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 10) |
| ] |
| |
| tmpdir = self.snapshot_dir |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, compression=compression)) |
| self.assertDatasetProduces(dataset, expected) |
| |
| # remove the original files and try to read the data back only from snapshot |
| self.removeTFRecords() |
| |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.apply( |
| snapshot.legacy_snapshot(tmpdir, compression=compression)) |
| self.assertDatasetProduces(dataset2, expected) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadShuffledSnapshotAfterWrite(self): |
| self.setUpTFRecord(num_files=10, num_records=50) |
| filenames = self._filenames |
| |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 50) |
| ] |
| |
| tmpdir = self.snapshot_dir |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, shard_size_bytes=100)) |
| self.assertDatasetProduces(dataset, expected) |
| |
| # remove the original files and try to read the data back only from snapshot |
| self.removeTFRecords() |
| |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, shard_size_bytes=100, shuffle_on_read=True)) |
| next2 = self.getNext(dataset2) |
| |
| res1 = self.evaluate(next2()) |
| res2 = self.evaluate(next2()) |
| res3 = self.evaluate(next2()) |
| res4 = self.evaluate(next2()) |
| res5 = self.evaluate(next2()) |
| |
| # make sure that we don't read the file back in the same order. |
| self.assertNotEqual([res1, res2, res3, res4, res5], expected[0:5]) |
| |
| # make sure all the elements are still there |
| dataset3 = core_readers._TFRecordDataset(filenames) |
| dataset3 = dataset3.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, shard_size_bytes=100, shuffle_on_read=True)) |
| self.assertDatasetProduces(dataset3, expected, assert_items_equal=True) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testReadShuffledSnapshotWithSeedAfterWrite(self): |
| self.setUpTFRecord(num_files=10, num_records=50) |
| filenames = self._filenames |
| |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 50) |
| ] |
| |
| tmpdir = self.snapshot_dir |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot(tmpdir, shard_size_bytes=10)) |
| self.assertDatasetProduces(dataset, expected) |
| |
| # remove the original files and try to read the data back only from snapshot |
| self.removeTFRecords() |
| |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, |
| shard_size_bytes=10, |
| shuffle_on_read=True, |
| shuffle_seed=123456)) |
| next2 = self.getNext(dataset2) |
| |
| dataset3 = core_readers._TFRecordDataset(filenames) |
| dataset3 = dataset3.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, |
| shard_size_bytes=10, |
| shuffle_on_read=True, |
| shuffle_seed=123456)) |
| next3 = self.getNext(dataset3) |
| |
| # make sure that the items are read back in the same order for both datasets |
| for _ in range(500): |
| res2 = self.evaluate(next2()) |
| res3 = self.evaluate(next3()) |
| self.assertEqual(res2, res3) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(compression=[ |
| snapshot.COMPRESSION_NONE, snapshot.COMPRESSION_GZIP, |
| snapshot.COMPRESSION_SNAPPY |
| ]))) |
| def testReadSnapshotParallelAfterWrite(self, compression): |
| self.setUpTFRecord(5, 500) |
| filenames = self._filenames |
| |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 5) |
| for r in range(0, 500) |
| ] |
| |
| tmpdir = self.snapshot_dir |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, |
| shard_size_bytes=1024 * 1024, |
| num_reader_threads=2, |
| reader_buffer_size=10, |
| compression=compression)) |
| self.assertDatasetProduces(dataset, expected, assert_items_equal=True) |
| |
| # remove the original files and try to read the data back only from |
| # snapshot. |
| self.removeTFRecords() |
| |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, |
| shard_size_bytes=1024 * 1024, |
| num_reader_threads=2, |
| reader_buffer_size=10, |
| compression=compression)) |
| self.assertDatasetProduces(dataset2, expected, assert_items_equal=True) |
| |
| # Not testing Snappy here because Snappy reads currently require a lot of |
| # memory. |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.times( |
| combinations.combine(compression=[ |
| snapshot.COMPRESSION_NONE, snapshot.COMPRESSION_GZIP |
| ]), |
| combinations.combine(threads=2, size=[1, 2]) + |
| combinations.combine(threads=8, size=[1, 4, 8])))) |
| def testReadSnapshotBackAfterMultiThreadedWrite(self, compression, threads, |
| size): |
| self.setUpTFRecord() |
| filenames = self._filenames |
| |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 10) |
| ] |
| |
| tmpdir = self.snapshot_dir |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, |
| compression=compression, |
| num_writer_threads=threads, |
| writer_buffer_size=size)) |
| self.assertDatasetProduces(dataset, expected) |
| |
| # remove the original files and try to read the data back only from |
| # snapshot |
| self.removeTFRecords() |
| |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.apply( |
| snapshot.legacy_snapshot(tmpdir, compression=compression)) |
| self.assertDatasetProduces(dataset2, expected, assert_items_equal=True) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testSameFingerprintWithDifferentInitializationOrder(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset1 = dataset_ops.Dataset.range(0, 100) |
| dataset2 = dataset_ops.Dataset.range(100, 200) |
| dataset3 = dataset_ops.Dataset.range(200, 300) |
| |
| dataset = dataset1.concatenate(dataset2).concatenate(dataset3) |
| dataset = dataset.apply(snapshot.legacy_snapshot(tmpdir)) |
| self.assertDatasetProduces(dataset, list(range(300))) |
| |
| dataset4 = dataset_ops.Dataset.range(200, 300) |
| dataset5 = dataset_ops.Dataset.range(100, 200) |
| dataset6 = dataset_ops.Dataset.range(0, 100) |
| |
| dataset = dataset6.concatenate(dataset5).concatenate(dataset4) |
| dataset = dataset.apply(snapshot.legacy_snapshot(tmpdir)) |
| self.assertDatasetProduces(dataset, list(range(300))) |
| |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testExpiredSnapshotRewrite(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset1 = dataset_ops.Dataset.range(1000) |
| dataset1 = dataset1.apply( |
| snapshot.legacy_snapshot(tmpdir, pending_snapshot_expiry_seconds=1)) |
| next1 = self.getNext(dataset1) |
| |
| # Don't finish reading dataset1, so it is never finalized |
| for _ in range(500): |
| self.evaluate(next1()) |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| time.sleep(2) |
| |
| # Creating dataset2 after we run through dataset1 due to eager mode, where |
| # the snapshot state is determined immediately upon dataset creation. We |
| # only want to determine the snapshot state for dataset2 after the first |
| # snapshot has expired. |
| dataset2 = dataset_ops.Dataset.range(1000) |
| dataset2 = dataset2.apply( |
| snapshot.legacy_snapshot(tmpdir, pending_snapshot_expiry_seconds=1)) |
| next2 = self.getNext(dataset2) |
| |
| for _ in range(500): |
| self.evaluate(next2()) |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 2, 1) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testSnapshotArgsCreateNewSnapshot(self): |
| tmpdir = self.snapshot_dir |
| |
| dataset1 = dataset_ops.Dataset.range(1000) |
| dataset1 = dataset1.apply( |
| snapshot.legacy_snapshot(tmpdir, shard_size_bytes=10000)) |
| next1 = self.getNext(dataset1) |
| |
| for _ in range(1000): |
| self.evaluate(next1()) |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, 1) |
| |
| # Create second snapshot with a different shard_size_bytes |
| dataset2 = dataset_ops.Dataset.range(1000) |
| dataset2 = dataset1.apply( |
| snapshot.legacy_snapshot(tmpdir, shard_size_bytes=20000)) |
| next2 = self.getNext(dataset2) |
| |
| for _ in range(1000): |
| self.evaluate(next2()) |
| self.assertSnapshotDirectoryContains(tmpdir, 2, 1, 1) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(compression=[ |
| snapshot.COMPRESSION_NONE, snapshot.COMPRESSION_GZIP, |
| snapshot.COMPRESSION_SNAPPY |
| ]))) |
| def testSpecifyShardSize(self, compression): |
| tmpdir = self.snapshot_dir |
| |
| dataset = dataset_ops.Dataset.from_tensor_slices([1.0]) |
| dataset = dataset.map(lambda x: gen_array_ops.broadcast_to(x, [1024, 1024])) |
| dataset = dataset.repeat(10) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot( |
| tmpdir, shard_size_bytes=10 * 1024 * 1024, compression=compression)) |
| next_fn = self.getNext(dataset) |
| |
| for _ in range(10): |
| self.evaluate(next_fn()) |
| |
| num_files = 1 |
| if compression == snapshot.COMPRESSION_NONE: |
| num_files = 3 |
| self.assertSnapshotDirectoryContains(tmpdir, 1, 1, num_files) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testAdditionalOperationsAfterReadBack(self): |
| self.setUpTFRecord() |
| filenames = self._filenames |
| |
| expected = [ |
| b"Record %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 10) |
| ] |
| |
| tmpdir = self.snapshot_dir |
| dataset = core_readers._TFRecordDataset(filenames) |
| dataset = dataset.apply(snapshot.legacy_snapshot(tmpdir)) |
| self.assertDatasetProduces(dataset, expected) |
| |
| # remove the original files and try to read the data back only from snapshot |
| self.removeTFRecords() |
| |
| dataset2 = core_readers._TFRecordDataset(filenames) |
| dataset2 = dataset2.apply(snapshot.legacy_snapshot(tmpdir)) |
| self.assertDatasetProduces(dataset2, expected) |
| |
| expected_after = [ |
| b"cord %d of file %d" % (r, f) # pylint:disable=g-complex-comprehension |
| for f in range(0, 10) |
| for r in range(0, 10) |
| ] |
| |
| dataset3 = core_readers._TFRecordDataset(filenames) |
| dataset3 = dataset3.apply(snapshot.legacy_snapshot(tmpdir)) |
| dataset3 = dataset3.map(lambda x: string_ops.substr_v2(x, 2, 1000)) |
| self.assertDatasetProduces(dataset3, expected_after) |
| |
| |
| class SnapshotCheckpointTest(checkpoint_test_base.CheckpointTestBase, |
| parameterized.TestCase): |
| |
| def _build_snapshot_dataset(self, repeat=False): |
| |
| def ds_fn(): |
| self._snapshot_dir = os.path.join(self.get_temp_dir(), "snapshot") |
| if not os.path.exists(self._snapshot_dir): |
| os.mkdir(self._snapshot_dir) |
| |
| dataset = dataset_ops.Dataset.range(100) |
| dataset = dataset.snapshot(path=self._snapshot_dir) |
| if repeat: |
| dataset = dataset.repeat(2) |
| return dataset |
| |
| return ds_fn |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testCheckpointBeforeEpochEndNoRepeat(self): |
| ds_fn = self._build_snapshot_dataset(repeat=False) |
| outputs = self.gen_outputs(ds_fn, [], 50, verify_exhausted=False) |
| self.assertSequenceEqual(outputs, range(50)) |
| outputs.extend( |
| self.gen_outputs(ds_fn, [], 50, ckpt_saved=True, verify_exhausted=True)) |
| self.assertSequenceEqual(outputs, range(100)) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testCheckpointBeforeOneEpochWithReading(self): |
| ds_fn = self._build_snapshot_dataset(repeat=True) |
| |
| # Generate 50 entries from iterator and save checkpoint. |
| outputs = self.gen_outputs(ds_fn, [], 50, verify_exhausted=False) |
| self.assertSequenceEqual(outputs, list(range(50))) |
| |
| # Restore from checkpoint and produce the rest of the elements from the |
| # iterator. |
| t = self.gen_outputs(ds_fn, [], 150, ckpt_saved=True, verify_exhausted=True) |
| outputs.extend(t) |
| self.assertSequenceEqual( |
| outputs, |
| list(range(50)) + list(range(50, 100)) + list(range(100))) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testCheckpointBeforeOneEpochThenRunAFewSteps(self): |
| ds_fn = self._build_snapshot_dataset(repeat=False) |
| outputs = self.gen_outputs( |
| ds_fn, [10], 20, verify_exhausted=False, save_checkpoint_at_end=False) |
| self.assertSequenceEqual(outputs, range(20)) |
| |
| outputs = outputs[:10] |
| outputs.extend( |
| self.gen_outputs(ds_fn, [], 90, ckpt_saved=True, verify_exhausted=True)) |
| self.assertSequenceEqual(outputs, range(100)) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testCheckpointAfterOneEpoch(self): |
| ds_fn = self._build_snapshot_dataset(repeat=True) |
| |
| # Generate 110 entries from iterator and save checkpoint. |
| outputs = self.gen_outputs(ds_fn, [], 110, verify_exhausted=False) |
| self.assertSequenceEqual(outputs, list(range(100)) + list(range(10))) |
| |
| # Restore from checkpoint and produce the rest of the elements from the |
| # iterator. |
| t = self.gen_outputs(ds_fn, [], 90, ckpt_saved=True, verify_exhausted=True) |
| outputs.extend(t) |
| self.assertSequenceEqual( |
| outputs, |
| list(range(100)) + list(range(10)) + list(range(10, 100))) |
| |
| @combinations.generate(test_base.default_test_combinations()) |
| def testCheckpointAfterOneEpochRunFewSteps(self): |
| ds_fn = self._build_snapshot_dataset(repeat=True) |
| |
| # Generate 120 entries from iterator and save checkpoint at 110. |
| outputs = self.gen_outputs( |
| ds_fn, [110], 120, verify_exhausted=False, save_checkpoint_at_end=False) |
| self.assertSequenceEqual(outputs, list(range(100)) + list(range(20))) |
| |
| # Restore from checkpoint and produce the rest of the elements from the |
| # iterator. |
| outputs = outputs[:110] |
| t = self.gen_outputs(ds_fn, [], 90, ckpt_saved=True, verify_exhausted=True) |
| outputs.extend(t) |
| self.assertSequenceEqual( |
| outputs, |
| list(range(100)) + list(range(10)) + list(range(10, 100))) |
| |
| |
| class LegacySnapshotCheckpointTest( |
| checkpoint_test_base.CheckpointTestBase, parameterized.TestCase): |
| |
| def _build_snapshot_dataset(self, |
| num_threads=1, |
| repeat=False, |
| pending_snapshot_expiry_seconds=-1, |
| shard_size_bytes=None): |
| |
| def ds_fn(): |
| self.snapshot_dir = os.path.join(self.get_temp_dir(), "snapshot") |
| if not os.path.exists(self.snapshot_dir): |
| os.mkdir(self.snapshot_dir) |
| dataset = dataset_ops.Dataset.range(1000) |
| dataset = dataset.apply( |
| snapshot.legacy_snapshot( |
| self.snapshot_dir, |
| num_writer_threads=num_threads, |
| writer_buffer_size=2 * num_threads, |
| num_reader_threads=num_threads, |
| reader_buffer_size=2 * num_threads, |
| pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds, |
| shard_size_bytes=shard_size_bytes)) |
| if repeat: |
| dataset = dataset.repeat(2) |
| return dataset |
| |
| return ds_fn |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(pending_snapshot_expiry_seconds=[None, 1]))) |
| def testSnapshotBeforeEpochEnd(self, pending_snapshot_expiry_seconds): |
| ds_fn = self._build_snapshot_dataset( |
| pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds) |
| outputs = self.gen_outputs(ds_fn, [], 100, verify_exhausted=False) |
| self.assertSequenceEqual(outputs, range(100)) |
| outputs.extend( |
| self.gen_outputs( |
| ds_fn, [], 900, ckpt_saved=True, verify_exhausted=False)) |
| self.assertSequenceEqual(outputs, range(1000)) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(pending_snapshot_expiry_seconds=[None, 1]))) |
| def testCheckpointBeforeOneEpochThenRunFewStepsSmallShardMultiThread( |
| self, pending_snapshot_expiry_seconds): |
| ds_fn = self._build_snapshot_dataset( |
| pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds, |
| shard_size_bytes=100) |
| |
| outputs = [] |
| with ops.Graph().as_default() as g: |
| init_op, get_next_op, saver = self._build_graph(ds_fn) |
| with self.session(graph=g) as sess: |
| self._initialize(init_op, sess) |
| start = 0 |
| end = 100 |
| num_iters = end - start |
| for _ in range(num_iters): |
| outputs.append(sess.run(get_next_op)) |
| self._save(sess, saver) |
| start = 100 |
| end = 400 |
| num_iters = end - start |
| for _ in range(num_iters): |
| outputs.append(sess.run(get_next_op)) |
| self.assertSequenceEqual(outputs, range(400)) |
| |
| outputs = outputs[:100] |
| outputs.extend( |
| self.gen_outputs( |
| ds_fn, [], 900, ckpt_saved=True, verify_exhausted=False)) |
| self.assertSequenceEqual(outputs, range(1000)) |
| fp_dir_list = os.listdir(self.snapshot_dir) |
| self.assertLen(list(fp_dir_list), 2) |
| for d in fp_dir_list: |
| if not d.endswith("-graph.pbtxt"): |
| fp_dir = os.path.join(self.snapshot_dir, d) |
| run_dir_list = os.listdir(fp_dir) |
| self.assertLen(list(run_dir_list), 2) |
| for e in run_dir_list: |
| if e != "snapshot.metadata": |
| run_dir = os.path.join(fp_dir, e) |
| self.assertLen(list(os.listdir(run_dir)), 258) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(pending_snapshot_expiry_seconds=[None, 1]))) |
| def testCheckpointBeforeOneEpochThenRunFewSteps( |
| self, pending_snapshot_expiry_seconds): |
| ds_fn = self._build_snapshot_dataset( |
| pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds) |
| |
| # Generate 200 entries from iterator but save checkpoint after producing |
| # 100. |
| outputs = self.gen_outputs( |
| ds_fn, [100], 200, verify_exhausted=False, save_checkpoint_at_end=False) |
| self.assertSequenceEqual(outputs, range(200)) |
| |
| outputs = outputs[:100] |
| outputs.extend( |
| self.gen_outputs( |
| ds_fn, [], 900, ckpt_saved=True, verify_exhausted=False)) |
| self.assertSequenceEqual(outputs, range(1000)) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(pending_snapshot_expiry_seconds=[None, 1]))) |
| def testCheckpointBeforeOneEpochThenRunFewStepsMultipleThreads( |
| self, pending_snapshot_expiry_seconds): |
| ds_fn = self._build_snapshot_dataset( |
| num_threads=2, |
| pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds) |
| |
| # Generate 200 entries from iterator but save checkpoint after producing |
| # 100. |
| outputs = self.gen_outputs( |
| ds_fn, [100], 200, verify_exhausted=False, save_checkpoint_at_end=False) |
| self.assertSequenceEqual(outputs, range(200)) |
| |
| outputs = outputs[:100] |
| outputs.extend( |
| self.gen_outputs( |
| ds_fn, [], 900, ckpt_saved=True, verify_exhausted=False)) |
| self.assertSequenceEqual(outputs, range(1000)) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(pending_snapshot_expiry_seconds=[None, 1]))) |
| def testCheckpointAfterOneEpoch(self, pending_snapshot_expiry_seconds): |
| ds_fn = self._build_snapshot_dataset( |
| repeat=True, |
| pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds) |
| |
| # Generate 1100 entries from iterator and save checkpoint. |
| outputs = self.gen_outputs(ds_fn, [], 1100, verify_exhausted=False) |
| self.assertSequenceEqual(outputs, list(range(1000)) + list(range(100))) |
| |
| # Restore from checkpoint and produce the rest of the elements from the |
| # iterator. |
| t = self.gen_outputs( |
| ds_fn, [], 900, ckpt_saved=True, verify_exhausted=False) |
| outputs.extend(t) |
| self.assertSequenceEqual( |
| outputs, |
| list(range(1000)) + list(range(100)) + list(range(900))) |
| |
| @combinations.generate( |
| combinations.times( |
| test_base.default_test_combinations(), |
| combinations.combine(pending_snapshot_expiry_seconds=[None, 1]))) |
| def testCheckpointAfterOneEpochThenRunFewSteps( |
| self, pending_snapshot_expiry_seconds): |
| ds_fn = self._build_snapshot_dataset( |
| repeat=True, |
| pending_snapshot_expiry_seconds=pending_snapshot_expiry_seconds) |
| |
| # Generate 200 entries from iterator but save checkpoint after producing |
| # 100. |
| outputs = self.gen_outputs( |
| ds_fn, [1100], |
| 1200, |
| verify_exhausted=False, |
| save_checkpoint_at_end=False) |
| self.assertSequenceEqual( |
| outputs, |
| list(range(1000)) + list(range(100)) + list(range(100))) |
| |
| outputs = outputs[:1100] |
| t = self.gen_outputs( |
| ds_fn, [], 900, ckpt_saved=True, verify_exhausted=False) |
| outputs.extend(t) |
| self.assertSequenceEqual( |
| outputs, (list(range(1000)) + list(range(100)) + list(range(900)))) |
| |
| |
| if __name__ == "__main__": |
| test.main() |
