【三维重建】【深度学习】NeuS代码Pytorch实现--训练阶段代码解析(下)

【三维重建】【深度学习】NeuS代码Pytorch实现–训练阶段代码解析(下)

论文提出了一种新颖的神经表面重建方法，称为NeuS，用于从2D图像输入以高保真度重建对象和场景。在NeuS中建议将曲面表示为有符号距离函数(SDF)的零级集，并开发一种新的体绘制方法来训练神经SDF表示，因此即使没有掩模监督，也可以实现更准确的表面重建。NeuS在高质量的表面重建方面的性能优于现有技术，特别是对于具有复杂结构和自遮挡的对象和场景。本篇博文将根据代码执行流程解析训练阶段具体的功能模块代码。

---

前言

在详细解析NeuS网络之前,首要任务是搭建NeuS【win10下参考教程】所需的运行环境，并完成模型的训练和测试，展开后续工作才有意义。
本博文将对NeuS训练阶段涉及的剩余部分功能代码模块进行解析，其他代码模块后续的博文将会陆续讲解。

博主将各功能模块的代码在不同的博文中进行了详细的解析，点击【win10下参考教程】，博文的目录链接放在前言部分。

---

NeRF网络

在exp_runner.py文件内class Runner的__init__函数中完成初始化，并将其传递给NeuS，以方便后续SDF网络的使用。

# nerf网络
self.nerf_outside = NeRF(**self.conf['model.nerf']).to(self.device)

在models/renderer.py文件内的render_core_outside函数上调用。

# 背景的概率密度和rbg
density, sampled_color = nerf(pts, dirs)

NeRF的定义在models/fields.py文件内。

class NeRF(nn.Module):
    def __init__(self,
                 D=8,           # 网络深度
                 W=256,         # 网络宽度
                 d_in=3,        # 采样点(世界坐标系下)的位置 输入channel
                 d_in_view=3,   # 单位方向向量 输入channel
                 multires=0,    # 采样点编码长度
                 multires_view=0,   # 单位方向向量编码长度
                 output_ch=4,   # 输出channel
                 skips=[4],     # 网络中间插入新输入的位置
                 use_viewdirs=False):   # 是否使用单位方向向量做额外输入
        super(NeRF, self).__init__()
        self.D = D
        self.W = W
        self.d_in = d_in
        self.d_in_view = d_in_view
        self.input_ch = 3
        self.input_ch_view = 3
        self.embed_fn = None
        self.embed_fn_view = None

        # 对采样点进行位置编码
        if multires > 0:
            embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
            self.embed_fn = embed_fn
            self.input_ch = input_ch
        # 对单位向量进行位置编码
        if multires_view > 0:
            embed_fn_view, input_ch_view = get_embedder(multires_view, input_dims=d_in_view)
            self.embed_fn_view = embed_fn_view
            self.input_ch_view = input_ch_view

        self.skips = skips
        self.use_viewdirs = use_viewdirs

        # 搭建网络层
        self.pts_linears = nn.ModuleList(
            [nn.Linear(self.input_ch, W)] +
            [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + self.input_ch, W) for i in range(D - 1)])

        self.views_linears = nn.ModuleList([nn.Linear(self.input_ch_view + W, W // 2)])

        if use_viewdirs:
            self.feature_linear = nn.Linear(W, W)
            self.alpha_linear = nn.Linear(W, 1)     # α
            self.rgb_linear = nn.Linear(W // 2, 3)  # rgb
        else:
            self.output_linear = nn.Linear(W, output_ch)

    def forward(self, input_pts, input_views):
        # 采样点(世界坐标系下)的位置编码
        if self.embed_fn is not None:
            input_pts = self.embed_fn(input_pts)
        # ,和单位方向向量编码
        if self.embed_fn_view is not None:
            input_views = self.embed_fn_view(input_views)

        h = input_pts
        for i, l in enumerate(self.pts_linears):
            h = self.pts_linears[i](h)
            h = F.relu(h)
            if i in self.skips:         # 在网络中间层插入新输入
                h = torch.cat([input_pts, h], -1)

        # 使用单位方向向量做额外输入
        if self.use_viewdirs:
            alpha = self.alpha_linear(h)
            feature = self.feature_linear(h)
            h = torch.cat([feature, input_views], -1)

            for i, l in enumerate(self.views_linears):
                h = self.views_linears[i](h)
                h = F.relu(h)

            rgb = self.rgb_linear(h)
            return alpha, rgb
        else:
            assert False

NeRF网络结构及其执行流程如下图所示。

与原始的nerf不同的地方在于，输入的γ(x)不仅仅表示对3D坐标点做位置编码(60)，还额外增加了距离系数(80)。

---

背景颜色渲染

mask模式没有背景颜色渲染过程
在models/renderer.py文件的render函数内。

# 背景模型
if self.n_outside > 0:
    # 前景采样点(粗&精)和背景采样点归并
    z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
    z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
    # 获得背景的相关数据
    ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf)
    # 采样点的颜色
    background_sampled_color = ret_outside['sampled_color']
    # 采样点的权重因子
    background_alpha = ret_outside['alpha']

render_core_outside通过nerf渲染前景以外的背景颜色。

def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
    """
    渲染背景
    """
    batch_size, n_samples = z_vals.shape

    # 相邻俩个采样点之间的距离
    dists = z_vals[..., 1:] - z_vals[..., :-1]  # [batch_size, n_smaples-1]
    dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)    # [batch_size, n_smaples]
    # 相邻俩个采样点的中点,看作新采样点
    mid_z_vals = z_vals + dists * 0.5

    # Section midpoints
    # 世界坐标系下相邻采样点的位置
    pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]    # [batch_size, n_samples, 3]
    # 原点到世界坐标系下采样点的距离,保留距离值大于1部分(球体外部的),距离小于1(在球体内部的)的则设置为1
    dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)     # [batch_size, n_samples, 1]
    # 对球体外的采样点做归一化处理,球体外的采样点做相当于没做处理,并额外增加了距离系数
    pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)       # [batch_size, n_samples, 4]
    # 单位方向向量
    dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)      # [batch_size, n_samples, 3]
    pts = pts.reshape(-1, 3 + int(self.n_outside > 0))              # [batch_size*n_samples, 4] 不是4不会调用该函数
    dirs = dirs.reshape(-1, 3)   # [batch_size*n_samples, 3]
    # 采样点的概率密度和rbg
    density, sampled_color = nerf(pts, dirs)
    # rgb归一化处理
    sampled_color = torch.sigmoid(sampled_color)        # [batch_size*n_samples, 3]
    # 权重因子
    alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)    # [batch_size,n_samples]
    alpha = alpha.reshape(batch_size, n_samples)    # [batch_size,n_samples]
    # 权重
    weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]    # [batch_size,n_samples]
    sampled_color = sampled_color.reshape(batch_size, n_samples, 3)     # [batch_size,n_samples,3]
    # 背景颜色=采样点的颜色*权重
    color = (weights[:, :, None] * sampled_color).sum(dim=1)    # [batch_size,1,3]
    if background_rgb is not None:
        # 额外添加背景颜色
        color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))

    return {
        'color': color,         # 背景颜色  # [batch_size,1,3]
        'sampled_color': sampled_color,     # 采样点的颜色     [batch_size*n_samples, 3]
        'alpha': alpha,         # 采样点的权因子               # [batch_size,n_samples]
        'weights': weights,     # 采样点的权重                 # [batch_size,n_samples]
    }

这里对部分世界坐标系下的采样点做了归一化，具有就是对球体外的采样点做了归一化，使得他们的采样点位置都落在球体表面，球体内的采样点位置不变。然后将处理完的采样点位置和与之对应单位方向向量输入到nerf网络，获得采样点的颜色和采样点的概率密度。归一化过程如图所示。
球体外采样点失去了原本的位置信息，因此额外添加了距离系数即距离的倒数与3D坐标位置一起输入到nerf中。

这里本来是相邻采样点的中点的归一化，但这不影响理解，就将中点视作新采样点。

---

RenderingNetwork网络

在exp_runner.py文件内class Runner的__init__函数中完成初始化，并将其传递给NeuS，以方便后续RenderingNetwork网络的使用。

# 渲染网络
self.color_network = RenderingNetwork(**self.conf['model.rendering_network']).to(self.device)

在models/renderer.py文件内的render_core函数上调用。

# 采样点的rgb
sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)

RenderingNetwork的定义在models/fields.py文件内。

class RenderingNetwork(nn.Module):
    def __init__(self,
                 d_feature,     # sdf输出的隐藏特征
                 mode,          # 模式选择
                 d_in,          # 输入channel(3D坐标+单位方向向量+梯度信息)
                 d_out,         # 输出channel
                 d_hidden,      # 中间网络层channel
                 n_layers,      # 网络层数
                 weight_norm=True,  # 权重归一化
                 multires_view=0,   # 位置编码长度
                 squeeze_out=True):     # 归一化
        super().__init__()
        self.mode = mode
        self.squeeze_out = squeeze_out
        # 各层的channel
        dims = [d_in + d_feature] + [d_hidden for _ in range(n_layers)] + [d_out]
        self.embedview_fn = None

        # 单位方向向量编码
        if multires_view > 0:
            # 编码函数 编码输出的channel
            embedview_fn, input_ch = get_embedder(multires_view)
            self.embedview_fn = embedview_fn
            dims[0] += (input_ch - 3)

        # 网络深度
        self.num_layers = len(dims)

        for l in range(0, self.num_layers - 1):
            out_dim = dims[l + 1]
            lin = nn.Linear(dims[l], out_dim)
            if weight_norm:
                # 权重归一化处理
                lin = nn.utils.weight_norm(lin)

            setattr(self, "lin" + str(l), lin)

        self.relu = nn.ReLU()

    def forward(self, points, normals, view_dirs, feature_vectors):
        # 3D坐标点、梯度信息、单位方向向量、sdf隐藏特征
        if self.embedview_fn is not None:
            # 单位方向向量编码
            view_dirs = self.embedview_fn(view_dirs)

        rendering_input = None

        if self.mode == 'idr':
            rendering_input = torch.cat([points, view_dirs, normals, feature_vectors], dim=-1)
        elif self.mode == 'no_view_dir':
            # 没有单位方向向量
            rendering_input = torch.cat([points, normals, feature_vectors], dim=-1)
        elif self.mode == 'no_normal':
            # 没有梯度信息
            rendering_input = torch.cat([points, view_dirs, feature_vectors], dim=-1)

        x = rendering_input

        for l in range(0, self.num_layers - 1):
            lin = getattr(self, "lin" + str(l))
            x = lin(x)
            # 除了最后一层,激活函数都是relu
            if l < self.num_layers - 2:
                x = self.relu(x)
        if self.squeeze_out:
            x = torch.sigmoid(x)
        return x

RenderingNetwork网络结构及其执行流程如下图所示。

输入只对单位方向向量做了编码γ(d)。

---

前景颜色渲染

在models/renderer.py文件的render函数内。

ret_fine = self.render_core(rays_o,
                            rays_d,
                            z_vals,
                            sample_dist,
                            self.sdf_network,
                            self.deviation_network,
                            self.color_network,
                            background_rgb=background_rgb,
                            background_alpha=background_alpha,
                            background_sampled_color=background_sampled_color,
                            cos_anneal_ratio=cos_anneal_ratio)

render_core通过RenderingNetwork渲染前景颜色。

def render_core(self,
                rays_o,
                rays_d,
                z_vals,
                sample_dist,
                sdf_network,
                deviation_network,
                color_network,
                background_alpha=None,
                background_sampled_color=None,
                background_rgb=None,
                cos_anneal_ratio=0.0):
    batch_size, n_samples = z_vals.shape
    # 相邻俩个采样点之间的距离
    dists = z_vals[..., 1:] - z_vals[..., :-1]  # [batch_size, n_smaples-1]
    dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)    # [batch_size, n_smaples]
    # 相邻俩个采样点的中点,看作新采样点
    mid_z_vals = z_vals + dists * 0.5
    # 世界坐标系下相邻采样点的位置
    pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # [batch_size, n_smaples, 3]
    # 单位方向向量
    dirs = rays_d[:, None, :].expand(pts.shape)

    pts = pts.reshape(-1, 3)     # [batch_size*n_samples, 3]
    dirs = dirs.reshape(-1, 3)   # [batch_size*n_samples, 3]

    sdf_nn_output = sdf_network(pts)    # [batch_size*n_samples,257]
    # sdf值
    sdf = sdf_nn_output[:, :1]          # [batch_size,n_samples,1]
    # sdf隐藏特征
    feature_vector = sdf_nn_output[:, 1:]               # [batch_size,n_samples,256]
    # sdf相关的梯度数据
    gradients = sdf_network.gradient(pts).squeeze()     # [batch_size*n_samples,3]
    # 采样点的rgb
    sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)   # [batch_size,n_samples,3]
    # 单参数网络
    inv_s = deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)       # [1,1]
    inv_s = inv_s.expand(batch_size * n_samples, 1)         # [batch_size*n_samples,1]

    # 单位方向向量和梯度信息相乘
    true_cos = (dirs * gradients).sum(-1, keepdim=True)     # [batch_size*n_samples,3]=>[batch_size*n_samples,1]
    # 个人理解:单独的梯度信是为了能够强化sdf部分权重的训练更新 iter_cos<=0
    iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
                 F.relu(-true_cos) * cos_anneal_ratio)                   # [batch_size*n_samples,1]
    # 估计采样点的sdf取值区间(末端) 小于sdf
    estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5     # [batch_size*n_samples,1]
    # 估计采样点的sdf取值区间(始端) 大于sdf
    estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5     # [batch_size*n_samples,1]

    # 归一化处理
    # 采样点sdf区段的始端
    prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)    # [batch_size*n_samples,1]
    # 采样点sdf区段的末端
    next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)    # [batch_size*n_samples,1]

    # 采样点的sdf区段的长度
    p = prev_cdf - next_cdf
    # 采样点sdf始端值
    c = prev_cdf
    # 采样点权重因子=区段的长度/始端值,范围在0~1
    alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)     # [batch_size,n_samples]
    # 采样点的长度
    pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)   # [batch_size,n_samples]
    # 在球体内部
    inside_sphere = (pts_norm < 1.0).float().detach()   # [batch_size,n_samples]
    # 放宽限制的球体内部
    relax_inside_sphere = (pts_norm < 1.2).float().detach()

    if background_alpha is not None:
        # background_alpha和background_sampled_color是nerf网络输出的关于背景的采样点的权重因子和rgb
        # 采样点的权重因子保留在球体内部的部分,在球体外的用关于背景的采样点(前景部分的)权重因子替换
        alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
        # 前景采样点权重因子和关于背景的采样点(背景部分)的权重因子拼接
        alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)     # [batch_size,n_samples+n_outside]

        # 采样点的rgb保留在球体内部的部分,在球体外的用关于背景的采样点(前景部分的)rgb
        sampled_color = sampled_color * inside_sphere[:, :, None] +\
                        background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]     # [batch_size,n_samples,3]
        # 前景采样点rgb和关于背景的采样点(背景部分)的rgb拼接
        sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)  # [batch_size,n_samples+n_outside,3]


    # 采样点权重
    weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]    # [batch_size,n_samples+n_outside]
    # 权重总和
    weights_sum = weights.sum(dim=-1, keepdim=True)     # [batch_size,1]
    # 前景的rgb
    color = (sampled_color * weights[:, :, None]).sum(dim=1)    # [batch_size,1,3]
    if background_rgb is not None:
        # 额外背景rgb
        color = color + background_rgb * (1.0 - weights_sum)

    # 梯度的信息的二范式的平方
    gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,     # [batch_size,n_samples]
                                        dim=-1) - 1.0) ** 2
    # 保留在放宽限制的球体内部采样点的梯度
    gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)

    return {
        'color': color,                     # 前景颜色  [batch_size,n_samples+n_outside,3]
        'sdf': sdf,                         # 采样点(中点)sdf  [batch_size*n_samples,1]
        'dists': dists,                     # 相邻采样点的距离  [batch_size,n_samples]
        'gradients': gradients.reshape(batch_size, n_samples, 3),       # 梯度信息 [batch_size,n_samples，3]
        's_val': 1.0 / inv_s,               # 可训练参数 [batch_size*n_samples,1]
        'mid_z_vals': mid_z_vals,           # 采样点(中点)位置         [batch_size,n_samples+n_outside]
        'weights': weights,                 # 采样点(中点)权重                     [batch_size,n_samples+n_outside]
        'cdf': c.reshape(batch_size, n_samples),       # 采样点(中点)sdf始端值     [batch_size,n_samples]
        'gradient_error': gradient_error,              # 保留在放宽限制的球体内部采样点的梯度  [batch_size,n_samples]
        'inside_sphere': inside_sphere                 # 采样点是否在球体内 [batch_size,n_samples]
    }

前景采样点的权重因子示意图，以及加入背景采样点后的权重因子示意图。

前景颜色渲染过程，对于前景的采样点，inside_sphere的绿色部分表示在球体内的有效采样点，红色部分表示在球体外的采样点。保留绿色采样点部分的权重因子(alpha)，红色采样点部分的权重因子用在背景颜色渲染过程(nerf)提供的权重因子(background_alpha)替换，并且加入背景采样点的权重因子(background_alpha)。

加不加入背景权重因子取决的使用者采用的模式，假如是mask模式则没有背景采样点。
强调一点，背景颜色渲染过程中是提供包括前景采样点和背景采样点的权重因子，前景采样点权重因子可以理解成备胎，用于取代前景颜色渲染过程产生的超出球体的采样点的无效权重因子。

---

可训练参数

在models/renderer.py文件内的render_core函数上调用。

# 可训练参数网络
inv_s = deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)       # [1,1]

RenderingNetwork的定义在models/fields.py文件内。

class SingleVarianceNetwork(nn.Module):
    def __init__(self, init_val):
        super(SingleVarianceNetwork, self).__init__()
        self.register_parameter('variance', nn.Parameter(torch.tensor(init_val)))

    def forward(self, x):
        return torch.ones([len(x), 1]) * torch.exp(self.variance * 10.0)

---

更新NeuS网络

train函数这里只展示了光训练的代码部分，它部分代码包括了模型的保存、新视角的渲染和图像mesh模型的生成将在放到测试章节继续讲解。

def train(self):
    self.writer = SummaryWriter(log_dir=os.path.join(self.base_exp_dir, 'logs'))
    # 更新学习率
    self.update_learning_rate()
    # 训练需要迭代的次数
    res_step = self.end_iter - self.iter_step
    # 对图像序号进行随机排序
    image_perm = self.get_image_perm()

    for iter_i in tqdm(range(res_step)):
        # 随机生成rays
        data = self.dataset.gen_random_rays_at(image_perm[self.iter_step % len(image_perm)], self.batch_size)
        # rays_o(光心)、rays_v(单位方向向量)、true_rgb(rays穿过图像的rgb)、mask(前后景标签)
        rays_o, rays_d, true_rgb, mask = data[:, :3], data[:, 3: 6], data[:, 6: 9], data[:, 9: 10]
        # 最近点和最远点
        near, far = self.dataset.near_far_from_sphere(rays_o, rays_d)

        # 使用背景色
        background_rgb = None
        if self.use_white_bkgd:
            background_rgb = torch.ones([1, 3])

        # mask权重系数:这里的目的是是否需要区分背景和前景
        if self.mask_weight > 0.0:
            mask = (mask > 0.5).float()
        else:
            mask = torch.ones_like(mask)
        # 为前景的像素点总数
        mask_sum = mask.sum() + 1e-5

        render_out = self.renderer.render(rays_o, rays_d, near, far,
                                          background_rgb=background_rgb,
                                          cos_anneal_ratio=self.get_cos_anneal_ratio())

        color_fine = render_out['color_fine']       # 前景颜色
        s_val = render_out['s_val']                 # 可训练参数
        cdf_fine = render_out['cdf_fine']           # 采样点df值
        gradient_error = render_out['gradient_error']  # 梯度误差
        weight_max = render_out['weight_max']          # 最大权重采样带你
        weight_sum = render_out['weight_sum']          # 采样点权重总数

        # 颜色Loss(有mask,只计算mask部分)
        color_error = (color_fine - true_rgb) * mask
        color_fine_loss = F.l1_loss(color_error, torch.zeros_like(color_error), reduction='sum') / mask_sum

        # 评价指标
        psnr = 20.0 * torch.log10(1.0 / (((color_fine - true_rgb)**2 * mask).sum() / (mask_sum * 3.0)).sqrt())

        # sdf梯度loss
        eikonal_loss = gradient_error
        # 光线总权重值要接近mask
        mask_loss = F.binary_cross_entropy(weight_sum.clip(1e-3, 1.0 - 1e-3), mask)
        
        # 颜色loss+sdf梯度loss+标签loss
        loss = color_fine_loss +\
               eikonal_loss * self.igr_weight +\
               mask_loss * self.mask_weight

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

---

总结

尽可能简单、详细的介绍NeuS训练阶段剩余代码：nerf和renderingnetwork网络的结构和用途，以及neus网络的更新。后续会讲解测试验证阶段的代码