TensorFlow Probability计算Kalman滤波对数似然为何如此缓慢？

问题描述 #

我尝试用TensorFlow测试卡尔曼滤波器，按照官方文档定义LinearGaussianStateSpaceModel模型，生成样本并计算对数似然值。运行的代码如下：

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

ndims = 2
step_std = 1.0
noise_std = 5.0
model = tfd.LinearGaussianStateSpaceModel(
  num_timesteps=1000,
  transition_matrix=tf.linalg.LinearOperatorIdentity(ndims),
  transition_noise=tfd.MultivariateNormalDiag(
   scale_diag=step_std**2 * tf.ones([ndims])),
  observation_matrix=tf.linalg.LinearOperatorIdentity(ndims),
  observation_noise=tfd.MultivariateNormalDiag(
   scale_diag=noise_std**2 * tf.ones([ndims])),
  initial_state_prior=tfd.MultivariateNormalDiag(
   scale_diag=tf.ones([ndims])))

x = model.sample(1) # 从观测序列的先验分布采样
lp = model.log_prob(x) # 计算观测（批次）的边际似然
print(lp)

在Colab的GPU环境下，计算log_prob耗时30秒，请问该问题的原因是什么，以及如何优化性能？

耗时原因分析 #

• 默认实现的调度开销：LinearGaussianStateSpaceModel的log_prob依赖卡尔曼滤波做精确推断，默认实现未针对GPU做充分并行优化，序列式的滤波步骤无法充分利用GPU算力，产生不必要的调度延迟。
• 线性算子抽象的额外成本：使用tf.linalg.LinearOperatorIdentity作为转移/观测矩阵，虽然逻辑简洁，但线性算子的抽象层会引入额外的张量操作调度开销，在GPU上这种抽象无法直接触发底层高效的核函数计算。
• 未启用XLA编译：TensorFlow默认未开启XLA（加速线性代数），而XLA能通过算子融合、内存访问优化等方式，大幅提升序列式计算（如卡尔曼滤波）在GPU上的执行效率。

性能优化方案 #

1. 启用XLA编译 #

通过全局配置或tf.function的JIT编译加速计算：

# 全局启用XLA
tf.config.optimizer.set_jit(True)

# 或用tf.function单独包裹log_prob计算（更灵活）
@tf.function(jit_compile=True)
def compute_log_prob(model, x):
    return model.log_prob(x)

lp = compute_log_prob(model, x)

2. 替换线性算子为普通张量 #

将LinearOperatorIdentity替换为原生单位矩阵张量，消除抽象层开销：

transition_matrix = tf.eye(ndims)
observation_matrix = tf.eye(ndims)

model = tfd.LinearGaussianStateSpaceModel(
  num_timesteps=1000,
  transition_matrix=transition_matrix,
  transition_noise=tfd.MultivariateNormalDiag(
   scale_diag=step_std**2 * tf.ones([ndims])),
  observation_matrix=observation_matrix,
  observation_noise=tfd.MultivariateNormalDiag(
   scale_diag=noise_std**2 * tf.ones([ndims])),
  initial_state_prior=tfd.MultivariateNormalDiag(
   scale_diag=tf.ones([ndims])))

3. 批量处理样本 #

利用GPU并行计算优势，批量处理多个样本，降低单样本的调度开销：

# 生成32个样本的批次
x = model.sample(32)
# 一次计算整个批次的log_prob，平均耗时会显著降低
lp = compute_log_prob(model, x)

4. 升级依赖库版本 #

确保使用最新版的TensorFlow和TensorFlow Probability，新版本通常会修复GPU算子的性能瓶颈：

!pip install --upgrade tensorflow tensorflow-probability

内容的提问来源于stack exchange，提问作者Fly_back