Tensorflow概率：保存和加载模型答案

【问题标题】：Tensorflow Probability: saving and loading modelTensorflow概率：保存和加载模型
【发布时间】：2021-08-19 04:57:09
【问题描述】：

我正在尝试用 TensorFlow 概率拟合模型，例如：

input = Input(shape=(32,32,32,3))
x = tfp.layers.Convolution3DReparameterization(
        64, kernel_size=5, padding='SAME', activation=tf.nn.relu,
        data_format = 'channels_first')(input)
x = tf.keras.layers.MaxPooling3D(pool_size=(2, 2, 2),
                                 strides=(2, 2, 2),
                                 padding='SAME')(x)
x = tf.keras.layers.Flatten()(x)
output = tfp.layers.DenseFlipout(10)(x)

model3 = Model(input, output)
model3.save('tf_test_model3.h5')

当我将模型加载为model3 = load_model('tf_test_model3.h5') 时，我收到以下错误：

ValueError: Unknown layer: Conv3DReparameterization. Please ensure this object is passed to the `custom_objects` argument.

当我将它传递给 custom_objects 时：

custom_objects= {'Conv3DReparameterization': tfp.layers.Convolution3DReparameterization}
model3 = load_model('tf_test_model3.h5', custom_objects=custom_objects)

我收到以下错误： TypeError: 'str' object is not callable

我做错了什么？我该如何解决这个问题？

【问题讨论】：

TensorFlow 和 TFP 的版本是什么？
@Frightera Tensorflow 版本为 2.6.0，TFP 版本为 0.13.0
我只会保存权重并在重新创建模型后重新加载它们。其他错误，我有空再看看。
我试过了，但是在获取输出层的形状时又出现了另一个错误。 stackoverflow.com/questions/68843958/how-to-get-shape-of-layers

标签： python tensorflow keras tensorflow-probability

【解决方案1】：

您的代码中的错误是data_format = channels_first

您给出的输入形状为 (32,32,32,3)，它是 channel_last 类型的数据。提供编辑后的代码

import keras as tk
import tensorflow as tf
import tensorflow_probability as tfp
input = tf.keras.Input(shape=(32,32,32,3))
x = tfp.layers.Convolution3DReparameterization(
        64, kernel_size=5, padding='SAME', activation=tf.nn.relu,
        data_format = 'channels_last')(input) #changed  from channels_first
x = tf.keras.layers.MaxPooling3D(pool_size=(2, 2, 2),
                                 strides=(2, 2, 2),
                                 padding='SAME')(x)
x = tf.keras.layers.Flatten()(x)
output = tfp.layers.DenseFlipout(10)(x)

model3 = tk.Model(input, output)
model3.save('tf_test_model3.h5')

【讨论】：