Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/keras/src/saving/pickle_utils.py: 26%

1# Copyright 2021 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15"""Saving utilities to support Python's Pickle protocol.""" 

16import os 

17import tempfile 

18 

19import tensorflow.compat.v2 as tf 

20 

21from keras.src.saving import saving_lib 

22 

23 

24def deserialize_model_from_bytecode(serialized_model): 

25 """Reconstruct a Model from the output of `serialize_model_as_bytecode`. 

26 

27 Args: 

28 serialized_model: (bytes) return value from 

29 `serialize_model_as_bytecode`. 

30 

31 Returns: 

32 Keras Model instance. 

33 """ 

34 # Note: we don't use a RAM path for this because zipfile cannot write 

35 # to such paths. 

36 temp_dir = tempfile.mkdtemp() 

37 try: 

38 filepath = os.path.join(temp_dir, "model.keras") 

39 with open(filepath, "wb") as f: 

40 f.write(serialized_model) 

41 # When loading, direct import will work for most custom objects 

42 # though it will require get_config() to be implemented. 

43 # Some custom objects (e.g. an activation in a Dense layer, 

44 # serialized as a string by Dense.get_config()) will require 

45 # a custom_object_scope. 

46 model = saving_lib.load_model(filepath, safe_mode=False) 

47 except Exception as e: 

48 raise e 

49 else: 

50 return model 

51 finally: 

52 tf.io.gfile.rmtree(temp_dir) 

53 

54 

55def serialize_model_as_bytecode(model): 

56 """Convert a Keras Model into a bytecode representation for pickling. 

57 

58 Args: 

59 model: Keras Model instance. 

60 

61 Returns: 

62 Tuple that can be read by `deserialize_from_bytecode`. 

63 """ 

64 # Note: we don't use a RAM path for this because zipfile cannot write 

65 # to such paths. 

66 temp_dir = tempfile.mkdtemp() 

67 try: 

68 filepath = os.path.join(temp_dir, "model.keras") 

69 saving_lib.save_model(model, filepath) 

70 with open(filepath, "rb") as f: 

71 data = f.read() 

72 except Exception as e: 

73 raise e 

74 else: 

75 return data 

76 finally: 

77 tf.io.gfile.rmtree(temp_dir) 

78