一条指令跑通 JAX 转 PyTorch：openpi 权重迁移实战-北京尧图网络

一条指令跑通 JAX 转 PyTorchopenpi 权重迁移实战【免费下载链接】openpi项目地址: https://gitcode.com/GitHub_Trending/op/openpi你做 JAX 转 PyTorch 迁移第一步load_state_dict直接抛size mismatch注意力权重是三维的 einsum 形状怎么对都对不齐。下面基于 openpi 的examples/convert_jax_model_to_pytorch.py脚本把 pi0/pi05 的权重迁移完整跑一遍最终产出一个可部署的model.safetensors遇到不匹配也能直接定位到具体参数。JAX 转 PyTorch 迁移时权重经历了哪些变换完整转换链如下脚本读取 Orbax 检查点按模块重排权重最后写出 safetensors 转换产物。JAX 与 PyTorch 的权重对不上根因是三类结构差异JAX 卷积权重是[H, W, C_in, C_out]顺序PyTorch 是[C_out, C_in, H, W]JAX 版 Gemma 的 Q/O 投影存成四维 einsum 权重PyTorch 的 Linear 只收二维矩阵JAX 参数键是扁平字符串树PyTorch 是点号模块路径。convert_pi0_checkpoint先经restore_params把 Orbax 检查点还原成 float32 的 numpy 字典还原过程会尊重 JAX 模型加载时的 dtype 转换数值不漂移再交给两个 slice 函数逐模块重排。# slice_paligemma_state_dict重排 ViT 卷积核 state_dict[pytorch_key] state_dict.pop(jax_key).transpose(3, 2, 0, 1)这一行把 JAX 的[H, W, C_in, C_out]转成 PyTorch 的[C_out, C_in, H, W]。# 每层 q 投影einsum 四维权重 - Linear 二维矩阵 q_proj_weight_reshaped ( llm_attention_q_einsum[i].transpose(0, 2, 1) .reshape(num_attention_heads * head_dim, hidden_size) )这一步把 einsum 的四维注意力权重压成nn.Linear可用的二维(out_features, in_features)矩阵。# slice_gemma_state_dictpi05 是自适应归一化pi0 是 RMSNorm if pi05 in checkpoint_dir: state_dict[f{layer}.dense.weight] kernel.transpose() else: state_dict[f{layer}.weight] scalepi05 把每层的 RMSNorm scale 参数换成了自适应归一化的 Dense 层脚本靠检查点路径里是否带 pi05 区分这两种结构写入不同键名。ViT/LLM、动作专家、投影层三组权重最终合并成一个字典装入新实例化的PI0Pytorch按目标精度转换后经safetensors.torch.save_model写出model.safetensors同时落一份config.json记录 action_dim、精度等字段assets/目录也会一并拷贝过去。⚡ 转换脚本实操路径从检查点到 safetensors开始前确认你在仓库根目录后面所有命令都依赖uv run定位虚拟环境。1. 同步依赖并给 transformers 打补丁PyTorch 实现依赖 transformers 的三个补丁AdaRMS、激活精度控制、KV 缓存复用克隆后uv sync再把补丁文件拷进 venvgit clone https://gitcode.com/GitHub_Trending/op/openpi cd openpi uv sync cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/运行后你会看到Resolved 128 packages in 3s Installed 128 packages in 12s如果这里卡住cp -r报 No such file 时先跑uv pip show transformers确认 4.53.2 已装好再重试。产物带 .venv、且 transformers 已打补丁的仓库目录后面所有命令都在此目录下执行。2. 拉取 JAX 检查点并检查参数键结构先用--inspect_only确认检查点键结构与脚本预期一致避免转换白跑uv run examples/convert_jax_model_to_pytorch.py --checkpoint_dir ~/.cache/openpi/openpi-assets/checkpoints/pi05_droid --config_name pi05_droid --inspect_only输出是层级化的参数键列表img/embedding/kernel img/pos_embedding llm/embedder/input_embedding llm/layers/attn/q_einsum/w如果这里卡住检查点目录不存在时先运行uv run scripts/serve_policy.py droid让它从 gs://openpi-assets 下载完CtrlC 后回来执行本命令。产物确认参数树包含 img/视觉塔、llm/语言模型与动作专家、投影层三组键下一步复用同一路径。3. 执行转换并输出 safetensors 权重这条命令执行 JAX 转 PyTorch 的完整转换以 float32 还原保证数值输出精度默认 bfloat16与 JAX 推理精度对齐uv run examples/convert_jax_model_to_pytorch.py --checkpoint_dir ~/.cache/openpi/openpi-assets/checkpoints/pi05_droid --config_name pi05_droid --output_path ./pi05_droid_pytorch运行后你会看到Converting PI0 checkpoint from ... to ./pi05_droid_pytorch Model conversion completed successfully! Model saved to ./pi05_droid_pytorch如果这里卡住出现size mismatch时核对--config_name与检查点是否匹配pi0 和 pi05 的投影层结构不同交叉使用必报此错。产物./pi05_droid_pytorch/ 下三个文件——model.safetensors权重、config.json记录 action_dim、精度等、assets/归一化统计下一步把策略服务指向该目录。4. 让策略服务指向转换后的模型create_trained_policy会自动识别 PyTorch 格式API 与 JAX 版一致只换检查点路径uv run scripts/serve_policy.py policy:checkpoint --policy.configpi05_droid --policy.dir./pi05_droid_pytorch运行后你会看到INFO:root:Creating server (host: my-host, ip: 127.0.0.1)该行打印后服务在 8000 端口监听等待观测请求。如果这里卡住服务启动报 transformers 相关属性缺失回到第 1 节重跑cp -r补丁命令。产物可调用的策略服务客户端怎么接入看远程推理文档。报错关键字排障速查与延伸入口报错关键字根因修复动作size mismatch for ...--config_name与检查点不匹配用匹配的 config_name 重跑第 3 节命令Invalid precision--precision不支持 float16改用 bfloat16 或 float32ModuleNotFoundError: openpi未在仓库内以 uv 方式运行在仓库根目录用uv run ...执行命令transformers 缺 AdaRMS 相关属性transformers_replace 补丁未生效重跑第 1 节的cp -r命令gs://openpi-assets下载失败GCS 网络不可达手动下载检查点到本地目录--checkpoint_dir传本地路径转换报错不在上表、或要继续 openpi 模型迁移链路时去这三个地方CONTRIBUTING 贡献指南按模板提交转换报错 issue规范统计文档迁移后 state/action 归一化对不上时检查 norm_stats 的处理用下面 3 行代码自测转换后的检查点能否被策略正确加载from openpi.training import config as _config from openpi.policies import policy_config policy policy_config.create_trained_policy(_config.get_config(pi05_droid), ./pi05_droid_pytorch)【免费下载链接】openpi项目地址: https://gitcode.com/GitHub_Trending/op/openpi创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

一条指令跑通 JAX 转 PyTorch：openpi 权重迁移实战

相关推荐

猫抓浏览器扩展：3 步完成网页视频下载

openai-agents-python 沙箱实战：外发传真资料包策略（Outbound Fax Packet Requirements）的落地实现

Reflex 布局组件完全指南：Box、Stack、Flex、Grid 与 Container 的选型与实战

更多资讯

游戏引擎工具链设计：从模块化架构到AI辅助开发

PaddleOCR TIPC Linux 端补充训练功能测试完全指南：覆盖蒸馏、PACT 量化、FPGM 裁剪与自定义 OP 的训练链路验证

Label Studio Enterprise 2.36.0 版本全解析：Data Manager 批量列编辑、Custom Interfaces 沙箱与成员分析增强

变速箱箱体加工工艺与夹具设计全流程解析

PHP API接口开发实战：从规范到高性能优化

Odoo 19库存期初导入实战指南与优化技巧

今日更新

SpacetimeDB C++ 模块库实战：用 C++20 在数据库内构建 WebAssembly 模块

别踩雷！不是所有 AI 写作工具都靠谱，2026 高校认可工具精选

SSM+MySQL幼儿园管理系统源码解析：从部署到优化

本周精选

超人会飞不算本事：系统稳定依赖清晰规则与边界设计

超人VS蜘蛛侠：拆解超级IP的影响力与传播方法论

基于CNN的调制信号识别：MATLAB实现时频图分类实战

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？