Hugging Face 发布 TRELLIS.2 3D 生成模型 LoRA 微调指南

TL;DR：在本篇博客中，我们提供了微调脚本和容器，用于在你的 3D 资产上微调流行的 3D 模型 TRELLIS.2，既可以在本地运行，也可以在 Hugging Face Jobs 上运行。如果你对微调感兴趣，可以使用我们的脚本，只需几个参数（假设你已将数据集上传至 Hugging Face Hub），也可以通过这个仓库自行折腾。如果你想了解更多，我们讲解了脚本的工作原理。开始吧！

简介

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

TRELLIS 引入了一种通用的内部表示，称为统一潜在表示。它的解码器将这种表示转换为所需输出格式的 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 是一个受限（gated）模型
如果你使用自己的数据集，请在开始前将其网格上传到 Hub。

TRELLIS.2 的工作原理

在训练之前，我们来看看 TRELLIS.2 如何生成 3D 物体，以及我们要在哪些地方添加 LoRA 适配器。

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

对输入图像编码

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

稀疏结构生成（Sparse Structure Flow）

稀疏结构流 DiT 预测物体的粗略形状。它是一个扩散 transformer（DiT），生成一个大小为 16³ × 8 的小型潜变量。变分自编码器（VAE）解码器随后将这个潜变量转换为 64 × 64 × 64 的占用网格。网格中的每个体素表示该位置是否属于物体的一部分。体素就是三维空间中的像素。我们为这个模型添加 LoRA 适配器，因为它的输出会引导几何阶段。

O-Voxel 与结构化潜变量

O-Voxel 是 TRELLIS.2 的 3D 格式。它是一个稀疏体素网格，只保留物体表面上的体素。每个体素保留局部几何信息（例如网格顶点）和材质属性。由于这种格式很紧凑，CPU 只需几秒钟即可将网格转换为 O-Voxel，GPU 也只需几秒钟就能将 O-Voxel 转换回网格。

结构化稀疏潜变量来自 TRELLIS。它们将 3D 数据保存在稀疏张量中，从而加快计算速度。模型只使用网格中约 10% 的体素，因此不会在空白体素上浪费算力。于是，形状流 DiT 和纹理流 DiT 是稀疏 transformer。而稀疏结构流 DiT 是稠密 transformer。

几何生成（Shape Flow）

形状流 DiT 使用占用网格和 DINOv3 特征 token 来生成潜变量。另一个 VAE 解码器随后将这个潜变量转换为分辨率 512 × 512 × 512 的 O-Voxel 几何。

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

材质生成 + 网格创建

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

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

规划

微调哪些组件

我们只为稀疏结构流 DiT 和形状流 DiT 添加 LoRA 适配器。这些 DiT 控制着物体的形状和几何。我们冻结纹理流 DiT，以保留其现有能力。

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

如果我们在这些贴图上训练纹理流 DiT，它就能学会把光照表现为一种表面属性。它也可能忘掉之前学到的东西，这个问题被称为灾难性遗忘。

为什么要用 LoRA

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

你将运行的内容

一个 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，我们传入以下参数：

标志	含义
`-d`	分离运行：打印作业 ID 和 URL 后立即返回。要跟踪作业，请运行 `hf jobs logs -f <id>`
`--flavor a100-large`	硬件。我们使用 a100-large（80 GB）。要查看所有选项，请运行 hf jobs hardware
`--timeout 6h`	作业在此时间后停止。两条命令都可恢复运行。预处理会跳过已完成的物体。训练从 --push-to 中最新的检查点继续
`--secrets HF_TOKEN`	将你的令牌（来自 HF_TOKEN 或 hf auth login）作为加密的环境变量传给作业。令牌不会写入镜像中
`cr.hf.co/...:latest`	要运行的镜像。它之后的所有文本都是在容器中运行的命令

该容器还接受一些额外参数（无论在 Hugging Face Jobs 上运行还是用 Docker 运行），你可以在这里的文档中找到。预处理流水线中的步骤将在后文的这一节中说明。

LoRA 微调

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

我们使用 ‘ss’ 和 ‘shape’ 来表示我们训练 LoRA 的两种模式。‘ss’ 模式只微调 Sparse-Structure Flow DiT，而 ‘shape’ 模式则微调 Shape Flow 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 选择要训练的 flow DiT（ss 或 shape），STEPS 设置优化器步数。批次大小为 4 并配合 4 步梯度累积，可得到有效批次大小 16。然后，该片段加载所选 flow 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）自动混合精度执行前向传播。自动混合精度让 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 中的 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（稀疏结构流 DiT 使用 logitNormal，形状流 DiT 根据 TRELLIS.2 配置使用均匀分布）。
生成与数据形状匹配的随机噪声，并重塑 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 个小批次的梯度。梯度检查点通过在反向传播期间重新计算某些中间值来节省 GPU 内存。外层循环运行给定的优化器步骤数，而内层循环在一组批次上累积梯度。

对于每个批次，DINOv3 对条件图像进行编码。对于无分类器引导（CFG）丢弃，循环会将随机比例（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 配置文件包含完整的训练设置：ss_flow_lora_gso.json 用于稀疏结构流 DiT，slat_flow_img2shape_lora_gso.json 用于形状流 DiT。每次运行都在 Hugging Face Jobs 上使用一块 A100 80 GB。稀疏结构运行耗时 7.2 小时，形状运行耗时 3.6 小时。

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

流匹配损失随采样的物体、噪声和时间步而变化。要确定微调是否改善了几何质量，我们还必须在留出的物体上比较输出结果。

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

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

指标	含义
Chamfer 距离	从一个表面上的点到另一个表面上最近点的距离，双向测量。越低越好。
F-score @ 0.02	精确率（生成表面有多少接近参考）和召回率（生成表面覆盖了多少参考表面）的组合。这里匹配阈值设为扫描包围盒对角线的 2%。越高越好。
法向一致性	表面方向的吻合程度。表面法线是垂直于表面的向量。越高越好。

在图表中的九个物体对比中，基础模型的平均 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，也是可能的下一步。它们的表示和训练管线可能需要不同的方法。

感谢一路跟随，我们下期再见！👋🤗