import hypothesis.strategies as st from hypothesis import given, assume, settings import io import math import numpy as np import os import struct import unittest from pathlib import Path from typing import Dict, Generator, List, NamedTuple, Optional, Tuple, Type from caffe2.proto import caffe2_pb2 from caffe2.proto.caffe2_pb2 import BlobSerializationOptions from caffe2.python import core, test_util, workspace if workspace.has_gpu_support: DEVICES = [caffe2_pb2.CPU, workspace.GpuDeviceType] max_gpuid = workspace.NumGpuDevices() - 1 else: DEVICES = [caffe2_pb2.CPU] max_gpuid = 0 class MiniDBEntry(NamedTuple): key: str value_size: int # Utility class for other loading tests, don't add test functions here # Inherit from this test instead. If you add a test here, # each derived class will inherit it as well and cause test duplication class TestLoadSaveBase(test_util.TestCase): def __init__(self, methodName, db_type='minidb'): super(TestLoadSaveBase, self).__init__(methodName) self._db_type = db_type @settings(deadline=None) @given(src_device_type=st.sampled_from(DEVICES), src_gpu_id=st.integers(min_value=0, max_value=max_gpuid), dst_device_type=st.sampled_from(DEVICES), dst_gpu_id=st.integers(min_value=0, max_value=max_gpuid)) def load_save(self, src_device_type, src_gpu_id, dst_device_type, dst_gpu_id): workspace.ResetWorkspace() dtypes = [np.float16, np.float32, np.float64, np.bool, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16] arrays = [np.random.permutation(6).reshape(2, 3).astype(T) for T in dtypes] assume(core.IsGPUDeviceType(src_device_type) or src_gpu_id == 0) assume(core.IsGPUDeviceType(dst_device_type) or dst_gpu_id == 0) src_device_option = core.DeviceOption( src_device_type, src_gpu_id) dst_device_option = core.DeviceOption( dst_device_type, dst_gpu_id) for i, arr in enumerate(arrays): self.assertTrue(workspace.FeedBlob(str(i), arr, src_device_option)) self.assertTrue(workspace.HasBlob(str(i))) # Saves the blobs to a local db. tmp_folder = self.make_tempdir() op = core.CreateOperator( "Save", [str(i) for i in range(len(arrays))], [], absolute_path=1, db=str(tmp_folder / "db"), db_type=self._db_type) self.assertTrue(workspace.RunOperatorOnce(op)) # Reset the workspace so that anything we load is surely loaded # from the serialized proto. workspace.ResetWorkspace() self.assertEqual(len(workspace.Blobs()), 0) def _LoadTest(keep_device, device_type, gpu_id, blobs, loadAll): """A helper subfunction to test keep and not keep.""" op = core.CreateOperator( "Load", [], blobs, absolute_path=1, db=str(tmp_folder / "db"), db_type=self._db_type, device_option=dst_device_option, keep_device=keep_device, load_all=loadAll) self.assertTrue(workspace.RunOperatorOnce(op)) for i, arr in enumerate(arrays): self.assertTrue(workspace.HasBlob(str(i))) fetched = workspace.FetchBlob(str(i)) self.assertEqual(fetched.dtype, arr.dtype) np.testing.assert_array_equal( workspace.FetchBlob(str(i)), arr) proto = caffe2_pb2.BlobProto() proto.ParseFromString(workspace.SerializeBlob(str(i))) self.assertTrue(proto.HasField('tensor')) self.assertEqual(proto.tensor.device_detail.device_type, device_type) if core.IsGPUDeviceType(device_type): self.assertEqual(proto.tensor.device_detail.device_id, gpu_id) blobs = [str(i) for i in range(len(arrays))] # Load using device option stored in the proto, i.e. # src_device_option _LoadTest(1, src_device_type, src_gpu_id, blobs, 0) # Load again, but this time load into dst_device_option. _LoadTest(0, dst_device_type, dst_gpu_id, blobs, 0) # Load back to the src_device_option to see if both paths are able # to reallocate memory. _LoadTest(1, src_device_type, src_gpu_id, blobs, 0) # Reset the workspace, and load directly into the dst_device_option. workspace.ResetWorkspace() _LoadTest(0, dst_device_type, dst_gpu_id, blobs, 0) # Test load all which loads all blobs in the db into the workspace. workspace.ResetWorkspace() _LoadTest(1, src_device_type, src_gpu_id, [], 1) # Load again making sure that overwrite functionality works. _LoadTest(1, src_device_type, src_gpu_id, [], 1) # Load again with different device. _LoadTest(0, dst_device_type, dst_gpu_id, [], 1) workspace.ResetWorkspace() _LoadTest(0, dst_device_type, dst_gpu_id, [], 1) workspace.ResetWorkspace() _LoadTest(1, src_device_type, src_gpu_id, blobs, 1) workspace.ResetWorkspace() _LoadTest(0, dst_device_type, dst_gpu_id, blobs, 1) def saveFile( self, tmp_folder: Path, db_name: str, db_type: str, start_blob_id: int ) -> Tuple[str, List[np.ndarray]]: dtypes = [np.float16, np.float32, np.float64, np.bool, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16] arrays = [np.random.permutation(6).reshape(2, 3).astype(T) for T in dtypes] for i, arr in enumerate(arrays): self.assertTrue(workspace.FeedBlob(str(i + start_blob_id), arr)) self.assertTrue(workspace.HasBlob(str(i + start_blob_id))) # Saves the blobs to a local db. tmp_file = str(tmp_folder / db_name) op = core.CreateOperator( "Save", [str(i + start_blob_id) for i in range(len(arrays))], [], absolute_path=1, db=tmp_file, db_type=db_type) workspace.RunOperatorOnce(op) return tmp_file, arrays class TestLoadSave(TestLoadSaveBase): def testLoadSave(self): self.load_save() def testRepeatedArgs(self): dtypes = [np.float16, np.float32, np.float64, np.bool, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16] arrays = [np.random.permutation(6).reshape(2, 3).astype(T) for T in dtypes] for i, arr in enumerate(arrays): self.assertTrue(workspace.FeedBlob(str(i), arr)) self.assertTrue(workspace.HasBlob(str(i))) # Saves the blobs to a local db. tmp_folder = self.make_tempdir() op = core.CreateOperator( "Save", [str(i) for i in range(len(arrays))] * 2, [], absolute_path=1, db=str(tmp_folder / "db"), db_type=self._db_type) with self.assertRaises(RuntimeError): workspace.RunOperatorOnce(op) def testLoadExcessblobs(self): tmp_folder = self.make_tempdir() tmp_file, arrays = self.saveFile(tmp_folder, "db", self._db_type, 0) op = core.CreateOperator( "Load", [], [str(i) for i in range(len(arrays))] * 2, absolute_path=1, db=tmp_file, db_type=self._db_type, load_all=False) with self.assertRaises(RuntimeError): workspace.RunOperatorOnce(op) op = core.CreateOperator( "Load", [], [str(len(arrays) + i) for i in [-1, 0]], absolute_path=1, db=tmp_file, db_type=self._db_type, load_all=True) with self.assertRaises(RuntimeError): workspace.ResetWorkspace() workspace.RunOperatorOnce(op) op = core.CreateOperator( "Load", [], [str(len(arrays) + i) for i in range(2)], absolute_path=1, db=tmp_file, db_type=self._db_type, load_all=True) with self.assertRaises(RuntimeError): workspace.ResetWorkspace() workspace.RunOperatorOnce(op) def testTruncatedFile(self): tmp_folder = self.make_tempdir() tmp_file, arrays = self.saveFile(tmp_folder, "db", self._db_type, 0) with open(tmp_file, 'wb+') as fdest: fdest.seek(20, os.SEEK_END) fdest.truncate() op = core.CreateOperator( "Load", [], [str(i) for i in range(len(arrays))], absolute_path=1, db=tmp_file, db_type=self._db_type, load_all=False) with self.assertRaises(RuntimeError): workspace.RunOperatorOnce(op) op = core.CreateOperator( "Load", [], [], absolute_path=1, db=tmp_file, db_type=self._db_type, load_all=True) with self.assertRaises(RuntimeError): workspace.RunOperatorOnce(op) def testBlobNameOverrides(self): original_names = ['blob_a', 'blob_b', 'blob_c'] new_names = ['x', 'y', 'z'] blobs = [np.random.permutation(6) for i in range(3)] for i, blob in enumerate(blobs): self.assertTrue(workspace.FeedBlob(original_names[i], blob)) self.assertTrue(workspace.HasBlob(original_names[i])) self.assertEqual(len(workspace.Blobs()), 3) # Saves the blobs to a local db. tmp_folder = self.make_tempdir() with self.assertRaises(RuntimeError): workspace.RunOperatorOnce( core.CreateOperator( "Save", original_names, [], absolute_path=1, strip_prefix='.temp', blob_name_overrides=new_names, db=str(tmp_folder / "db"), db_type=self._db_type ) ) self.assertTrue( workspace.RunOperatorOnce( core.CreateOperator( "Save", original_names, [], absolute_path=1, blob_name_overrides=new_names, db=str(tmp_folder / "db"), db_type=self._db_type ) ) ) self.assertTrue(workspace.ResetWorkspace()) self.assertEqual(len(workspace.Blobs()), 0) self.assertTrue( workspace.RunOperatorOnce( core.CreateOperator( "Load", [], [], absolute_path=1, db=str(tmp_folder / "db"), db_type=self._db_type, load_all=1 ) ) ) self.assertEqual(len(workspace.Blobs()), 3) for i, name in enumerate(new_names): self.assertTrue(workspace.HasBlob(name)) self.assertTrue((workspace.FetchBlob(name) == blobs[i]).all()) # moved here per @cxj's suggestion load_new_names = ['blob_x', 'blob_y', 'blob_z'] # load 'x' into 'blob_x' self.assertTrue( workspace.RunOperatorOnce( core.CreateOperator( "Load", [], load_new_names[0:1], absolute_path=1, db=str(tmp_folder / "db"), db_type=self._db_type, source_blob_names=new_names[0:1] ) ) ) # we should have 'blob_a/b/c/' and 'blob_x' now self.assertEqual(len(workspace.Blobs()), 4) for i, name in enumerate(load_new_names[0:1]): self.assertTrue(workspace.HasBlob(name)) self.assertTrue((workspace.FetchBlob(name) == blobs[i]).all()) self.assertTrue( workspace.RunOperatorOnce( core.CreateOperator( "Load", [], load_new_names[0:3], absolute_path=1, db=str(tmp_folder / "db"), db_type=self._db_type, source_blob_names=new_names[0:3] ) ) ) # we should have 'blob_a/b/c/' and 'blob_x/y/z' now self.assertEqual(len(workspace.Blobs()), 6) for i, name in enumerate(load_new_names[0:3]): self.assertTrue(workspace.HasBlob(name)) self.assertTrue((workspace.FetchBlob(name) == blobs[i]).all()) def testMissingFile(self): tmp_folder = self.make_tempdir() tmp_file = tmp_folder / "missing_db" op = core.CreateOperator( "Load", [], [], absolute_path=1, db=str(tmp_file), db_type=self._db_type, load_all=True) with self.assertRaises(RuntimeError): try: workspace.RunOperatorOnce(op) except RuntimeError as e: print(e) raise def testLoadMultipleFilesGivenSourceBlobNames(self): tmp_folder = self.make_tempdir() db_file_1, arrays_1 = self.saveFile(tmp_folder, "db1", self._db_type, 0) db_file_2, arrays_2 = self.saveFile( tmp_folder, "db2", self._db_type, len(arrays_1) ) db_files = [db_file_1, db_file_2] blobs_names = [str(i) for i in range(len(arrays_1) + len(arrays_2))] workspace.ResetWorkspace() self.assertEqual(len(workspace.Blobs()), 0) self.assertTrue( workspace.RunOperatorOnce( core.CreateOperator( "Load", [], blobs_names, absolute_path=1, dbs=db_files, db_type=self._db_type, source_blob_names=blobs_names ) ) ) self.assertEqual(len(workspace.Blobs()), len(blobs_names)) for i in range(len(arrays_1)): np.testing.assert_array_equal( workspace.FetchBlob(str(i)), arrays_1[i] ) for i in range(len(arrays_2)): np.testing.assert_array_equal( workspace.FetchBlob(str(i + len(arrays_1))), arrays_2[i] ) def testLoadAllMultipleFiles(self): tmp_folder = self.make_tempdir() db_file_1, arrays_1 = self.saveFile(tmp_folder, "db1", self._db_type, 0) db_file_2, arrays_2 = self.saveFile( tmp_folder, "db2", self._db_type, len(arrays_1) ) db_files = [db_file_1, db_file_2] workspace.ResetWorkspace() self.assertEqual(len(workspace.Blobs()), 0) self.assertTrue( workspace.RunOperatorOnce( core.CreateOperator( "Load", [], [], absolute_path=1, dbs=db_files, db_type=self._db_type, load_all=True ) ) ) self.assertEqual(len(workspace.Blobs()), len(arrays_1) + len(arrays_2)) for i in range(len(arrays_1)): np.testing.assert_array_equal( workspace.FetchBlob(str(i)), arrays_1[i] ) for i in range(len(arrays_2)): np.testing.assert_array_equal( workspace.FetchBlob(str(i + len(arrays_1))), arrays_2[i] ) def testLoadAllMultipleFilesWithSameKey(self): tmp_folder = self.make_tempdir() db_file_1, arrays_1 = self.saveFile(tmp_folder, "db1", self._db_type, 0) db_file_2, arrays_2 = self.saveFile(tmp_folder, "db2", self._db_type, 0) db_files = [db_file_1, db_file_2] workspace.ResetWorkspace() self.assertEqual(len(workspace.Blobs()), 0) op = core.CreateOperator( "Load", [], [], absolute_path=1, dbs=db_files, db_type=self._db_type, load_all=True) with self.assertRaises(RuntimeError): workspace.RunOperatorOnce(op) def testLoadRepeatedFiles(self): tmp_folder = self.make_tempdir() tmp_file, arrays = self.saveFile(tmp_folder, "db", self._db_type, 0) db_files = [tmp_file, tmp_file] workspace.ResetWorkspace() self.assertEqual(len(workspace.Blobs()), 0) op = core.CreateOperator( "Load", [], [str(i) for i in range(len(arrays))], absolute_path=1, dbs=db_files, db_type=self._db_type, load_all=False) with self.assertRaises(RuntimeError): workspace.RunOperatorOnce(op) def testLoadWithDBOptions(self) -> None: tmp_folder = self.make_tempdir() tmp_file, arrays = self.saveFile(tmp_folder, "db", self._db_type, 0) db_files = [tmp_file, tmp_file] workspace.ResetWorkspace() self.assertEqual(len(workspace.Blobs()), 0) db_options = b"test_db_options" op = core.CreateOperator( "Load", [], [str(i) for i in range(len(arrays))], absolute_path=1, dbs=db_files, db_type=self._db_type, load_all=False, db_options=db_options, ) with self.assertRaises(RuntimeError): workspace.RunOperatorOnce(op) def create_test_blobs( self, size: int = 1234, feed: bool = True ) -> List[Tuple[str, np.ndarray]]: def int_array(dtype: Type[np.integer], size: int) -> np.ndarray: info = np.iinfo(dtype) return np.random.randint(info.min, info.max, size, dtype=dtype) def float_array(dtype: Type[np.floating], size: int) -> np.ndarray: return np.random.random_sample(size).astype(dtype) blobs = [ ("int8_data", int_array(np.int8, size)), ("int16_data", int_array(np.int16, size)), ("int32_data", int_array(np.int32, size)), ("int64_data", int_array(np.int64, size)), ("uint8_data", int_array(np.uint8, size)), ("uint16_data", int_array(np.uint16, size)), ("float16_data", float_array(np.float16, size)), ("float32_data", float_array(np.float32, size)), ("float64_data", float_array(np.float64, size)), ] if feed: for name, data in blobs: workspace.FeedBlob(name, data) return blobs def load_blobs( self, blob_names: List[str], dbs: List[str], db_type: Optional[str] = None ) -> None: workspace.ResetWorkspace() self.assertEqual(len(workspace.Blobs()), 0) load_op = core.CreateOperator( "Load", [], blob_names, absolute_path=1, dbs=dbs, db_type=db_type or self._db_type, ) self.assertTrue(workspace.RunOperatorOnce(load_op)) self.assertEqual(len(workspace.Blobs()), len(blob_names)) def load_and_check_blobs( self, blobs: List[Tuple[str, np.ndarray]], dbs: List[str], db_type: Optional[str] = None ) -> None: self.load_blobs([name for name, data in blobs], dbs, db_type) for name, data in blobs: np.testing.assert_array_equal(workspace.FetchBlob(name), data) def _read_minidb_entries( self, path: Path ) -> Generator[MiniDBEntry, None, None]: """Read the entry information out of a minidb file. """ header = struct.Struct("=ii") with path.open("rb") as f: while True: buf = f.read(header.size) if not buf: break if len(buf) < header.size: raise Exception("early EOF in minidb header") (key_len, value_len) = header.unpack(buf) if key_len < 0 or value_len < 0: raise Exception( f"invalid minidb header: ({key_len}, {value_len})" ) key = f.read(key_len) if len(key) < key_len: raise Exception("early EOF in minidb key") f.seek(value_len, io.SEEK_CUR) yield MiniDBEntry(key=key.decode("utf-8"), value_size=value_len) def _read_chunk_info(self, path: Path) -> Dict[str, List[MiniDBEntry]]: """Read a minidb file and return the names of each blob and how many chunks are stored for that blob. """ chunk_id_separator = "#%" results: Dict[str, List[MiniDBEntry]] = {} for entry in self._read_minidb_entries(path): parts = entry.key.rsplit(chunk_id_separator, 1) if len(parts) == 0: assert entry.key not in results results[entry.key] = [entry] else: blob_name = parts[0] results.setdefault(blob_name, []) results[blob_name].append(entry) return results def _test_save_with_chunk_size( self, num_elems: int, chunk_size: int, expected_num_chunks: int, ) -> None: tmp_folder = self.make_tempdir() tmp_file = str(tmp_folder / "save.output") blobs = self.create_test_blobs(num_elems) # Saves the blobs to a local db. save_op = core.CreateOperator( "Save", [name for name, data in blobs], [], absolute_path=1, db=tmp_file, db_type=self._db_type, chunk_size=chunk_size, ) self.assertTrue(workspace.RunOperatorOnce(save_op)) self.load_and_check_blobs(blobs, [tmp_file]) blob_chunks = self._read_chunk_info(Path(tmp_file)) for blob_name, chunks in blob_chunks.items(): self.assertEqual(len(chunks), expected_num_chunks) def testSaveWithChunkSize(self) -> None: num_elems = 1234 chunk_size = 32 expected_num_chunks = math.ceil(num_elems / chunk_size) self._test_save_with_chunk_size( num_elems=num_elems, chunk_size=chunk_size, expected_num_chunks=expected_num_chunks, ) def testSaveWithDefaultChunkSize(self) -> None: # This is the default value of the --caffe2_tensor_chunk_size flag from # core/blob_serialization.cc # # Test with just slightly more than this to ensure that 2 chunks are # used. default_chunk_size = 1000000 self._test_save_with_chunk_size( num_elems=default_chunk_size + 10, chunk_size=-1, expected_num_chunks=2, ) def testSaveWithNoChunking(self) -> None: default_chunk_size = 1000000 self._test_save_with_chunk_size( num_elems=default_chunk_size + 10, chunk_size=0, expected_num_chunks=1, ) def testSaveWithOptions(self) -> None: tmp_folder = self.make_tempdir() tmp_file = str(tmp_folder / "save.output") num_elems = 1234 blobs = self.create_test_blobs(num_elems) # Saves the blobs to a local db. save_op = core.CreateOperator( "Save", [name for name, data in blobs], [], absolute_path=1, db=tmp_file, db_type=self._db_type, chunk_size=40, options=caffe2_pb2.SerializationOptions( options=[ BlobSerializationOptions( blob_name_regex="int16_data", chunk_size=10 ), BlobSerializationOptions( blob_name_regex=".*16_data", chunk_size=20 ), BlobSerializationOptions( blob_name_regex="float16_data", chunk_size=30 ), ], ), ) self.assertTrue(workspace.RunOperatorOnce(save_op)) self.load_and_check_blobs(blobs, [tmp_file]) blob_chunks = self._read_chunk_info(Path(tmp_file)) # We explicitly set a chunk_size of 10 for int16_data self.assertEqual( len(blob_chunks["int16_data"]), math.ceil(num_elems / 10) ) # uint16_data should match the .*16_data pattern, and get a size of 20 self.assertEqual( len(blob_chunks["uint16_data"]), math.ceil(num_elems / 20) ) # float16_data should also match the .*16_data pattern, and get a size # of 20. The explicitly float16_data rule came after the .*16_data # pattern, so it has lower precedence and will be ignored. self.assertEqual( len(blob_chunks["float16_data"]), math.ceil(num_elems / 20) ) # int64_data will get the default chunk_size of 40 self.assertEqual( len(blob_chunks["int64_data"]), math.ceil(num_elems / 40) ) def testSaveWithDBOptions(self) -> None: num_elems = 1234 chunk_size = 32 expected_num_chunks = math.ceil(num_elems / chunk_size) tmp_folder = self.make_tempdir() tmp_file = str(tmp_folder / "save.output") blobs = self.create_test_blobs(num_elems) db_options = b"test_db_options" # Saves the blobs to a local db. save_op = core.CreateOperator( "Save", [name for name, data in blobs], [], absolute_path=1, db=tmp_file, db_type=self._db_type, chunk_size=chunk_size, db_options=db_options, ) self.assertTrue(workspace.RunOperatorOnce(save_op)) self.load_and_check_blobs(blobs, [tmp_file]) blob_chunks = self._read_chunk_info(Path(tmp_file)) for blob_name, chunks in blob_chunks.items(): self.assertEqual(len(chunks), expected_num_chunks) def testSaveFloatToBfloat16(self) -> None: tmp_folder = self.make_tempdir() tmp_file = str(tmp_folder / "save.output") # Create 2 blobs with the same float data float_data = np.random.random_sample(4000).astype(np.float32) workspace.FeedBlob("float1", float_data) workspace.FeedBlob("float2", float_data) blob_names = ["float1", "float2"] # Serialize the data, using bfloat16 serialization for one of the blobs save_op = core.CreateOperator( "Save", blob_names, [], absolute_path=1, db=tmp_file, db_type=self._db_type, options=caffe2_pb2.SerializationOptions( options=[ BlobSerializationOptions( blob_name_regex="float1", float_format=BlobSerializationOptions.FLOAT_BFLOAT16, ), ], ), ) self.assertTrue(workspace.RunOperatorOnce(save_op)) # As long as fbgemm was available for us to perform bfloat16 conversion, # the serialized data for float1 should be almost half the size of float2 if workspace.has_fbgemm: blob_chunks = self._read_chunk_info(Path(tmp_file)) self.assertEqual(len(blob_chunks["float1"]), 1, blob_chunks["float1"]) self.assertEqual(len(blob_chunks["float2"]), 1, blob_chunks["float2"]) self.assertLess( blob_chunks["float1"][0].value_size, 0.6 * blob_chunks["float2"][0].value_size ) self.load_blobs(blob_names, [tmp_file]) # float2 should be exactly the same as the input data np.testing.assert_array_equal(workspace.FetchBlob("float2"), float_data) # float2 should be close-ish to the input data np.testing.assert_array_almost_equal( workspace.FetchBlob("float1"), float_data, decimal=2 ) def testEstimateBlobSizes(self) -> None: # Create some blobs to test with float_data = np.random.random_sample(4000).astype(np.float32) workspace.FeedBlob("float1", float_data) workspace.FeedBlob("float2", float_data) workspace.FeedBlob( "float3", np.random.random_sample(2).astype(np.float32) ) workspace.FeedBlob( "ui16", np.random.randint(0, 0xffff, size=1024, dtype=np.uint16) ) # Estimate the serialized size of the data. # Request bfloat16 serialization for one of the float blobs, just to # exercise size estimation when using this option. options = caffe2_pb2.SerializationOptions( options=[ BlobSerializationOptions( blob_name_regex="float1", float_format=BlobSerializationOptions.FLOAT_BFLOAT16, chunk_size=500, ), ], ) get_blobs_op = core.CreateOperator( "EstimateAllBlobSizes", [], ["blob_names", "blob_sizes"], options=options, ) self.assertTrue(workspace.RunOperatorOnce(get_blobs_op)) blob_names = workspace.FetchBlob("blob_names") blob_sizes = workspace.FetchBlob("blob_sizes") sizes_by_name: Dict[str, int] = {} for idx, name in enumerate(blob_names): sizes_by_name[name.decode("utf-8")] = blob_sizes[idx] # Note that the output blob list will include our output blob names. expected_blobs = [ "float1", "float2", "float3", "ui16", "blob_names", "blob_sizes" ] self.assertEqual(set(sizes_by_name.keys()), set(expected_blobs)) def check_expected_blob_size( name: str, num_elems: int, elem_size: int, num_chunks: int = 1 ) -> None: # The estimation code applies a fixed 40 byte per-chunk overhead to # account for the extra space required for other fixed TensorProto # message fields. per_chunk_overhead = 50 expected_size = ( (num_chunks * (len(name) + per_chunk_overhead)) + (num_elems * elem_size) ) self.assertEqual( sizes_by_name[name], expected_size, f"expected size mismatch for {name}" ) check_expected_blob_size("ui16", 1024, 3) check_expected_blob_size("float2", 4000, 4) check_expected_blob_size("float3", 2, 4) # Our serialization options request to split float1 into 500-element # chunks when saving it. If fbgemm is available then the float1 blob # will be serialized using 2 bytes per element instead of 4 bytes. float1_num_chunks = 4000 // 500 if workspace.has_fbgemm: check_expected_blob_size("float1", 4000, 2, float1_num_chunks) else: check_expected_blob_size("float1", 4000, 4, float1_num_chunks) check_expected_blob_size("blob_names", len(expected_blobs), 50) check_expected_blob_size("blob_sizes", len(expected_blobs), 8) # Now actually save the blobs so we can compare our estimates # to how big the serialized data actually is. tmp_folder = self.make_tempdir() tmp_file = str(tmp_folder / "save.output") save_op = core.CreateOperator( "Save", list(sizes_by_name.keys()), [], absolute_path=1, db=tmp_file, db_type=self._db_type, options=options, ) self.assertTrue(workspace.RunOperatorOnce(save_op)) blob_chunks = self._read_chunk_info(Path(tmp_file)) saved_sizes: Dict[str, int] = {} for blob_name, chunks in blob_chunks.items(): total_size = sum(chunk.value_size for chunk in chunks) saved_sizes[blob_name] = total_size # For sanity checking, ensure that our estimates aren't # extremely far off for name in expected_blobs: estimated_size = sizes_by_name[name] saved_size = saved_sizes[name] difference = abs(estimated_size - saved_size) error_pct = 100.0 * (difference / saved_size) print( f"{name}: estimated={estimated_size} actual={saved_size} " f"error={error_pct:.2f}%" ) # Don't check the blob_names blob. It is a string tensor, and we # can't estimate string tensor sizes very well without knowing the # individual string lengths. (Currently it requires 102 bytes to # save, but we estimate 360). if name == "blob_names": continue # Check that we are within 100 bytes, or within 25% # We are generally quite close for tensors with fixed-width fields # (like float), but a little farther off for tensors that use varint # encoding. if difference > 100: self.assertLess(error_pct, 25.0) if __name__ == '__main__': unittest.main()