用 LoRA 微调 TRELLIS.2：让 3D 生成模型更懂你的数据

0 阅读

为什么通用模型搞不定特定物体

3D 生成技术在机器人、仿真和娱乐领域的应用越来越广，但有一个老问题始终存在：通用模型很难遵守特定领域的严格规则。比如你想生成一把椅子或一双鞋，通用模型给出的结果乍看没问题，细看却全是破绽——鞋底扭曲、椅腿不对称。这是因为现有的 3D 表示格式五花八门，从 GLB、PLY 文件到神经辐射场（NeRF）和 3D 高斯泼溅（3DGS），每种格式都有自己的脾气。

TRELLIS 系列模型试图解决这个问题。它引入了一种统一的潜在表示（unified latent representation），解码器可以将这种内部表示转换成任何你需要的 3D 输出格式。这意味着你不需要为每种格式单独训练一个模型。随后发布的 TRELLIS.2 进一步优化了 3D 结构的表示能力，但它依然是一个“通才”。如果你需要它在某个垂直领域（比如家具或鞋类）表现得更专业，直接用它往往不够精准。

我们做这个实验的目的很明确：用 Google Scanned Objects (GSO) 数据集中的 188 个鞋类物体，微调 TRELLIS.2，看看能不能让它生成的鞋底更靠谱。结果显示，经过 LoRA 微调后，模型在未见过的测试集上，F-score 从 0.276 提升到了 0.803，Chamfer 距离也从 0.174 降到了 0.033。这说明只要找对方法，小数据集也能让大模型学会“专精”。

TRELLIS.2 的工作流水线

在动手改代码之前，得先搞清楚 TRELLIS.2 是怎么把一个 2D 图片变成 3D 模型的。它的流水线分为五个主要阶段，我们要做的就是在合适的地方插入 LoRA 适配器。

1. 图像编码

首先，BiRefNet 负责把输入图片的背景去掉。接着，DINOv3-L 编码器将处理后的图片转化为特征 tokens。这些 tokens 就像是指南针，引导后续的 3D 生成过程。

2. 稀疏结构生成（Sparse Structure Flow）

这是一个扩散Transformer（DiT），它预测物体的粗略形状。它输出一个大小为 16³ × 8 的小潜在向量，然后通过变分自编码器（VAE）解码器变成一个 64 × 64 × 64 的占据网格（occupancy grid）。网格里的每个体素（voxel）代表该位置是否属于物体。因为这一阶段决定了物体的整体轮廓，我们在这里加入了 LoRA 适配器。

3. O-Voxel 与结构化潜在变量

O-Voxel 是 TRELLIS.2 的核心 3D 格式。它是一种稀疏体素网格，只保留物体表面的体素。每个体素存储局部几何信息（如网格顶点）和材质属性。这种格式非常高效，CPU 几秒就能把网格转成 O-Voxel，GPU 也能快速反向转换。由于模型只计算约 10% 的非空体素，计算效率大幅提升。

4. 几何生成（Shape Flow）

这一阶段的 Shape Flow DiT 利用占据网格和 DINOv3 特征 tokens，生成更精细的潜在变量。另一个 VAE 解码器将其转化为分辨率高达 512 × 512 × 512 的 O-Voxel 几何数据。这是决定物体细节形状的关键步骤，因此我们也在这里加入了 LoRA 适配器。

5. 材质生成与网格创建

最后，Texture Flow DiT 根据形状潜在变量和图像特征，生成表面的颜色、金属度、粗糙度和透明度。注意，我们没有在这一阶段加入 LoRA。原因是我们的数据集只有漫反射颜色图（diffuse color maps），缺乏训练所需的独立材质贴图。如果强行训练，模型可能会把光照效果误认为是表面属性，甚至出现“灾难性遗忘”，忘掉原本优秀的纹理生成能力。既然基座模型在颜色和纹理上已经表现不错，我们就冻结它，专注解决形状问题。

为什么要用 LoRA？

我们的实验数据集很小，只有 188 双鞋。对于这种规模的数据，全量微调（Full Fine-tuning）不仅成本高，还容易破坏模型原有的通用知识。LoRA（低秩适应）通过训练小型适配器矩阵，保持原始权重冻结，既减少了可训练参数量，又降低了显存需求。

在这次实验中，我们总共训练了约 930 万个适配器参数。具体配置如下：

  • Sparse-Structure Flow DiT：在注意力机制和前馈网络（MLP）层使用 rank 8 的适配器。
  • Shape Flow DiT：仅在注意力层使用 rank 16 的适配器。

Rank 值决定了适配器矩阵的大小，Rank 越高，适配器能学习的信息越多，但计算开销也越大。对于形状控制这种复杂任务，稍高的 Rank 是必要的。

实战指南：从预处理到训练

为了方便复现，我们提供了两个 Docker 镜像，分别用于预处理和微调。你可以在 Hugging Face Jobs 上运行，也可以在本地用 Docker 跑。

第一步：数据预处理

预处理的目标是把 GLB 或 OBJ 网格转换成 TRELLIS.2 训练所需的潜在表示，并渲染出条件图像。

预处理容器会调用 Blender（通过 xvfb-run无头运行）为每个物体渲染 16 个视角的图片。这些图片将作为训练时的条件输入。同时，容器会运行 TRELLIS.2 的预处理流水线，提取潜在变量。

在 Hugging Face Jobs 上运行的命令示例：

export HF_TOKEN=hf_xxx

hf jobs run -d --flavor a100-large --timeout 6h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-prep:latest \
  trellis2-prep \
    --meshes your-username/my-meshes \
    --push-to your-username/my-meshes-trellis2 \
    --private --holdout 0.05 \
    --views 16 --workers 8

这里有几个关键参数：

  • --flavor a100-large：使用 80GB 显存的 A100 显卡。
  • --holdout 0.05：保留 5% 的数据作为测试集，不参与训练。
  • --views 16：每个物体渲染 16 个视角。

对于 188 个物体，整个预处理过程在 A100 上大约耗时 1 小时。预处理好的数据会被推送到 Hugging Face Hub，方便后续训练直接读取。

第二步：LoRA 微调

训练容器加载预处理好的数据集，并根据配置文件启动训练。我们支持两种模式：ss（仅微调 Sparse-Structure Flow）和 shape（仅微调 Shape Flow）。

训练命令示例：

# 微调 Sparse-Structure Flow DiT
hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train --config ss \
    --latents your-username/my-meshes-trellis2 \
    --push-to your-username/my-trellis2-lora \
    --run-name ss_lora \
    --push-path ss_lora \
    --trackio-space your-username/trellis2-trackio

# 微调 Shape Flow DiT
hf jobs run -d --flavor a100-large --timeout 12h --secrets HF_TOKEN \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train --config shape \
    --latents your-username/my-meshes-trellis2 \
    --push-to your-username/my-trellis2-lora \
    --run-name shape_lora \
    --push-path shape_lora \
    --trackio-space your-username/trellis2-trackio \
    --set max_steps=3000 i_save=500

如果你想在本地运行，只需将 hf jobs run 替换为标准的 docker run 命令，并挂载相应的卷即可。

训练 pipeline 内部机制

注入适配器

TRELLIS.2 有自定义的前向传播逻辑，标准的 PeftModel() 包装器无法直接使用。我们调用 inject_adapter_in_model() 函数，直接将适配器插入模型层级中。为了数值稳定性， trainable 的适配器权重保持为 fp32，而前向传播使用 bf16 自动混合精度，以节省显存。

损失函数与训练循环

我们的目标函数是整流流匹配（Rectified Flow Matching），这与 Stable Diffusion 3 和 Flux 等图像生成模型使用的原理一致。TRELLIS.2 在预训练时也使用了同样的目标，因此非常适合微调。

训练循环的核心步骤如下：

  1. 采样时间步 t：根据配置，对 Sparse-Structure Flow 使用 logitNormal 分布，对 Shape Flow 使用均匀分布。
  2. 添加噪声：根据时间步 t，将干净数据与随机噪声混合。
  3. 预测速度：模型预测在 t × 1000 时刻的速度场。
  4. 计算损失：计算预测速度与目标速度（即数据的导数）之间的均方误差（MSE）。

为了防止梯度爆炸，我们在每次优化器步骤前对梯度进行裁剪。同时，使用梯度累积（gradient accumulation），每 4 个小批次更新一次权重，等效批大小为 16。

结果评估：鞋底终于像样了

我们分别在 Sparse-Structure Flow 和 Shape Flow 上训练了两个适配器。在 A100 上，Sparse-Structure 的训练耗时 7.2 小时，Shape Flow 耗时 3.6 小时。

虽然训练损失曲线显示 Loss 在下降，但这并不能直接反映几何质量。真正的考验在于那 9 个被保留下来的测试物体（held-out objects）。

对比结果显示，基座模型生成的鞋子鞋底往往扁平或扭曲，而经过 LoRA 微调后的模型，生成的鞋底结构与真实扫描数据高度吻合。定量指标也证实了这一点：

  • F-score @ 0.02：从 0.276 提升至 0.803。这个指标综合了精确率和召回率，越高越好。
  • Chamfer Distance：从 0.174 降至 0.033。这个指标衡量生成表面与参考表面之间的距离，越低越好。
  • Normal Consistency：从 0.530 提升至 0.804。这反映了表面法线方向的一致性，越高说明表面越平滑自然。

在所有 9 个测试物体上，微调后的模型 F-score 均有显著提升。这证明了我们策略的有效性：冻结纹理阶段，专注微调几何结构阶段，既能避免数据不足带来的副作用，又能精准提升特定品类的生成质量。

接下来可以做什么？

这次实验验证了 LoRA 在 3D 生成模型微调中的可行性。核心经验在于：选择与问题和数据相匹配的训练目标。对于鞋类这种几何结构复杂但纹理相对标准的物体，微调结构阶段收益最大。

未来，这套方法可以扩展到其他 3D 生成模型，如 Pixal3D 和 UniLat3D。此外，随着高斯泼溅和世界生成模型（如 HunyuanWorld 2.0 和 NVIDIA Lyra 2.0）的发展，如何针对这些新架构设计高效的微调方案，也是值得探索的方向。它们的表示形式和训练流水线可能不同，但“小数据、精微调”的思路依然适用。

如果你手头有特定的 3D 资产数据集，不妨试试这套流程。代码和预训练适配器已开源，欢迎在社区中交流反馈。