解决Keras模型序列化错误

奇怪的模型序列化错误,原因居然是...

解决Keras模型序列化错误

封面图片:JOSHUA COLEMAN


在经历过TensorFlow 2中子类化模型的一堆坑之后,我尽可能的使用Keras Fucntional API来构建模型。然而这几天一个Keras模型序列化的错误困扰了我好久,最终的解决方案却异常简单。

错误现象

在开始训练模型时会看到一行警告:

WARNING:tensorflow:Model failed to serialize as JSON. Ignoring... ('Not JSON Serializable:', <tf.Tensor: shape=(), dtype=int32, numpy=58)
警告信息

这个警告并不影响训练,也不影响checkpoint文件的保存。但是在导出模型为`saved_model`格式时,该警告会变成错误,导致保存失败。

理论上,只要定位到无法序列化的对象,然后分析原因,想办法序列化就好。但是从错误信息只能看出这是一个INT32类型的Tensor,数值为58。没有它在代码中的变量名称。

这就很麻烦了,代码中的变量千千万,哪个才是罪魁祸首?

分析异常

饭要一口一口吃,问题要分割后一点一点解决。首先从错误信息来看,这是一个整形变量。代码中涉及到整形变量的地方一般是Tensor的形状、通道数、某些自定义的计时器、循环步数等。这样可以排除一部分。其次,该变量的形状为(),这意味着它不太可能是Tensor形状(一般为[bs, height, width, channels])。最关键的是它的数值为58。在神经网络的构建过程中,涉及到形状的数值都尽可能的被8整除。这个奇怪的58——正好是116/2,而116,是我当前所构建网络的第一层卷积层通道数。所以,它应当是通道数。

异常原因

但是,一个整形为何无法序列化?

Google之后,我在官方GitHub上发现了一篇相关讨论。作者遇到了和我类似情况。不同之处在于他在loss函数中创建了一个类型为float32的tf.variable。

TypeError: ‘Not JSON Serializable’ while doing tf.keras.Model.save and using keras variable in loss_weights in tf.keras.Model.compile · Issue #28799 · tensorflow/tensorflow
System information Have I written custom code: NA OS Platform and Distribution: Ubuntu 16.04 LTS Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: NA Ten...

这时,我恍然大悟。

问题并非是“整型”无法序列化,而是Tensor无法序列化!而我的这行代码因为使用了`tf.math`函数误将将整型通道数变为了整型Tensor。

num_identity = tf.cast(tf.math.ceil(num_input_channels * split), tf.int32)
引起错误的代码

解决方案

原因找到,解决方案自然也有了。此时的通道数本质是一个Eager Tensor,所以,在其后使用`.numpy()`方法获得其数值即可。

num_identity = tf.cast(tf.math.ceil(num_input_channels * split), tf.int32).numpy()
修复后的代码

现在,模型可以正常保存了。


以上内容涉及的模型为ShuffleNet v2,如果你感兴趣的话欢迎查阅源码:

yinguobing/models
A playground for friendly deep neural network models. - yinguobing/models