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

文章标题：3D 生成模型的终极微调指南

正文：

TLDR; 在这篇博客中，我们提供了微调脚本和容器，用于在你的 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 在特定物体类别上的性能。我们将说明在训练期间哪些 pipeline 组件需要微调，哪些保持不变。该 pipeline 将结构、详细几何形状和材质分开。因此，我们只微调对任务重要的组件。

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

我们将使用什么？

LoRA

LoRA，即低秩适应，通过训练小型适配器矩阵来微调模型，同时保持原始权重冻结。这减少了可训练参数的数量，以及它们的梯度和优化器状态所需的内存。我们使用 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 有五个主要阶段：去除图像背景、编码图像、预测粗结构、生成详细几何形状、生成材质。然后 pipeline 将结果导出为带纹理的网格。

输入图像编码

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

稀疏结构生成（稀疏结构流）

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

O-Voxel 与结构化 Latent

O-Voxel 是 TRELLIS.2 的 3D 格式。它是一个稀疏体素栅格，仅保留物体表面上的体素。每个体素保留局部几何信息（如网格顶点）和材质属性。由于该格式体积小，CPU 可在数秒内将网格转换为 O-Voxel，GPU 也可在数秒内将 O-Voxel 转换回网格。

结构化稀疏 latent 来自 TRELLIS。它们将 3D 数据存储在稀疏张量中，从而加快计算速度。模型仅使用栅格中约 10% 的体素，因此不会在空体素上消耗算力。因此，形状流 DiT 和纹理流 DiT 是稀疏 Transformer，而稀疏结构流 DiT 是稠密 Transformer。

几何生成（形状流）

形状流 DiT 利用占据栅格和 DINOv3 特征 token 生成 latent。随后，另一个 VAE 解码器将该 latent 转换为分辨率为 512 × 512 × 512 的 O-Voxel 几何。

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

材质生成 + 网格创建

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

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

规划

哪些组件需要微调

我们仅向稀疏结构流 DiT 和形状流 DiT 添加 LoRA 适配器。这两个 DiT 控制物体的形状和几何。我们冻结纹理流 DiT，以保持其现有能力。

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

如果我们在这些贴图上训练纹理流 DiT，它可能会学会将光照表现为表面属性，也可能遗忘之前学到的内容，即灾难性遗忘问题。

为什么使用 LoRA

我们的子集包含 GSO 中 1,030 个对象里的 188 个鞋子对象。对于这么小的数据集，LoRA 比全量微调更合适。全量微调的成本高得多。它还可能覆盖模型已经擅长的内容（其学习到的先验）。我们的配置为两个目标 DiT 总共训练约 9.3 million 个适配器参数。我们在 sparse-structure flow DiT 的注意力层和前馈（MLP）层中使用 rank 8 适配器。我们在 shape flow 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，我们传递以下参数：

这个容器还接受几个更多参数（无论在 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
      

        # 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
      

用于稀疏结构 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

用于 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

## 训练流水线内部

### 预处理为 latents

我们使用 TRELLIS.2 的冻结编码器把输入网格转换为训练用的 latents。我们保存这些 latents，这样后续运行就不必再次编码网格。

![Preprocessing](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/3d-lora-finetuning/preprocessing.png)

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

![Conditioning](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/3d-lora-finetuning/conditioning_views.png)

> 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()。这个函数把适配器直接添加到模型上，并保留其自定义接口。

![Preparing LoRA Adapters](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/3d-lora-finetuning/lora_adapters.png)

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

为了数值稳定性，我们将可训练的适配器权重保持为 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 在预训练中使用了相同的目标函数，因此它是微调的一个好选择。关于理论，请阅读这篇关于 flow matching 的博客。

![训练循环](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/3d-lora-finetuning/training_loop.png)

我们只更新稀疏结构和形状 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 配置，稀疏结构 flow DiT 使用 logitNormal，形状 flow 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

每个批次中，DINOv3 对条件图像进行编码。对于无分类器引导（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 的损失则大致保持在同一水平。这些曲线显示了训练目标如何变化，但它们并不直接衡量生成几何的质量。

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

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

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

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

感谢你一路跟随，我们下一篇见！👋🤗