DeepSDF 工程源码阅读：网络、损失与数据流

本文记录本地工程 ~/other/temp/DeepSDF 的关键实现，重点回答三个复现时最容易卡住的问题：

MLP 网络具体如何构建，跳跃连接和权重归一化放在哪里。
Clamped L1 Loss 在代码里如何通过 torch.clamp 截断远离表面的 SDF。
训练时一个 batch 如何同时拿到形状的 latent code 索引、三维坐标和真实 SDF。

源码阅读范围：

model/model_sdf.py
model/train_sdf.py
data/dataset_sdf.py
utils/utils_deepsdf.py
config_files/train_sdf.yaml

默认配置先看清楚

这个仓库的默认训练配置在 config_files/train_sdf.yaml：

batch_size: 16
clamp: true
clamp_value: 0.1

num_layers: 8
inner_dim: 256
latent_size: 64
skip_connections: True

因此，当前源码默认的单个查询点输入不是论文常见说明里的 128 + 3 = 131 维，而是：

\mathbf{x}_{\text{model}} = [\mathbf{z}_i,\ x,\ y,\ z] \in \mathbb{R}^{64+3} = \mathbb{R}^{67}

也就是说，每条 SDF 样本由两部分组成：

z_i：当前形状的 latent code，默认 64 维。
(x, y, z)：空间中的查询点坐标，3 维。

备注

源码注释里仍然写着 128 和 131，这是旧配置说明。按当前 train_sdf.yaml 运行时，实际输入维度是 67。

网络定义：MLP、Skip Connection 与 Weight Normalization

DeepSDF MLP architecture

网络定义在 model/model_sdf.py 的 SDFModel 中。初始化时先计算输入维度：

dim_coords = 3
input_dim = self.latent_size + dim_coords
self.skip_tensor_dim = copy.copy(input_dim)

在默认配置下：

latent_size = 64
dim_coords = 3
input_dim = 67
skip_tensor_dim = 67

MLP 主干怎么搭

主干层在一个循环里创建：

layers = []
for _ in range(num_layers - num_extra_layers):
    layers.append(
        nn.Sequential(
            nn.utils.weight_norm(nn.Linear(input_dim, inner_dim)),
            nn.ReLU()
        )
    )
    input_dim = inner_dim
self.net = nn.Sequential(*layers)

如果 num_layers=8 且开启 skip connection：

num_extra_layers = 2
len(self.net) = num_layers - num_extra_layers = 6

所以 self.net 中有 6 个 Linear + ReLU 模块，每个 Linear 都包了一层 nn.utils.weight_norm(...)。

Weight Normalization 放在哪里

源码只对 self.net 里的线性层使用权重归一化：

nn.utils.weight_norm(nn.Linear(input_dim, inner_dim))

它没有用在：

self.final_layer
self.skip_layer

权重归一化的作用是把线性层权重重参数化，常见形式是：

\mathbf{w} = g \frac{\mathbf{v}}{\|\mathbf{v}\|}

也就是把“方向”与“长度”拆开优化。对 DeepSDF 这类 MLP 来说，它有助于稳定隐式场的训练；这个仓库的 train_sdf.py 注释也提到 MPS 设备当时不能处理 weight normalization。

Skip Connection 怎么实现

前向传播里先保存原始输入：

input_data = x.clone().detach()

然后如果开启 skip connection：

for i in range(3):
    x = self.net[i](x)
x = self.skip_layer(x)
x = torch.hstack((x, input_data))
for i in range(self.num_layers - 5):
    x = self.net[3 + i](x)
sdf = self.final_layer(x)

按默认维度展开：

输入 x:              B x 67
self.net[0..2]:      B x 256
skip_layer:          B x 189
拼回 input_data:     B x (189 + 67) = B x 256
self.net[3..5]:      B x 256
final_layer:         B x 1

这里的关键是 skip_layer 不是直接把 256 维隐藏特征和 67 维输入拼成 323 维，而是先把隐藏特征压到：

256 - 67 = 189

再拼回原始输入，最后仍然保持 256 维，保证后续 self.net[3] 能继续接收 inner_dim=256 的输入。

为什么这里要做 skip connection？

DeepSDF 的输入中包含两个很重要的信息：

当前形状的身份信息：latent code z_i
当前查询点的位置：(x, y, z)

如果网络很深，原始坐标和 latent code 经过多层非线性变换后，低层输入信息可能被弱化。Skip connection 的作用就是在网络中间再次把原始输入拼回来，让后半段 MLP 继续直接看到：

[z_i, x, y, z]

这样做的直觉是：网络既能学习前几层抽出的高维特征，也不会忘掉“这个点到底在哪里、属于哪个形状”。

为什么 input_data = x.clone().detach() 值得注意？

detach() 会切断这条 skip 输入分支对原始输入张量的梯度追踪。这里的原始输入 x 本身是由 latent code 和坐标拼出来的，其中 latent code 是可训练变量。

因此，这个实现中：

前 3 层主干仍然会把梯度传回 latent code。
但是中间直接拼回来的 input_data 分支不会通过 shortcut 再额外把梯度传回输入。

这和“不 detach 的 skip connection”略有不同。复现论文时如果追求严格一致，可以专门对比 detach() 与不 detach() 的训练差异。

损失函数：Clamped L1 Loss 在代码里怎么截断

Clamped L1 gradient intuition

loss 定义在 utils/utils_deepsdf.py：

def SDFLoss_multishape(sdf, prediction, x_latent, sigma):
    l1 = torch.mean(torch.abs(prediction - sdf))
    l2 = sigma**2 * torch.mean(torch.linalg.norm(x_latent, dim=1, ord=2))
    loss = l1 + l2
    return loss, l1, l2

它由两部分组成：

\mathcal{L} = \underbrace{\frac{1}{B}\sum_b |\hat{s}_b - s_b|}_{\text{SDF 重建 L1}} + \underbrace{\sigma^2 \frac{1}{B}\sum_b \|\mathbf{z}_b\|_2}_{\text{latent code 正则}}

其中：

prediction 是网络预测的 SDF。
sdf 是真实 SDF。
x_latent 是当前 batch 查到的 latent code。
sigma_regulariser=0.01，所以正则权重是 0.0001。

`torch.clamp` 在哪里用

真实 SDF 在构造 DataLoader 前被截断：

if self.train_cfg['clamp']:
    data.data['sdf'] = torch.clamp(
        data.data['sdf'],
        -self.train_cfg['clamp_value'],
        self.train_cfg['clamp_value']
    )

预测 SDF 在训练循环里也被截断：

predictions = self.model(x)
if self.train_cfg['clamp']:
    predictions = torch.clamp(
        predictions,
        -self.train_cfg['clamp_value'],
        self.train_cfg['clamp_value']
    )

默认 clamp_value=0.1，所以训练时实际进入 L1 的值都落在：

[-0.1,\ 0.1]

为什么 clamp 会“切断过大距离的梯度”

对预测值 $\hat{s}$ 做：

\operatorname{clamp}(\hat{s}, -\delta, \delta)

可以理解为：

\operatorname{clamp}(\hat{s}) = \begin{cases} -\delta, & \hat{s} < -\delta \\ \hat{s}, & -\delta \le \hat{s} \le \delta \\ \delta, & \hat{s} > \delta \end{cases}

它对预测值的导数大致是：

\frac{\partial \operatorname{clamp}(\hat{s})}{\partial \hat{s}} = \begin{cases} 0, & \hat{s} < -\delta \\ 1, & -\delta < \hat{s} < \delta \\ 0, & \hat{s} > \delta \end{cases}

所以当网络预测已经跑到 0.1 外面时，torch.clamp 输出会卡在边界上，L1 loss 对这部分预测的梯度基本传不回去。

这使训练更关注物体表面附近的 SDF，因为 DeepSDF 最终真正要恢复的是：

f_\theta(\mathbf{z}, \mathbf{x}) = 0

也就是零水平集，而不是远离表面的大距离值。

一个具体数值例子

设 clamp_value = 0.1，真实 SDF 已经被截断，预测也会被截断。

原始预测 $\hat{s}$	clamp 后预测	真实 SDF	L1
`0.50`	`0.10`	`0.08`	`0.02`
`0.06`	`0.06`	`0.08`	`0.02`
`-0.30`	`-0.10`	`-0.07`	`0.03`

第二行在 [-0.1, 0.1] 内，梯度可以正常推动预测从 0.06 靠近 0.08。

第一行原始预测是 0.50，但 clamp 后固定为 0.10。从 loss 看，它只知道当前参与计算的是 0.10，而不是 0.50。因此这类远离表面的预测不会强迫网络精确拟合大距离。

警告

utils/utils_deepsdf.py 里还有一个自定义 clamp() 函数，但训练主流程没有用它；训练和验证实际使用的是 PyTorch 的 torch.clamp(...)。

数据加载器：Batch 如何拿到 latent 索引、坐标和 SDF

DeepSDF batch data flow

数据加载逻辑在 data/dataset_sdf.py：

def __getitem__(self, idx):
    latent_class = self.data['samples_latent_class'][idx, :]
    sdf = self.data['sdf'][idx]
    return latent_class, sdf

这里 latent_class 这个变量名略容易误导。它实际不是只有一个类别编号，而是一整行：

[shape_id, x, y, z]

也就是：

第 0 列：这个 SDF 样本属于哪个形状。
第 1 到 3 列：查询点坐标。

DataLoader 会把多条样本堆成 batch：

train_loader = DataLoader(
    train_data,
    batch_size=self.train_cfg['batch_size'],
    shuffle=True,
    drop_last=True
)

默认 batch_size=16，所以进入训练循环时：

batch[0]: B x 4, 每行 [shape_id, x, y, z]
batch[1]: B x 1, 每行 [sdf]

`generate_xy()` 如何组装模型输入

训练前调用：

x, y, latent_codes_indices_batch, latent_codes_batch = self.generate_xy(batch)

函数内部做三步：

latent_classes_batch = batch[0][:, 0].view(-1, 1).to(torch.long)
coords = batch[0][:, 1:]
latent_codes_batch = self.latent_codes[latent_classes_batch.view(-1)]

x = torch.hstack((latent_codes_batch, coords))
y = batch[1]

张量形状是：

latent_classes_batch: B x 1
coords:               B x 3
latent_codes_batch:   B x 64
x:                    B x 67
y:                    B x 1

也就是说，一个 batch 里可以混合来自不同形状的 SDF 样本。每条样本都通过自己的 shape_id 去查对应 latent code，然后与自己的坐标拼接。

训练时到底优化了什么

训练初始化时，模型参数和 latent code 是两个不同的优化对象：

self.optimizer_model = optim.Adam(
    self.model.parameters(),
    lr=self.train_cfg['lr_model'],
    weight_decay=0
)

self.latent_codes = utils_deepsdf.generate_latent_codes(
    self.train_cfg['latent_size'],
    samples_dict
)
self.optimizer_latent = optim.Adam(
    [self.latent_codes],
    lr=self.train_cfg['lr_latent'],
    weight_decay=0
)

所以训练 DeepSDF 时不是只训练一个 decoder，而是同时训练：

共享的 SDF decoder 参数 $\theta$ 。
每个训练形状自己的 latent code $\mathbf{z}_i$ 。

训练循环里也可以看到两个优化器一起更新：

    self.optimizer_model.zero_grad()
    self.optimizer_latent.zero_grad()

    x, y, latent_codes_indices_batch, latent_codes_batch = self.generate_xy(batch)

    predictions = self.model(x)
    loss_value, loss_rec, loss_latent = SDFLoss_multishape(...)

loss_value.backward()

self.optimizer_latent.step()
self.optimizer_model.step()

整体数据流是：

(\text{shape id}, x,y,z) \rightarrow \mathbf{z}_i \rightarrow [\mathbf{z}_i,x,y,z] \rightarrow f_\theta \rightarrow \hat{s} \rightarrow \mathcal{L}_{SDF}

训练阶段：网络权重和所有训练形状的 latent codes 如何同时优化？

DeepSDF 的训练可以理解成同时学习两类未知量：

\theta

和：

\{\mathbf{z}_1,\mathbf{z}_2,\ldots,\mathbf{z}_N\}

其中：

$\theta$ 是共享 decoder 的网络权重。
$\mathbf{z}_i$ 是第 $i$ 个训练形状自己的 latent code。
$N$ 是训练集中形状数量。

代码里这两类变量对应两个独立 optimizer：

self.optimizer_model = optim.Adam(
    self.model.parameters(),
    lr=self.train_cfg['lr_model'],
    weight_decay=0
)

self.optimizer_latent = optim.Adam(
    [self.latent_codes],
    lr=self.train_cfg['lr_latent'],
    weight_decay=0
)

注意这里不是“交替训练一个 epoch 网络、一个 epoch latent code”，而是在每个 batch 里同时反向传播，然后两个 optimizer 都执行一步。

一个 batch 的更新顺序是：

self.optimizer_model.zero_grad()
self.optimizer_latent.zero_grad()

x, y, latent_codes_indices_batch, latent_codes_batch = self.generate_xy(batch)
predictions = self.model(x)
loss_value, loss_rec, loss_latent = SDFLoss_multishape(...)

loss_value.backward()

self.optimizer_latent.step()
self.optimizer_model.step()

这段代码的含义是：

generate_xy(batch) 用 batch 中的 shape_id 查出对应的 latent code。
把查到的 latent code 和坐标拼成 [z_i, x, y, z]。
decoder 用这个输入预测 SDF。
loss 同时依赖 decoder 参数 $\theta$ 和当前 batch 用到的 latent codes。
loss_value.backward() 会同时计算两类梯度：

\frac{\partial \mathcal{L}}{\partial \theta}

和：

\frac{\partial \mathcal{L}}{\partial \mathbf{z}_i}

optimizer_latent.step() 更新训练形状的 latent table。
optimizer_model.step() 更新共享 decoder 的网络权重。

因此，训练阶段的本质是：

\min_{\theta,\{\mathbf{z}_i\}} \sum_i \sum_j \left| f_\theta(\mathbf{z}_i,\mathbf{x}_{ij}) - s_{ij} \right| + \lambda \|\mathbf{z}_i\|

也就是：让一个共享 decoder 学会“如何根据 latent code 和查询点预测 SDF”，同时也为每个训练形状找到一个合适的 latent code。

推理时到底优化了什么

DeepSDF 最反直觉的地方在推理阶段：面对一个新形状时，通常不再更新 decoder 权重，而是只优化一个新 latent code。

在这个仓库里，对应函数是 model/model_sdf.py 里的：

def infer_latent_code(self, cfg, pointcloud, sdf_gt, writer, latent_code_initial):
    latent_code = latent_code_initial.clone().detach().requires_grad_(True)
    optim = torch.optim.Adam([latent_code], lr=cfg['lr'])

这里新建的 optimizer 只接收：

[latent_code]

没有传入：

self.model.parameters()

这就说明推理阶段优化的是新形状的 latent code，而不是网络权重。

推理阶段：为什么冻结网络，只反向传播更新 latent code？

训练完成后，decoder 已经学到一个共享函数：

f_\theta(\mathbf{z},\mathbf{x}) \rightarrow s

可以把它理解成一个“形状空间解释器”：

输入不同的 latent code $\mathbf{z}$ ，代表不同形状。
输入查询点 $\mathbf{x}$ ，询问这个点到该形状表面的 signed distance。
输出 SDF 值 $s$ 。

推理时来了一个未知形状。我们没有这个形状的 latent code，所以先随机初始化一个：

latent_code = latent_code_initial.clone().detach().requires_grad_(True)

然后把它复制到所有观测点上：

latent_code_tile = torch.tile(latent_code, (pointcloud.shape[0], 1))
x = torch.hstack((latent_code_tile, pointcloud))

此时每个观测点的输入都是：

[同一个新 latent code, x, y, z]

接着用已经训练好的 decoder 预测 SDF：

predictions = self(x)

再和观测到的真实/近似 SDF 做 loss：

loss_value, l1, l2 = utils_deepsdf.SDFLoss_multishape(
    sdf_gt,
    predictions,
    x[:, :self.latent_size],
    sigma=cfg['sigma_regulariser']
)
loss_value.backward()
optim.step()

因为 optimizer 只管理 latent_code：

optim = torch.optim.Adam([latent_code], lr=cfg['lr'])

所以反向传播虽然会经过 decoder，但 optim.step() 只会更新这个新 latent code。

推理阶段的优化目标可以写成：

\min_{\mathbf{z}} \sum_j \left| f_{\theta^\*}(\mathbf{z},\mathbf{x}_j) - s_j \right| + \lambda \|\mathbf{z}\|

其中 $\theta^\*$ 是训练好的、固定不变的 decoder 权重。

这个过程的直觉是：

decoder 已经学会了“合理形状长什么样”。
新形状只需要在 latent space 里找到一个合适的位置。
反向传播不是为了改网络，而是为了搜索那个 latent code。
找到 $\mathbf{z}$ 后，就能查询任意 3D 点的 SDF，再用 Marching Cubes 提取 mesh。

这就是 DeepSDF 的精髓：训练时学习一个形状空间；推理时在这个形状空间里优化新形状的坐标。

复现时需要特别留意的源码细节

1. 维度注释和配置不一致

model_sdf.py 和 train_sdf.py 的注释里写了：

latent: 128
input: 131

但默认配置是：

latent_size: 64
input: 67

复现实验时应以 train_sdf.yaml 为准。

2. Loss multiplier 的写法比较危险

训练代码里写的是：

loss_value, loss_rec, loss_latent = \
    self.train_cfg['loss_multiplier'] * SDFLoss_multishape(...)

当前 YAML 中 loss_multiplier: 1，所以不会改变 tuple。但如果把它改成其他整数，Python 的 tuple * int 会复制 tuple，而不是逐项乘 loss；如果改成浮点数则会报错。

更稳妥的写法应该是：

loss_value, loss_rec, loss_latent = SDFLoss_multishape(...)
loss_value = self.train_cfg['loss_multiplier'] * loss_value

3. 验证函数里引用了全局 `train_cfg`

validate() 中有一段：

if train_cfg['clamp']:
    predictions = torch.clamp(predictions, -train_cfg['clamp_value'], train_cfg['clamp_value'])

这里用的是全局变量 train_cfg，而不是 self.train_cfg。如果以后把 Trainer 当模块导入使用，可能会因为全局 train_cfg 不存在而出问题。

更稳妥的是：

if self.train_cfg['clamp']:
    predictions = torch.clamp(
        predictions,
        -self.train_cfg['clamp_value'],
        self.train_cfg['clamp_value']
    )

4. `Dataset` 在初始化时直接把全部数据搬到 GPU

SDFDataset.__init__() 里有：

value = torch.from_numpy(samples_dict[obj_idx][key]).float().to(device)

这意味着数据集初始化时会直接把所有 SDF 样本放到 GPU 或当前设备上。小实验很方便，但大数据集会明显占显存。

如果数据规模变大，更常见的做法是：

Dataset 保持 CPU tensor。
每个 batch 在训练循环里 .to(device)。

小结

这个工程的 DeepSDF 训练逻辑可以概括为：

每条样本保存为 [shape_id, x, y, z] + sdf。
Batch 进入训练时，用 shape_id 查该形状的 latent code。
拼接 [latent_code, x, y, z] 作为 MLP 输入。
MLP 中间通过 skip connection 把原始输入再次拼回来。
主干线性层使用 weight normalization。
预测值和真实 SDF 都被 clamp 到 [-0.1, 0.1]。
Loss 是 clamped L1 reconstruction loss 加 latent code 正则。

对后续实验来说，最值得先改动并观察的变量是：

latent_size
inner_dim
clamp_value
是否保留 input_data.detach()
Dataset 是否提前把全部数据放到 GPU

默认配置先看清楚​

网络定义：MLP、Skip Connection 与 Weight Normalization​

MLP 主干怎么搭​

Weight Normalization 放在哪里​

Skip Connection 怎么实现​

损失函数：Clamped L1 Loss 在代码里怎么截断​

torch.clamp 在哪里用​

为什么 clamp 会“切断过大距离的梯度”​

数据加载器：Batch 如何拿到 latent 索引、坐标和 SDF​

generate_xy() 如何组装模型输入​

训练时到底优化了什么​

推理时到底优化了什么​

复现时需要特别留意的源码细节​

1. 维度注释和配置不一致​

2. Loss multiplier 的写法比较危险​

3. 验证函数里引用了全局 train_cfg​

4. Dataset 在初始化时直接把全部数据搬到 GPU​

小结​