WOMD Motion 官方教程：Colab 报错排查与解决方案

这份笔记总结了在 Google Colab 运行 Waymo Open Dataset（WOMD）Motion 官方教程时遇到的问题、原因和最终可复用的解决方案。

1. 最终可工作的环境

本次验证通过的组合如下：

项目	版本/设置
Colab 运行时版本	`2025.07`
Python	`3.11.13`
`waymo-open-dataset-tf-2-12-0`	`1.6.7`
TensorFlow	`2.13.0`
JAX	`0.4.13`
jaxlib	`0.4.13`
教程测试文件的地图采样数	`num_map_samples = 20000`

注意：虽然包名中含有 tf-2-12-0，但本次安装 1.6.7 后实际使用的是 TensorFlow 2.13.0。判断兼容性时应以安装后的依赖和版本输出为准，不要只看包名。

2. 问题一：无法安装 `jaxlib==0.4.13`

报错现象

在较新的 Colab 默认运行时中执行：

%pip install waymo-open-dataset-tf-2-12-0==1.6.7

出现类似错误：

Could not find a version that satisfies the requirement jaxlib==0.4.13
No matching distribution found for jaxlib==0.4.13

当时运行时使用的是 Python 3.13。

原因

WOMD 1.6.7 固定依赖较旧的 jax==0.4.13 和 jaxlib==0.4.13。
jaxlib 0.4.13 没有适用于 Python 3.13 的 wheel。
这个旧版 CPU wheel 也不能再单纯依赖普通 PyPI 索引自动找到。

解决方案

先在 Colab 中选择兼容的旧运行时：

代码执行程序 → 更改运行时类型 → 运行时版本 → 2025.07

确认它使用 Python 3.11，然后执行：

%pip install --no-cache-dir \
  "https://storage.googleapis.com/jax-releases/nocuda/jaxlib-0.4.13-cp311-cp311-manylinux2014_x86_64.whl" \
  "waymo-open-dataset-tf-2-12-0==1.6.7"

安装完成后必须重启会话：

代码执行程序 → 重新启动会话

3. 问题二：重启后出现 JAX CUDA 插件错误

报错现象

ERROR: jax._src.xla_bridge: Jax plugin configuration error
...
from jax._src.lib import triton
ImportError: cannot import name 'triton'

同时版本检查可能仍然显示：

Python: 3.11.13
TensorFlow: 2.13.0
JAX: 0.4.13
jaxlib: 0.4.13
Waymo import: OK

原因

Colab 运行时中残留了较新的 CUDA 12 JAX 插件，而教程安装的是旧版 JAX 0.4.13。插件会尝试从旧版 JAX 中导入不存在或不兼容的 triton 接口，因此打印红色错误。

这个错误不一定会阻止 WOMD 的 TensorFlow 教程继续运行，但应清理掉不兼容插件，避免后续出现隐蔽问题。

解决方案

%pip uninstall -y jax-cuda12-plugin jax-cuda12-pjrt

然后再次重启会话。Motion 教程主要使用 TensorFlow，保留 JAX 的 CPU 后端即可。

4. 问题三：TFRecord 文件不存在

报错现象

NotFoundError:
/content/waymo-od-motion/tutorial/.../tfexample.tfrecord;
No such file or directory

原因

官方教程中的 FILENAME 路径含有字面量 ...，它只是占位符，不是真实文件路径。教程仓库不会自动附带完整 WOMD 数据集，需要用户自行下载数据并填写实际路径。

另外，TFRecordDataset() 使用延迟读取：创建 dataset 时可能不会报错，直到执行 next() 真正读取文件时才会出现 NotFoundError。

快速验证教程的解决方案

下载官方仓库中的单样本 TFExample 测试文件：

!wget -q -O /content/motion_data_one_example.tfrecord \
  https://raw.githubusercontent.com/waymo-research/waymo-open-dataset/master/src/waymo_open_dataset/utils/testdata/motion_data_one_example.tfrecord

FILENAME = "/content/motion_data_one_example.tfrecord"

检查文件：

import os

print(os.path.exists(FILENAME))
print(os.path.getsize(FILENAME))

本次样本的正常结果为：

True
1182920

对于完整数据集，这个解析单元格需要指向 training_tfexample.tfrecord... 一类的 TFExample 数据分片，不能指向 Scenario Proto 分片。

5. 问题四：`roadgraph_samples/xyz` 无法解析

报错现象

InvalidArgumentError:
Key: roadgraph_samples/xyz. Can't parse serialized Example.

原因

当前教程代码可能设置：

num_map_samples = 30000

但官方测试文件 motion_data_one_example.tfrecord 中只有 20,000 个 roadgraph samples：

roadgraph_samples/xyz = 60000 个 float = 20000 × 3

解析器却按照 [30000, 3] 读取，因此形状不匹配。

解决方案

把定义改成：

num_map_samples = 20000

然后必须重新运行整个特征定义单元格，包括：

roadgraph_features = {...}
state_features = {...}
traffic_light_features = {...}

features_description = {}
features_description.update(roadgraph_features)
features_description.update(state_features)
features_description.update(traffic_light_features)

最后再运行创建和解析数据集的单元格。

只修改 num_map_samples 变量而不重建 roadgraph_features 是无效的，因为旧字典中的 FixedLenFeature 已经保存了原来的形状。

验证：

print(parsed["roadgraph_samples/xyz"].shape)

预期输出：

(20000, 3)

6. 当前代码解析的是 TFExample，还是 Scenario Proto？

当前这套代码使用：

tf.io.FixedLenFeature(...)
tf.io.parse_single_example(data, features_description)

因此它解析的是 TFExample 表示形式。

两种 WOMD 数据表示的区别如下：

对比项	TFExample 表示	Scenario Proto 表示
外层文件	TFRecord	TFRecord
每条记录的 payload	序列化的 `tf.train.Example`	序列化的 `Scenario` protobuf
解析方法	`tf.io.parse_single_example`	`scenario_pb2.Scenario.FromString`
解析结果	Tensor 字典	层次化 `Scenario` 对象
是否需要手写字段形状	是，需要 `features_description`	否，由 proto schema 定义

需要特别注意：

两者的文件扩展名都可能是 .tfrecord。
不能仅凭扩展名判断数据格式，必须知道每条 record 内部存放的 payload。
tf.train.Example 本身也是一种 protobuf；Waymo 文档中的 “TFExample” 和 “Scenario Proto” 指的是两套不同的数据表示。

最小 Scenario Proto 示例

下载官方单 Scenario 测试文件：

!wget -q -O /content/motion_data_one_scenario.tfrecord \
  https://raw.githubusercontent.com/waymo-research/waymo-open-dataset/master/src/waymo_open_dataset/utils/testdata/motion_data_one_scenario.tfrecord

解析：

import tensorflow as tf
from waymo_open_dataset.protos import scenario_pb2

dataset = tf.data.TFRecordDataset(
    "/content/motion_data_one_scenario.tfrecord",
    compression_type="",
)
scenario_bytes = next(dataset.as_numpy_iterator())
scenario = scenario_pb2.Scenario.FromString(scenario_bytes)

print(scenario.scenario_id)
print(len(scenario.tracks))
print(len(scenario.map_features))

官方相关教程：

官方小样本的对应关系：

文件	内容格式	解析方式
`motion_data_one_example.tfrecord`	TFExample	`tf.io.parse_single_example`
`motion_data_one_scenario.tfrecord`	Scenario Proto	`Scenario.FromString`

7. 完整的恢复运行顺序

以后重新打开 Colab 时，可以按下面的顺序恢复：

将运行时版本切换为 2025.07（Python 3.11）。
运行安装单元格，卸载不兼容的 CUDA JAX 插件并安装 WOMD 1.6.7。
重启 Colab 会话。
运行 import 和版本检查单元格。
下载测试数据，或挂载 Google Drive 并填写完整数据集的实际路径。
使用测试 TFExample 时设置 num_map_samples = 20000。
从特征定义单元格开始依次重新运行，确保 features_description 已按新形状重建。
再运行数据解析、绘图、动画和 metrics 单元格。

版本检查代码

import sys
import tensorflow as tf
import jax
import jaxlib
import waymo_open_dataset

print("Python:", sys.version.split()[0])
print("TensorFlow:", tf.__version__)
print("JAX:", jax.__version__)
print("jaxlib:", jaxlib.__version__)
print("Waymo import: OK")

8. 如何保存代码块、输出和下次继续查看

保存 Notebook 和输出

打开 修改 → 笔记本设置。
确保“保存此笔记本时忽略代码单元格输出”没有被勾选。
选择 文件 → 在云端硬盘中保存副本。
建议命名为 WOMD_motion_tutorial_working.ipynb。
按 Ctrl+S，等待页面显示已保存。
还可以通过 文件 → 下载 → 下载 .ipynb 保存本地备份。

这样通常会保留：

Markdown 和代码单元格；
已运行单元格的文本、图片和动画输出；
单元格运行顺序。

不会随 Notebook 保存的内容

Colab 的运行时虚拟机是临时的。会话释放后，以下内容通常会消失：

通过 pip 临时安装的软件包；
Python 内存中的变量和模型；
/content 下临时下载的数据；
当前运行进程和 GPU/CPU 状态。

因此，应把下面两个单元格长期保留在 Notebook 顶部：

环境安装与版本检查单元格；
数据下载或 Google Drive 挂载单元格。

下次打开时重新执行它们，即可恢复工作环境。

9. 最终结论

这次连续报错并不是单一问题，而是四层兼容性问题叠加：

新版 Python 与旧版 jaxlib 不兼容；
Colab 自带的新 CUDA JAX 插件与旧 JAX 不兼容；
教程中的数据路径只是占位符；
当前教程 schema 的 30,000 个地图采样点与官方小测试文件的 20,000 个采样点不一致。

完成运行时降级、指定旧版 wheel、清理 CUDA 插件、提供真实 TFRecord 文件，并把 num_map_samples 调整为 20000 后，官方测试 TFExample 可以被正常解析。

posted @ 2026-09-28 20:41 MineralWaterAndCola 阅读(7) 评论(0) 收藏举报

刷新页面返回顶部

MineralWaterAndola

WOMD Motion 官方教程：Colab 报错排查与解决方案

1. 最终可工作的环境

2. 问题一：无法安装 `jaxlib==0.4.13`

报错现象

原因

解决方案

3. 问题二：重启后出现 JAX CUDA 插件错误

报错现象

原因

解决方案

推荐的一次性安装单元格

4. 问题三：TFRecord 文件不存在

报错现象

原因

快速验证教程的解决方案

5. 问题四：`roadgraph_samples/xyz` 无法解析

报错现象

原因

解决方案

6. 当前代码解析的是 TFExample，还是 Scenario Proto？

最小 Scenario Proto 示例

7. 完整的恢复运行顺序

版本检查代码

8. 如何保存代码块、输出和下次继续查看

保存 Notebook 和输出

不会随 Notebook 保存的内容

9. 最终结论

公告

MineralWaterAndola

WOMD Motion 官方教程：Colab 报错排查与解决方案

1. 最终可工作的环境

2. 问题一：无法安装 jaxlib==0.4.13

报错现象

原因

解决方案

3. 问题二：重启后出现 JAX CUDA 插件错误

报错现象

原因

解决方案

推荐的一次性安装单元格

4. 问题三：TFRecord 文件不存在

报错现象

原因

快速验证教程的解决方案

5. 问题四：roadgraph_samples/xyz 无法解析

报错现象

原因

解决方案

6. 当前代码解析的是 TFExample，还是 Scenario Proto？

最小 Scenario Proto 示例

7. 完整的恢复运行顺序

版本检查代码

8. 如何保存代码块、输出和下次继续查看

保存 Notebook 和输出

不会随 Notebook 保存的内容

9. 最终结论

公告

2. 问题一：无法安装 `jaxlib==0.4.13`

5. 问题四：`roadgraph_samples/xyz` 无法解析