3D 生成模型微调终极指南

TLDR：在本文中，我们提供微调脚本和容器，用于在你自己的 3D 资产上微调流行 3D 模型 TRELLIS.2，既可以在本地运行，也可以在 Hugging Face Jobs 上运行。如果你对微调感兴趣，只需少量参数即可使用我们的脚本（假设你已将数据集上传到 Hugging Face Hub），也可以用这个仓库自行调整。如果你想了解更多，我们会逐步讲解脚本的工作原理。开始吧！

引言#

3D 生成在机器人、仿真、媒体和娱乐领域有许多用途。在这些领域中，很难让生成的对象遵守严格规则，也难以标准化模型和数据表示。文件格式有很多，例如 GLB 和 PLY。表示方式也有很多，例如神经辐射场（NeRF）和 3D 高斯泼溅（3D Gaussian Splatting，3DGS）。每种方式都有不同的优缺点。

TRELLIS 引入了一种通用的内部表示，称为统一潜在表示（unified latent representation）。它的解码器将这种表示转换为所需输出格式的 3D 模型。因此，你不必为每种表示格式使用不同的模型。

同一个团队随后发布了 TRELLIS.2，它改进了模型对 3D 结构的表示。但这两个模型都没有提供简单的方法来将其适配到特定任务。

我们想改进什么#

像 TRELLIS.2 这样的通用模型，可能会在特定应用至关重要的细节上出错。例如，用户可能需要家具模型具有准确的几何结构。如果没有针对该任务进行微调，预训练模型可能生成不够准确的几何结构。我们在实验中使用 GSO 的一个子集，其中包含 188 个鞋类对象。在最初的测试中，预训练模型通常能生成看起来正确的鞋子，但鞋底不正确。我们测试 LoRA 能否让鞋底更准确。

我们的目标是使用 Google Scanned Objects（GSO）数据集，提高 TRELLIS.2 在特定对象类别上的性能。我们将解释在训练期间应当微调哪些流水线组件，哪些应保持不变。该流水线将结构、详细几何和材质分开。因此，我们只微调对任务重要的组件。

选择冻结什么，与选择微调什么同样重要。

我们将使用什么？#

LoRA#

LoRA 是低秩适应（low-rank adaptation）的缩写，它通过训练小型适配器矩阵、同时保持原始权重冻结来微调模型。这减少了可训练参数的数量，以及其梯度和优化器状态所需的内存。我们使用 Hugging Face 的 PEFT 库来添加和训练这些适配器。更多信息请阅读我们的 LoRA 指南。

Hugging Face Jobs#

Hugging Face Jobs 在云硬件上运行容器化工作负载。我们将其用于预处理和训练。如果你想使用自己的 GPU，请参见下面的本地 Docker 命令。开始之前，请确保你具备：

• 一个 HF_TOKEN，可以从这里获取
• 你的 Hugging Face 账户必须有权访问 DINOv3，因为 DINOv3 是门控模型
• 如果你使用自己的数据集，请在开始前将其网格上传到 Hub。

TRELLIS.2 如何工作#

在训练之前，让我们看看 TRELLIS.2 如何生成 3D 对象，以及我们将在哪里添加 LoRA 适配器。

TRELLIS.2 有五个主要阶段：移除图像背景、编码图像、预测粗略结构、生成详细几何，以及生成材质。然后流水线将结果导出为带纹理的网格。

编码输入图像#

首先，BiRefNet 移除输入图像的背景。然后 DINOv3-L 将图像编码为特征词元。生成模型使用这些词元来引导 3D 对象。

稀疏结构生成（Sparse Structure Flow）#

稀疏结构流（sparse-structure flow）DiT 预测对象的粗略形状。它是一个扩散 Transformer（diffusion transformer，DiT），生成大小为 16³ × 8 的小潜在表示。然后变分自编码器（variational autoencoder，VAE）解码器将该潜在表示转换为 64 × 64 × 64 的占用网格（occupancy grid）。网格中的每个体素（voxel）表示该位置是否属于对象。体素是三维空间中的像素。我们向该模型添加 LoRA 适配器，因为它的输出引导几何阶段。

O-Voxel 与结构化潜在表示#

O-Voxel 是 TRELLIS.2 的 3D 格式。它是一个稀疏体素网格，只保留对象表面上的体素。每个体素保留局部几何（例如网格顶点）和材质属性。由于格式紧凑，CPU 可以在几秒内将网格转换为 O-Voxel。GPU 可以在几秒内将 O-Voxel 转换回网格。

结构化稀疏潜在表示（structured sparse latents）来自 TRELLIS。它们将 3D 数据保存在稀疏张量中，从而加快计算。模型只使用网格中约 10% 的体素，因此不会在空体素上花费计算。因此，形状流 DiT 和纹理流 DiT 是稀疏 Transformer。稀疏结构流 DiT 则是稠密 Transformer。

几何生成（Shape Flow）#

形状流（shape flow）DiT 使用占用网格和 DINOv3 特征词元生成潜在表示。然后另一个 VAE 解码器将该潜在表示转换为分辨率为 512 × 512 × 512 的 O-Voxel 几何。

我们在这一阶段向 DiT 添加 LoRA 适配器，因为它控制对象的形状。

材质生成 + 网格创建#

在最后一个阶段，纹理流 DiT 使用形状潜在表示和 DINOv3 特征作为输入。它定义几何表面的材质和纹理。然后 VAE 解码器将潜在表示转换为每个体素的基础色、金属度、粗糙度和 alpha 值。我们不向这一阶段添加 LoRA 适配器。我们的数据集有漫反射颜色贴图，但缺少这一阶段训练所需的独立材质贴图。此外，我们的问题在于鞋子的形状。基础模型已经能生成良好的颜色和纹理。

最后，流水线将几何和材质属性组合成带纹理的网格，并导出为 GLB 文件。

规划#

微调哪些组件#

我们只向稀疏结构流 DiT 和形状流 DiT 添加 LoRA 适配器。这些 DiT 控制对象的形状和几何。我们冻结纹理流 DiT，以保持其当前能力。

我们的数据集包含带有烘焙光照效果的漫反射颜色贴图。要训练纹理流 DiT，我们需要独立的材质贴图。我们数据集中的颜色贴图将表面颜色与光照混在一起。因此，它们不是这项训练的正确目标。

如果我们在这些贴图上训练纹理流 DiT，它可能会学会将光照表现为表面属性。它也可能忘记之前学到的东西，这个问题称为灾难性遗忘（catastrophic forgetting）。

为什么使用 LoRA#

我们的子集包含 GSO 中 1,030 个对象里的 188 只鞋。对于这么小的数据集，LoRA 比全量微调更合适。全量微调成本高得多，还可能覆盖模型已经擅长的能力（其学到的先验）。我们的配置为两个目标 DiT 总共训练约 930 万个适配器参数。我们在稀疏结构流 DiT 的注意力和前馈（MLP）层中使用秩 8 适配器。我们在形状流 DiT 的注意力层中使用秩 16 适配器。秩决定适配器矩阵的大小。秩越高，适配器能学到的内容越多。

你将运行什么#

3D 模型环境可能需要许多依赖和安装步骤。我们提供两个 Docker 镜像（可以在 Jobs 或你自己的机器上运行），其中包含必要的软件。一个镜像用于预处理，一个用于微调。

要训练 LoRA 适配器并将它们上传到 Hub：

• 让容器访问你的数据。
• 为你的任务设置参数。
• 运行预处理容器。
• 运行微调容器。

然后我们会解释微调脚本、其设置，以及每个设置如何影响训练。

实验概览#

我们将逐步介绍实验的三个部分：

• 准备数据。将 GLB 或 OBJ 网格转换为 TRELLIS.2 用于训练的潜在表示，并渲染输入图像。
• 训练适配器。针对每个 DiT 运行一个任务，使用准备好的数据集和 LoRA 配置。
• 解释代码。检查容器如何加载数据、添加适配器并训练模型。

我们会保存预处理后的数据集，以便所有训练运行都可以复用它。

操作演示#

在 Hugging Face Jobs 上运行#

预处理

预处理容器使用 Blender 为每个对象渲染 16 个视图。LoRA 训练将使用这些视图作为条件图像。Blender 通过 xvfb-run 工具在无显示环境下运行。容器还会通过预处理流水线处理每个对象。参见这里的鞋类数据集，了解一个中间输出示例。对于我们的 188 个对象，完整运行在 a100-large 任务上花费了约 1 小时。这个配套仓库中也提供了预处理流水线的精简版本。

export HF_TOKEN=hf_xxx

# Replace your-username with your Hugging Face username

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

对于 Hugging Face Jobs，我们传递以下参数：

无论运行在 Hugging Face Jobs 还是 Docker 上，此容器都接受更多参数，可在此处找到文档。预处理流水线所采取的步骤将在后面的章节中说明。

LoRA 微调

训练容器加载准备好的数据集以及配置预设或配置文件。它使用 Trackio 记录训练指标。使用 Trackio 仪表盘监控运行。

我们用 ss 和 shape 表示训练 LoRA 的两种模式。ss 模式只微调稀疏结构流 DiT，而 shape 模式微调形状流 DiT。在底层，它们只是引用每个阶段要使用的不同 config.json 文件。有关过程的更多细节，请参见后面的章节。

export HF_TOKEN=hf_xxx # Replace your-username with your Hugging Face username

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

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

如果设置了 --push-to，容器会将适配器检查点上传到 Hub。为简洁起见，容器接受的参数也在此处记录。同一仓库中还提供了一个精简脚本。

如果你想试用并基于我们已处理的数据集进行微调，可以随意使用以下命令。

export HF_TOKEN=hf_xxx

# Replace your-username with your Hugging Face username

# For the 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 suvadityamuk/gso-footwear-trellis2-latents \
  --push-to your-username/trellis2-footwear-lora \
  --run-name ss_lora \
  --push-path ss_lora \
  --trackio-space your-username/trellis2-trackio \
  --set max_steps=3000 i_save=500

# For the 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 suvadityamuk/gso-footwear-trellis2-latents \
  --push-to your-username/trellis2-footwear-lora \
  --run-name shape_lora \
  --push-path shape_lora \
  --trackio-space your-username/trellis2-trackio \
  --set max_steps=3000 i_save=500

在本地使用 Docker 运行#

预处理

命令和标志与 Hugging Face Jobs 相同。只有 Docker 选项不同。

# For preprocessing into latents
docker run --gpus all --rm --shm-size 16g -e HF_TOKEN=hf_xxx \
  -v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
  cr.hf.co/suvadityamuk/trellis2-finetuning-prep:latest \
  trellis2-prep --meshes your-username/my-meshes --push-to your-username/my-meshes-trellis2

LoRA 微调

# For Sparse-structure Flow DiT
docker run --gpus all --rm --shm-size 16g \
  -v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
  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 --trackio-space your-username/trellis2-trackio

# For Shape Flow DiT
docker run --gpus all --rm --shm-size 16g \
  -v hf-cache:/root/.cache/huggingface -v "$PWD/work":/workspace/work \
  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 --trackio-space your-username/trellis2-trackio

如果你想在自己选择的机器上本地使用我们已预处理的数据集，请使用以下命令：

export HF_TOKEN=hf_xxx # Replace your-username with your Hugging Face username

# For Sparse-structure Flow DiT
docker run --gpus all --rm --shm-size 16g -e HF_TOKEN \
  -v hf-cache:/root/.cache/huggingface \
  -v "$PWD/work":/workspace/work \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train \
  --config ss \
  --latents suvadityamuk/gso-footwear-trellis2-latents \
  --push-to your-username/trellis2-footwear-lora \
  --trackio-space your-username/trellis2-trackio \
  --set max_steps=3000 i_save=500

# For Shape Flow DiT
docker run --gpus all --rm --shm-size 16g -e HF_TOKEN \
  -v hf-cache:/root/.cache/huggingface \
  -v "$PWD/work":/workspace/work \
  cr.hf.co/suvadityamuk/trellis2-finetuning-train:latest \
  trellis2-train \
  --config shape \
  --latents suvadityamuk/gso-footwear-trellis2-latents \
  --push-to your-username/trellis2-footwear-lora \
  --trackio-space your-username/trellis2-trackio \
  --set max_steps=3000 i_save=500

我们训练好的 LoRA 适配器位于此仓库：suvadityamuk/trellis-2-finetune-gso

训练流水线内部#

预处理为潜在表示#

我们使用 TRELLIS.2 的冻结编码器将输入网格转换为训练潜在表示。我们保存这些潜在表示，以便后续运行不必再次编码网格。

我们还在 Blender 中为每个对象渲染 16 个视图。训练期间，我们为每个样本随机选择一个视图作为条件图像。我们保留 188 个对象中的 9 个用于评估，并在其余 179 个上训练。这些留出对象可以显示适配器是否在模型训练期间未见过的对象上有效。

Blender 为某个对象渲染的部分视图

第一个片段用于设置数据。FLOW 选择要训练的流 DiT（ss 或 shape），STEPS 设置优化器步数。批大小为 4、梯度累积步数为 4，得到有效批大小 16。然后，该片段加载所选流 DiT 的 LoRA 配置文件，以 TRELLIS.2 格式创建数据集，移除留出对象，并为微调创建数据加载器。

cfg = json.load(
    open(
        {
            "ss": "/workspace/TRELLIS.2/configs/lora/ss_flow_lora_gso.json",
            "shape": "/workspace/TRELLIS.2/configs/lora/slat_flow_img2shape_lora_gso.json",
        }[FLOW]
    )
)

key, latents = {
    "ss": ("ss_latent", "ss_latents/ss_enc_conv3d_16l8_fp16_64"),
    "shape": ("shape_latent", "shape_latents/shape_enc_next_dc_f16c32_fp16_512"),
}[FLOW]

roots = json.dumps(
    {
        "gso": {
            key: f"{DATA}/{latents}",
            "render_cond": f"{DATA}/renders_cond",
            "meta": DATA,
        }
    }
)

ds = getattr(datasets, cfg["dataset"])(roots, **cfg["dataset_args"])

holdout = set(open(f"{DATA}/holdout.txt").read().split())
ds.instances = [(root, sha) for root, sha in ds.instances if sha not in holdout]

batches = cycle(
    DataLoader(
        ds,
        batch_size=BATCH,
        shuffle=True,
        drop_last=True,
        num_workers=4,
        collate_fn=getattr(ds, "collate_fn", None),
    )
)

准备 LoRA 适配器#

我们使用 peft 库将 LoRA 适配器插入选定的层。TRELLIS.2 有自定义的前向传播，因此标准的 PeftModel() 包装器无法使用。相反，我们调用 inject_adapter_in_model()。该函数直接将适配器添加到模型中，并保留其自定义接口。

我们将可训练的适配器权重保持为 32 位浮点数（fp32），以保证数值稳定性，并在前向传播中使用 bfloat16（bf16）自动转换（autocasting）。自动转换让 PyTorch 对某些操作使用较低精度，从而减少内存使用。自动转换不会改变冻结主干权重的存储方式。

model = models.from_pretrained(cfg["ckpt"])
for block in model.blocks:
    block.use_checkpoint = True

inject_adapter_in_model(
    LoraConfig(
        r=cfg["rank"], lora_alpha=cfg["alpha"], target_modules=cfg["lora_targets"]
    ),
    model,
)

params = []
for name, p in model.named_parameters():
    p.requires_grad_("lora_" in name)
    if p.requires_grad:
        p.data = p.data.float()  # fp32 master weights for LoRA
        params.append(p)

model.cuda().train()

训练#

值得注意的是，我们这里的目标函数是经典的整流流匹配（rectified flow matching）（见于 Stable Diffusion 3、Flux 2 等流行的图像生成模型）。TRELLIS.2 在预训练中使用了相同的目标，因此它是微调的良好选择。关于理论，请阅读这篇关于流匹配的博客。

我们只更新稀疏结构 DiT 和形状 DiT 中的 LoRA 参数。所有原始模型权重保持冻结。

下面的片段展示了一个阶段的核心训练循环。它们比容器代码更简单，并且不包含某些设置步骤。完整工作流请使用容器命令。

我们初始化优化器并加载 DINOv3，以从输入图像中提取特征。

opt = torch.optim.AdamW(
    params, lr=cfg["lr"], weight_decay=cfg["weight_decay"], betas=(0.9, 0.95)
)

dino = DinoV3FeatureExtractor(
    cfg["image_cond_model"], image_size=cfg["image_size"]
).cuda()

然后我们在训练循环中遵循以下步骤：

• 为每个样本采样一个时间步 t（根据 TRELLIS.2 配置，稀疏结构流 DiT 使用 logitNormal，形状流 DiT 使用 uniform）。
• 生成与数据形状匹配的随机噪声，并重塑 t 以便广播。
• 根据 t 将干净数据与噪声混合（其中 t = 0 是噪声最小的数据，t = 1 是纯噪声）。
• 在 t × 1000 处预测速度，以匹配预训练条件。
• 将损失计算为预测速度与目标速度（x_t 的一阶导数）之间的均方误差。

def flow_loss(x0, cond):
    s = cfg["t_schedule"]

    t = (
        torch.rand(x0.shape[0])
        if s["name"] == "uniform"
        else torch.sigmoid(
            torch.randn(x0.shape[0]) * s["args"]["std"] + s["args"]["mean"]
        )
    )
    t = t.cuda()

    noise = (
        x0.replace(torch.randn_like(x0.feats))
        if cfg["sparse"]
        else torch.randn_like(x0)
    )

    tv = t.view(-1, *[1] * (len(x0.shape) - 1))

    x_t = (1 - tv) * x0 + (cfg["sigma_min"] + (1 - cfg["sigma_min"]) * tv) * noise

    pred, target = model(x_t, t * 1000, cond), (1 - cfg["sigma_min"]) * noise - x0

    return (
        F.mse_loss(pred.feats, target.feats)
        if cfg["sparse"]
        else F.mse_loss(pred, target)
    )

训练循环计算损失并更新适配器权重。梯度累积在每次优化器步骤之前累加 4 个小批次的梯度。梯度检查点（gradient checkpointing）通过在反向传播期间重新计算某些中间值来节省 GPU 内存。外层循环运行给定的优化器步数，而内层循环在一组批次上累积梯度。

对于每个批次，DINOv3 编码条件图像。对于无分类器引导（classifier-free guidance，CFG）的 dropout，循环会将随机比例（p_uncond）样本的条件设置为零。然后，它使用 bf16 自动转换计算流匹配损失。内层循环结束后，它会裁剪梯度以防止更新过大，并执行一次优化器步骤。它每 10 步记录一次损失，每 500 步保存一次适配器检查点。

for step in range(1, STEPS + 1):
    for _ in range(ACCUM):
        batch = next(batches)
        x0 = recursive_to_device(batch["x_0"], "cuda")

        with torch.no_grad():
            cond = dino(batch["cond"])
            cond[torch.rand(len(cond), device="cuda") < cfg["p_uncond"]] = 0

        with torch.autocast("cuda", dtype=torch.bfloat16):
            loss = flow_loss(x0, cond) / ACCUM

        loss.backward()

    torch.nn.utils.clip_grad_norm_(params, cfg["grad_clip"])

    opt.step()
    opt.zero_grad()

结果#

我们训练了两个适配器：一个用于稀疏结构 DiT，一个用于形状 DiT。容器镜像中的 LoRA 配置文件包含完整的训练设置：稀疏结构流 DiT 使用 ss_flow_lora_gso.json，形状流 DiT 使用 slat_flow_img2shape_lora_gso.json。每次运行在 Hugging Face Jobs 上使用一块 A100 80 GB。稀疏结构运行耗时 7.2 小时，形状运行耗时 3.6 小时。

稀疏结构流 DiT 的训练损失在开始时快速下降，然后下降得更慢。形状流 DiT 的损失保持在大致相同的水平。这些曲线展示了训练目标如何变化，但它们并不直接衡量生成几何的质量。

流匹配损失会随采样的对象、噪声和时间步变化。要了解微调是否改善了几何，我们还必须在留出对象上比较输出。

下图比较了 9 个留出对象中 5 个的输出。它展示了输入图像、基础模型生成的输出、经过旋转调整后的基础输出、使用鞋类适配器的输出，以及参考扫描。在这些示例中，微调模型的鞋底更接近扫描结果。

我们还使用三个指标在测试集上评估几何：

在图表的九对象比较中，基础模型的平均 F-score 为 0.276，而使用鞋类适配器的模型为 0.803。Chamfer 距离从 0.174 降至 0.033，法线一致性从 0.530 升至 0.804。使用两个适配器后，F-score 在所有九个留出对象上都更好。

所有运行的统计数据请参见这个 Trackio space。

未来方向#

呼，我们完成了！我们通过向控制几何的两个模型组件添加 LoRA 适配器，为鞋类生成微调了 TRELLIS.2。我们准备了一次数据，在独立任务中训练适配器，并在留出对象上评估输出。

主要经验是选择一个与问题和现有数据相匹配的训练目标。对于这个实验，我们微调了结构和形状阶段，冻结了纹理阶段。

接下来，我们可以研究同样的选择是否有助于其他 3D 生成模型，例如 Pixal3D 和 UniLat3D。

高斯泼溅和世界生成模型，例如 HunyuanWorld 2.0 和 NVIDIA Lyra 2.0，也是可能的下一步。它们的表示和训练流水线可能需要不同的方法。

感谢你读到这里，我们下一篇见！👋🤗