gpflow 模型的奇怪 `pickle`/`gpflow.utilities.freeze` 行为

Question

我一直在尝试（粗略地）在玩具数据集上训练和保存 gpflow SVGP 模型，主要遵循 this notebook example

使用pickle 保存模型后（我很欣赏不建议这样做，但我不认为这是这里的主要问题），我发现了一些不寻常的，我认为是意外行为： 如果我们在尝试picklemodel 之前没有调用gpflow.utilities.freeze(model)，那么我们会得到一个错误。 如果我们确实调用gpflow.utilities.freeze(model)（丢弃返回的冻结模型），那么model 可以被腌制而不会出错。

重现

最小的、可重复的示例

import numpy as np
import gpflow
import tensorflow as tf
import pickle
rng = np.random.RandomState(123)

N = 10000  # Number of training observations
X = rng.rand(N, 1)
Y = rng.randn(N, 1)
data = (X, Y)

n_inducing_vars = 100
Z = X[:n_inducing_vars]
minibatch_size = 100
n_iterations = 100

#Define model object
model = gpflow.models.SVGP(gpflow.kernels.Matern12(), gpflow.likelihoods.Bernoulli(), inducing_variable=Z, num_data=N)

#Create minibatch object
data_minibatch = (
tf.data.Dataset.from_tensor_slices(data).prefetch(
    N).repeat().shuffle(N).batch(minibatch_size)
    )
data_minibatch_it = iter(data_minibatch)
model_objective = model.training_loss_closure(data_minibatch_it)

#Define optimiser
optimizer = tf.keras.optimizers.Adam(0.001)
#Optimise both variational parameters and kernel hyperparameters.
for step in range(n_iterations):
    optimizer.minimize(model_objective,
                       var_list=model.trainable_variables
                       )

freeze = False
if not freeze:
    # pickle doesn't work
    pickle.dump(model, open('test1', 'wb'))
else:
    # if following code is executed, pickle works fine
    _ = gpflow.utilities.freeze(model)  # ignore return value
    pickle.dump(model, open('test1', 'wb'))

堆栈跟踪或错误消息

TypeError                                 Traceback (most recent call last)
<ipython-input-6-3d5f537ca994> in <module>
----> 1 pickle.dump(model, open('test1', 'wb'))

TypeError: can't pickle HashableWeakRef objects

预期行为

并不是说我希望泡菜在第一个实例中起作用，因为我知道这不是保存tensorflow 相关对象的推荐方法。但是，我当然不希望它在第一次失败但在第二次成功。从代码库来看，我不认为 gpflow.utilities.freeze(model) 应该改变 model，它似乎正在这样做。

系统信息

• 使用 GPflow 版本 2.0.0 ... 2.0.4 测试
• TensorFlow 版本：2.1.0，tensorflow_probability 0.9.0
• Python 版本：Python 3.6.9

我猜想在model 上调用freeze 时，实际上莫名其妙地将model 转换为“冻结”模型，然后该模型具有使其能够被腌制的“恒定”属性（https://gpflow.readthedocs.io/en/master/notebooks/intro_to_gpflow2.html#TensorFlow-saved_model） .

非常感谢您对此问题的任何澄清。

请注意，我在gpflow github (https://github.com/GPflow/GPflow/issues/1493) 上以issue 的形式发布了此问题，但决定应在此处向更广泛的 gpflow 社区广播此问题。

Answer 1

此行为适用于使用 tensorflow_probability 的双射器的任何代码/模型，并且不限于 SVGP 模型。在 GPflow 中，双射器用于约束参数，例如以确保核方差和长度尺度始终为正。

基本解释是 tensorflow_probability 的双射器保留了他们操作过的张量的缓存，例如，这允许他们准确地恢复以下示例中的原始张量：

import tensorflow as tf
import tensorflow_probability as tfp
bij = tfp.bijectors.Exp()
x = tf.constant(1.2345)
y = bij.forward(x)
assert bij.inverse(y) is x  # actual object identity, not just numerical equivalence

然而，这些缓存使用 HashableWeakRef 对象，这些对象无法腌制 - 甚至无法复制（使用 Python 标准库的 copy.deepcopy 函数）。

只有在您实际通过双射器运行张量时才会填充缓存 - 如果您只是创建模型而不优化它，您可以腌制（或复制）它就好了。当然，这通常不是很有用。

要解决此问题并允许复制“使用过的”（例如经过训练的）模型，我们有 gpflow.utilities.reset_cache_bijectors()。这由gpflow.utilities.deepcopy() 调用以允许复制。而gpflow.utilities.freeze() 反过来需要进行深度复制，以便为您提供冻结副本，而不是原地冻结模型，这解释了轻微的副作用。

所以不是freeze 是您成功腌制它所必需的；在酸洗之前添加对reset_cache_bijectors(model)的调用就足够了，将示例中的代码替换为

if not freeze:
    gpflow.utilities.reset_cache_bijectors(model)  # with this added call, pickle *does* work
    pickle.dump(model, open('test1', 'wb'))

最终，这是一个只能在上游通过 tensorflow_probability 在他们自己的代码中“正确”修复的问题。可以在awav 的pull request to tensorflow_probability 中找到更多详细信息，旨在解决此问题。

附带说明，正如markvdw 指出的那样，您可能会发现存储使用gpflow.utilities.read_values() 获得的模型的所有参数值更容易（它返回值的参数键字典），您可以存储以您喜欢的任何方式，并通过首先重新创建对象然后使用gpflow.utilities.multiple_assign() 分配参数来重新加载。

Answer 2

让我们看看下面几行发生了什么：

if not freeze:
    # pickle doesn't work
    pickle.dump(model, open('test1', 'wb'))  # Line 1
else:
    # if following code is executed, pickle works fine
    gpflow.utilities.freeze(model)           # Line 2
    pickle.dump(model, open('test1', 'wb'))  # Line 3

在Line 1 中，经过训练的模型包含Parameters 实例，这些实例将 TensorFlow 概率双射器作为从受约束空间到不受约束空间并返回的变换器。 TFP 双射器缓存所有正向和反向计算。双射器的缓存是用映射实现的，其中键是正向和反向函数的张量输入，值是返回的张量对象。不幸的是，张量（例如 np.arrays）不能被散列，为此 TFP 为张量实现了一个包装器HashableWeakRef。错误消息 "TypeError: can't pickle HashableWeakRef objects" 具有误导性。这实际上意味着HashableWeakRef python 无法复制该实例，仅仅是因为它是对尚未创建的对象的引用。因此，这些对象不能被腌制。

在Line 2 中，freeze 方法由两个调用组成：第一个调用删除双射器的内容，即缓存，第二个调用是copy.deepcopy。 freeze 背后的魔力在于它删除了引用。是的，它修改了现有对象，但它既不影响急切计算也不影响tf.functioned 函数。清理使deepcopy成为可能。

Line 3 有效，因为该对象没有要复制的引用。

这个问题有很长的报告记录，并尝试在 GPflow 中修复它： TFP#547、TFP#944、GPflow#1479、GPflow#1293、GPflow#1338

这是 TensorFlow 概率的建议修复：TFP#947