def test_non_grouped_weights_throws(self): groups = [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'float32') }] with self.assertRaises(Exception): write_weights(groups, TMP_DIR)
def test_bad_numpy_array_dtype_throws(self): groups = [ [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'float64') }] ] with self.assertRaises(Exception): write_weights(groups, TMP_DIR)
def test_bad_weights_entry_throws_no_data(self): groups = [ [{ 'name': 'weight1', 'nodata': np.array([1, 2, 3], 'float32') }] ] with self.assertRaises(Exception): write_weights(groups, TMP_DIR)
def test_duplicate_weight_name_throws(self): groups = [ [{ 'name': 'duplicate', 'data': np.array([1, 2, 3], 'float32') }], [{ 'name': 'duplicate', 'data': np.array([4, 5, 6], 'float32') }] ] with self.assertRaises(Exception): write_weights(groups, TMP_DIR)
def test_1_group_1_weight_sharded(self): groups = [ [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'float32') }] ] # Shard size is smaller than the size of the array so it gets split between # multiple files. manifest_json = write_weights(groups, TMP_DIR, shard_size_bytes = 2 * 4) manifest = json.loads(manifest_json) self.assertTrue( os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')), 'weights_manifest.json does not exist') self.assertEqual( manifest, [{ 'paths': ['group1-shard1of2', 'group1-shard2of2'], 'weights': [{ 'name': 'weight1', 'shape': [3], 'dtype': 'float32' }] }]) shard_1_path = os.path.join(TMP_DIR, 'group1-shard1of2') shard_1 = np.fromfile(shard_1_path, 'float32') np.testing.assert_array_equal(shard_1, np.array([1, 2], 'float32')) shard_2_path = os.path.join(TMP_DIR, 'group1-shard2of2') shard_2 = np.fromfile(shard_2_path, 'float32') np.testing.assert_array_equal(shard_2, np.array([3], 'float32'))
def test_1_group_1_weight(self): groups = [ [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'float32') }] ] manifest_json = write_weights(groups, TMP_DIR, shard_size_bytes = 4 * 4) manifest = json.loads(manifest_json) self.assertTrue( os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')), 'weights_manifest.json does not exist') self.assertEqual( manifest, [{ 'paths': ['group1-shard1of1'], 'weights': [{ 'name': 'weight1', 'shape': [3], 'dtype': 'float32' }] }]) weights_path = os.path.join(TMP_DIR, 'group1-shard1of1') weight1 = np.fromfile(weights_path, 'float32') np.testing.assert_array_equal(weight1, np.array([1, 2, 3], 'float32'))
def test_no_write_manfest(self): groups = [ [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'float32') }] ] manifest_json = write_weights(groups, TMP_DIR, write_manifest=False) manifest = json.loads(manifest_json) self.assertFalse( os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')), 'weights_manifest.json exists, but expected not to exist') self.assertEqual( manifest, [{ 'paths': ['group1-shard1of1'], 'weights': [{ 'name': 'weight1', 'shape': [3], 'dtype': 'float32' }] }]) weights_path = os.path.join(TMP_DIR, 'group1-shard1of1') weight1 = np.fromfile(weights_path, 'float32') np.testing.assert_array_equal(weight1, np.array([1, 2, 3], 'float32'))
def test_1_group_2_packed_sharded_multi_dtype(self): groups = [ [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'int32') }, { 'name': 'weight2', 'data': np.array([4.1, 5.1], 'float32') }] ] # Shard size is smaller than the sum of the weights so they get packed and # then sharded. The two buffers will get sharded into 3 files, with shapes # [2], [2], and [1]. The second shard is a mixed dtype. manifest_json = write_weights(groups, TMP_DIR, shard_size_bytes = 2 * 4) manifest = json.loads(manifest_json) self.assertTrue( os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')), 'weights_manifest.json does not exist') self.assertEqual( manifest, [{ 'paths': ['group1-shard1of3', 'group1-shard2of3', 'group1-shard3of3'], 'weights': [{ 'name': 'weight1', 'shape': [3], 'dtype': 'int32' }, { 'name': 'weight2', 'shape': [2], 'dtype': 'float32' }] }]) shard_1_path = os.path.join(TMP_DIR, 'group1-shard1of3') shard_1 = np.fromfile(shard_1_path, 'int32') np.testing.assert_array_equal(shard_1, np.array([1, 2], 'int32')) # Shard 2 has a mixed dtype so we parse the bytes directly. shard_2_path = os.path.join(TMP_DIR, 'group1-shard2of3') with open(shard_2_path, 'rb') as f: shard_2_bytes = f.read() shard_2_int = np.frombuffer(shard_2_bytes[:4], 'int32') np.testing.assert_array_equal(shard_2_int, np.array([3], 'int32')) shard_2_float = np.frombuffer(shard_2_bytes[4:], 'float32') np.testing.assert_array_equal(shard_2_float, np.array([4.1], 'float32')) shard_3_path = os.path.join(TMP_DIR, 'group1-shard3of3') shard_3 = np.fromfile(shard_3_path, 'float32') np.testing.assert_array_equal(shard_3, np.array([5.1], 'float32'))
def test_1_group_2_weights_packed(self): groups = [ [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'float32') }, { 'name': 'weight2', 'data': np.array([4, 5], 'float32') }] ] # Shard size is larger than the sum of the two weights so they get packed. manifest_json = write_weights(groups, TMP_DIR, shard_size_bytes = 8 * 4) manifest = json.loads(manifest_json) self.assertTrue( os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')), 'weights_manifest.json does not exist') self.assertEqual( manifest, [{ 'paths': ['group1-shard1of1'], 'weights': [{ 'name': 'weight1', 'shape': [3], 'dtype': 'float32' }, { 'name': 'weight2', 'shape': [2], 'dtype': 'float32' }] }]) weights_path = os.path.join(TMP_DIR, 'group1-shard1of1') weights = np.fromfile(weights_path, 'float32') np.testing.assert_array_equal(weights, np.array([1, 2, 3, 4, 5], 'float32'))
def test_empty_groups_throws(self): groups = [] with self.assertRaises(Exception): write_weights(groups, TMP_DIR)
def test_no_weights_groups_throws(self): groups = None with self.assertRaises(Exception): write_weights(groups, TMP_DIR)
def test_2_groups_4_weights_sharded_packed(self): groups = [ # Group 1 [{ 'name': 'weight1', 'data': np.array([1, 2, 3], 'float32') }, { 'name': 'weight2', 'data': np.array([[4, 5], [6, 7]], 'float32') }], # Group 2 [{ 'name': 'weight3', 'data': np.array([1, 2, 3, 4], 'int32') }, { 'name': 'weight4', 'data': np.array([[1.1, 1.2], [1.3, 1.4], [1.5, 1.6]], 'float32') }] ] manifest_json = write_weights(groups, TMP_DIR, shard_size_bytes = 4 * 4) manifest = json.loads(manifest_json) self.assertTrue( os.path.isfile(os.path.join(TMP_DIR, 'weights_manifest.json')), 'weights_manifest.json does not exist') self.assertEqual( manifest, [{ 'paths': ['group1-shard1of2', 'group1-shard2of2'], 'weights': [{ 'name': 'weight1', 'shape': [3], 'dtype': 'float32' }, { 'name': 'weight2', 'shape': [2, 2], 'dtype': 'float32' }] }, { 'paths': ['group2-shard1of3', 'group2-shard2of3', 'group2-shard3of3'], 'weights': [{ 'name': 'weight3', 'shape': [4], 'dtype': 'int32' }, { 'name': 'weight4', 'shape': [3, 2], 'dtype': 'float32' }] }]) group0_shard_1_path = os.path.join(TMP_DIR, 'group1-shard1of2') group0_shard_1 = np.fromfile(group0_shard_1_path, 'float32') np.testing.assert_array_equal(group0_shard_1, np.array([1, 2, 3, 4], 'float32')) group0_shard_2_path = os.path.join(TMP_DIR, 'group1-shard2of2') group0_shard_2 = np.fromfile(group0_shard_2_path, 'float32') np.testing.assert_array_equal(group0_shard_2, np.array([5, 6, 7], 'float32')) group1_shard_1_path = os.path.join(TMP_DIR, 'group2-shard1of3') group1_shard_1 = np.fromfile(group1_shard_1_path, 'int32') np.testing.assert_array_equal(group1_shard_1, np.array([1, 2, 3, 4], 'int32')) group2_shard_2_path = os.path.join(TMP_DIR, 'group2-shard2of3') group2_shard_2 = np.fromfile(group2_shard_2_path, 'float32') np.testing.assert_array_equal(group2_shard_2, np.array([1.1, 1.2, 1.3, 1.4], 'float32')) group2_shard_3_path = os.path.join(TMP_DIR, 'group2-shard3of3') group2_shard_3 = np.fromfile(group2_shard_3_path, 'float32') np.testing.assert_array_equal(group2_shard_3, np.array([1.5, 1.6], 'float32'))