CARLA–PyTorch JIT 轮子动力学推理模块

该文档详细说明了代码中各模块的功能与调用流程，方便开发者快速理解与使用。

目录

1. 项目概述

本模块在 CARLA 仿真环境 中实现了 轮子动力学计算 与 PyTorch JIT 模型推理 之间的无缝桥接。通过这一模块，开发者可以在真实感强的仿真环境中使用深度学习模型对轮子的物理行为进行推理，并对其进行动态调整。这一桥梁的搭建为复杂的自动驾驶仿真提供了更高效的计算与推理能力。

该模块的各个函数关系图

该模块功能如下：

  1. 将 C++ 原始数据封装成 PyTorch 张量：将 CARLA 仿真中的轮子动力学数据转化为 PyTorch 张量，方便模型处理。

  2. 调用 TorchScript 模型进行推理：加载并使用已训练的 TorchScript 模型执行前向推理，计算轮子物理行为。

  3. 从输出张量中提取结果：从推理结果中获取粒子力和轮子力，并填充到 WheelOutput 中。
  4. 支持多种推理模式：提供 CPU、动态推理及 CUDA 加速模式，适应不同硬件环境和计算需求。

2. 依赖与头文件

  • 源码参考 pytorch.h 源码
  • torch/torch.h / torch/script.h：PyTorch C++ API 与 JIT 接口
  • torchscatter & torchcluster：验证扩展库的 CUDA 支持
  • tensorexpr_fuser：控制 JIT Fuse 优化
  • CUDACachingAllocator：用于显式清理 CUDA 缓存

3. 全局辅助函数

3.1 add_mark

  • 作用：空壳函数，用于在代码中插入标记，便于日志或调试

  • 参数：text — 标记文本，不作实际处理

  • 源码链接：更多细节请参考 pytorch.cpp

4. carla::learning 命名空间

该命名空间封装了与轮子动力学及模型推理相关的核心函数。

4.1 test_learning

  • 源码参考: pytorch.cpp 源码

  • 功能：打印 torchcluster 和 torchscatter 的 CUDA 版本，验证环境

  • 示例输出：

cuda version X.Y

4.2 GetWheelTensorInputs

  • 源码参考: pytorch.cpp 源码
  • 输入：WheelInput 结构体，包含：
  • particles_positions (float*)
  • particles_velocities (float*)
  • wheel_positions (float[3])
  • wheel_orientation (float[4])
  • wheel_linear_velocity (float[3])
  • wheel_angular_velocity (float[3])
  • num_particles (int)
  • 处理：
  • 使用 torch::from_blob 将原始指针包装成 at::Tensor（不复制数据）
  • 构造 std::vectortorch::jit::IValue，按顺序存放所有张量
  • 输出：Value 元组，用于 Module::forward()

4.3 GetWheelTensorOutput

  • 源码参考:pytorch.cpp 源码
  • 输入：
  • particle_forces：形状 [num_particles, 3] 张量
  • wheel_forces：长度 6 张量（力 xyz + 扭矩 xyz）
  • 处理：
  • 从 wheel_forces.data_ptr<float>() 读取 6 个值，赋给 WheelOutput 对应字段
  • 遍历 particle_forces.data_ptr<float>()，按 xyz 顺序填充 WheelOutput::_particle_forces 向量
  • 输出：填充后的 WheelOutput 结构体

4.4 GetWheelTensorOutputDynamic

  • 源码参考: pytorch.cpp 源码
  • 功能：与 GetWheelTensorOutput 相同，用于“动态”推理后的后处理

5. NeuralModelImpl 结构体

  • 源码参考: pytorch.cpp 源码
  • 成员变量：
  • module：已加载的 TorchScript 模型
  • particles_*_tensors：可选的缓存张量队列（当前未深入使用）
  • 成员函数：
  • GetWheelTensorInputsCUDA：CUDA 版本的 GetWheelTensorInputs

5.1 GetWheelTensorInputsCUDA

  • 源码参考: pytorch.cpp 源码
  • 功能：与 GetWheelTensorInputs 相同，但创建后立即 .cuda() 并附加 wheel_idx 信息
  • 用途：并行推理时区分不同车轮

6. NeuralModel 类接口

该类对外暴露模型加载、输入设置、推理和结果获取接口。

源码参考: pytorch.cpp 源码

6.1 构造与析构

NeuralModel::NeuralModel()
NeuralModel::~NeuralModel()
  • 功能：使用 std::make_unique<NeuralModelImpl>() 初始化内部实现；自动释放

6.2 LoadModel

void NeuralModel::LoadModel(char* filename, int device)
  • 功能：
  • torch::jit::setTensorExprFuserEnabled(false) 关闭 TensorExpr Fuse
  • Model->module = torch::jit::load(filename) 加载模型
  • （可选）将模型转移到 cuda:device
  • 异常处理：捕获并打印 c10::Error

6.3 SetInputs

void NeuralModel::SetInputs(Inputs input)
  • 功能：保存传入的 Inputs（包含四个 WheelInput、操作指令、地形类型、verbose 标志）

6.4 推理：Forward

  • 源码参考: pytorch.cpp 源码
  • 流程：
  • 调用 GetWheelTensorInputs 构建四组车轮张量
  • 驾驶命令（steering, throttle, braking）及 optional terrain_type、verbose
  • 执行 module.forward(TorchInputs)
  • 拆解输出元组为 8 个张量，调用 GetWheelTensorOutput 填充 _output

6.5 动态推理：ForwardDynamic

  • 区别：使用 GetWheelTensorOutputDynamic，并在末尾调用 CUDACachingAllocator::emptyCache() 清理 CUDA 缓存

6.6 CUDA 推理：ForwardCUDATensors

  • 区别：
  • 使用 GetWheelTensorInputsCUDA 将所有张量 .cuda()
  • 驾驶指令等也 .cuda()
  • 其余逻辑同 Forward

6.7 结果获取：GetOutputs

  • 源码参考: pytorch.cpp 源码
  • 功能：返回内部 _output，包含四个 WheelOutput

7. 整体流程

  1. 初始化与加载
NeuralModel model;
model.LoadModel("mymodel.pt", /*device=*/0);
  1. 设置输入
Inputs inp = /* 构造 WheelInput + 操作指令 */;
model.SetInputs(inp);
  1. 选择推理模式
model.Forward();                // CPU 模式
// 或
model.ForwardDynamic();         // 动态 + 缓存清理
// 或
model.ForwardCUDATensors();     // 全 CUDA 模式
  1. 获取输出
auto &outs = model.GetOutputs();
// outs.wheel0 ... outs.wheel3