Cannot Save Tensorflow_probability model

Question

I'm trying to save the tensorflow model bellow:

Model: "sequential_117"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 fcl_0 (DenseFlipout)        (None, 47)                1175      
                                                                 
 BN_0 (BatchNormalization)   (None, 47)                188       
                                                                 
 fcl_1 (DenseFlipout)        (None, 14)                1330      
                                                                 
 final_layer (DenseFlipout)  (None, 1)                 29        
                                                                 
=================================================================
Total params: 2,722
Trainable params: 2,628
Non-trainable params: 94
_________________________________________________________________

But I'm facing this error : TypeError: cannot pickle 'module' object

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_14328/3501717249.py in <module>
      1 # save model
----> 2 best_model.save("model_BNN")
      

[...]

~\Anaconda3\lib\copy.py in deepcopy(x, memo, _nil)
    159                     reductor = getattr(x, "__reduce_ex__", None)
    160                     if reductor is not None:
--> 161                         rv = reductor(4)
    162                     else:
    163                         reductor = getattr(x, "__reduce__", None)

TypeError: cannot pickle 'module' object

Does anyone have an idea about this issue ? Is that a bug of tensorflow probability ?

Thanks in advance.

Answer 1

model.save in Tensorflow takes only the name of the save object, and not an extension. This is because it saves a bunch of things in a folder with that name.

best_model.save("model_BNN") should work for your case if you are trying to save the entire model.

Check the documentation here. If you want to save only the weights of your trained model, you can use this