PyTorch3D 實作可微分體積渲染與最佳化

– 12 分鐘閱讀

隨著深度學習技術的發展，可微分體積渲染已成為一個熱門研究方向。本文將介紹如何利用 PyTorch3D 這個強大的工具，實作可微分體積渲染，並探討其在神經網路訓練中的應用。首先，我們會定義體積資料的結構，包含密度和顏色資訊，並使用張量表示。接著，我們將介紹體積渲染器、取樣器以及如何結合這些模組進行渲染。最後，我們將深入探討如何使用 Huber 損失函式來最佳化體積渲染網路，使其能夠根據多視角影像學習和重建三維場景。

在 PyTorch3D 中，體積資料通常以密度和特徵（例如顏色）的形式表示，並儲存在 Volumes 結構中。我們可以使用多層感知器或卷積神經網路來預測這些體積屬性。為了從體積中取樣資訊，我們需要定義一個 VolumeSampler，它可以根據相機引數和射線方向，在體積空間中取樣密度和特徵值。取樣過程可以使用不同的插值方法，例如雙線性插值。接著，raymarcher 模組會根據取樣點的密度和顏色計算每個畫素的顏色值，最終生成渲染影像。為了訓練體積渲染網路，我們可以使用可微分渲染技術，將渲染過程嵌入到神經網路中，並使用例如 Huber 損失函式來計算渲染影像與目標影像之間的差異，進而透過梯度下降最佳化網路引數。這個過程可以讓網路學習如何從多視角影像中重建三維場景的密度和顏色分佈。

體積資料結構

體積資料通常由一組3D格點組成，每個格點都有一個密度值和一個顏色值。密度值表示格點所在位置的物體密度，顏色值表示格點所在位置的物體顏色。體積資料可以用一個4D張量表示，形狀為（批次大小，通道數，高度，寬度，深度）。

import torch

batch_size = 10
densities = torch.zeros([batch_size, 1, 64, 64, 64]).to(device)
colors = torch.zeros([batch_size, 3, 64, 64, 64]).to(device)
voxel_size = 0.1

體積渲染器

體積渲染器是用於渲染體積資料的模組。它可以根據體積資料和相機引數生成最終的渲染影像。體積渲染器可以用不同的方式實作，例如使用GPU加速的體積渲染演算法或使用神經網路進行渲染。

class Volumes:
    def __init__(self, densities, features, voxel_size):
        self.densities = densities
        self.features = features
        self.voxel_size = voxel_size

volumes = Volumes(densities=densities, features=colors, voxel_size=voxel_size)

體積取樣器

體積取樣器是用於從體積資料中取樣點的模組。它可以根據相機引數和體積資料生成取樣點的密度值和顏色值。體積取樣器可以用不同的方式實作，例如使用最近鄰居插值或雙線性插值。

class VolumeSampler:
    def __init__(self, volumes, sample_mode):
        self.volumes = volumes
        self.sample_mode = sample_mode

    def __call__(self, ray_bundle):
        # 根據相機引數和體積資料生成取樣點的密度值和顏色值
        rays_densities, rays_features = self.sample(ray_bundle)
        return rays_densities, rays_features

volume_sampler = VolumeSampler(volumes=volumes, sample_mode="bilinear")
rays_densities, rays_features = volume_sampler(ray_bundle)

可微分體積渲染

可微分體積渲染是指體積渲染的過程可以被微分的技術。這意味著體積渲染的結果可以被用於計算梯度和最佳化神經網路。可微分體積渲染可以用於實作端到端的學習和最佳化，例如學習體積資料的表示或最佳化體積渲染的引數。

# 定義體積渲染的損失函式
def loss(rays_densities, rays_features):
    # 計算體積渲染的損失
    loss = torch.mean((rays_densities - rays_features) ** 2)
    return loss

# 最佳化體積渲染的引數
optimizer = torch.optim.Adam(volume_sampler.parameters(), lr=0.001)
for epoch in range(100):
    # 前向傳播
    rays_densities, rays_features = volume_sampler(ray_bundle)
    loss_value = loss(rays_densities, rays_features)
    
    # 反向傳播
    optimizer.zero_grad()
    loss_value.backward()
    optimizer.step()

圖表翻譯：

  graph LR
    A[體積資料] -->|密度值和顏色值|> B[體積渲染器]
    B -->|渲染影像|> C[最終結果]
    C -->|損失函式|> D[最佳化器]
    D -->|梯度下降|> E[體積渲染引數]
    E -->|更新|> B

6. 印出 rays_densities 和 rays_features 的形狀

print('rays_densities 的形狀 = ', rays_densities.shape)
print('rays_features 的形狀 = ', rays_features.shape)

7. 結果

結果如下所示。首先，我們注意到包的大小是 10，也就是說它由 10 個相機組成，這解釋了張量的第一個維度。接下來，每個影像畫素都有一個射線，而相機的解析度是 64x64。每個射線上的點數是 50，這解釋了張量的第四個維度。每個密度可以由一個數字表示，而每個顏色需要三個數字來表示 RGB 值：

rays_densities 的形狀 = torch.Size([10, 64, 64, 50, 1])
rays_features 的形狀 = torch.Size([10, 64, 64, 50, 3])

8. 儲存密度和顏色

最後，我們儲存密度和顏色，因為它們將在下一節中使用：

torch.save({'rays_densities': rays_densities, 'rays_features': rays_features}, 'volume_sampling.pt')

現在你已經對體積取樣有了一個清晰的理解。你知道它是什麼以及它的好處。在下一節中，你將學習如何使用這些密度和顏色來生成一組相機的 RGB 顏色影像。

射線漫遊器的觀察

現在我們已經有了射線選擇器中所有點的顏色和密度值，我們需要確定如何使用它們來生成投影影像的畫素值。

射線漫遊器的觀察

在這個部分，我們將討論如何將射線點的密度和顏色值轉換為影像的 RGB 值。這個過程模擬了影像形成的物理過程。

在這個部分，我們將討論一個非常簡單的模型，其中影像每個畫素的 RGB 值是對應射線點的顏色值的加權和。如果我們將密度視為佔用或不透明度的機率，那麼每個射線點的入射光強度等於 a = ∏i=1(1 - p i)，其中 p i 是密度。給定這個點被某個物體佔用機率為 p i，我們可以計算出反射光的預期強度為 w i = ap i。我們使用 w i 作為顏色值的加權。通常，權重會使用 softmax 操作進行標準化，以確保所有權重的總和為 1。

PyTorch3D 庫中包含了多個射線漫遊器的實作。下面的原始碼位於 GitHub 上的書籍儲存庫中，檔名為 understand_ray_marcher.py。

1. 匯入必要的套件

import torch
from pytorch3d.renderer.implicit.raymarching import EmissionAbsorptionRaymarcher

2. 載入射線點的密度和顏色值

checkpoint = torch.load('volume_sampling.pt')
rays_densities = checkpoint.get('rays_densities')
rays_features = checkpoint.get('rays_features')

3. 定義射線漫遊器和傳入密度和顏色值

ray_marcher = EmissionAbsorptionRaymarcher(...)

… (未完成)

內容解密：

以上程式碼的作用是定義一個射線漫遊器，並傳入射線點的密度和顏色值。射線漫遊器的作用是模擬影像形成的物理過程，將射線點的密度和顏色值轉換為影像的 RGB 值。

圖表翻譯：

此圖示為射線漫遊器的過程，包括射線點的密度和顏色值的轉換為影像的 RGB 值。圖中顯示了射線點的密度和顏色值如何被使用來計算影像的 RGB 值。

  flowchart TD
    A[射線點的密度和顏色值] --> B[射線漫遊器]
    B --> C[影像的 RGB 值]
    C --> D[影像形成]

疾馳於無限可能之路：體驗卷積神經網路的魅力

在人工智慧的世界中，卷積神經網路（Convolutional Neural Networks, CNNs）是一種非常重要的技術，尤其是在影像和視覺領域。今天，我們將探討卷積神經網路的基本原理和其在實際應用中的魅力。

卷積神經網路的基本原理

卷積神經網路是一種特殊的神經網路，主要用於處理影像和視覺資料。它的基本原理是使用卷積層（Convolutional Layers）和池化層（Pooling Layers）來提取影像中的特徵。

卷積層的工作原理

卷積層的工作原理是使用一組卷積核（Filters）來掃描影像，提取影像中的區域性特徵。每個卷積核都有一個大小和深度，大小決定了掃描的範圍，深度決定了提取的特徵維度。

池化層的工作原理

池化層的工作原理是使用一組池化函式（Pooling Functions）來降低影像的維度，減少資料的量。常用的池化函式包括最大池化（Max Pooling）和平均池化（Average Pooling）。

卷積神經網路的實際應用

卷積神經網路在實際應用中有非常廣泛的使用，包括影像分類、物體偵測、分割和生成等。

影像分類

影像分類是卷積神經網路的一個基本應用，包括將影像分類為不同的類別。例如，將影像分類為狗、貓、車等。

物體偵測

物體偵測是卷積神經網路的一個重要應用，包括偵測影像中的物體位置和類別。例如，偵測影像中的人、車、腳踏車等。

分割

分割是卷積神經網路的一個重要應用，包括將影像分割為不同的區域。例如，將影像分割為前景和背景。

生成

生成是卷積神經網路的一個重要應用，包括生成新的影像。例如，生成新的狗、貓、車等影像。

未來，卷積神經網路將繼續在影像和視覺領域中發揮重要作用。隨著技術的進步，卷積神經網路將能夠處理更加複雜的影像和視覺資料，實作更加精確的分類、偵測、分割和生成等。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision
import torchvision.transforms as transforms

# 定義卷積神經網路模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = nn.functional.relu(nn.functional.max_pool2d(self.conv1(x), 2))
        x = nn.functional.relu(nn.functional.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)

# 初始化模型、損失函式和最佳化器
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 訓練模型
for epoch in range(10):
    for x, y in train_loader:
        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()
    print('Epoch {}: Loss = {:.4f}'.format(epoch+1, loss.item()))

圖表翻譯：

此圖示為卷積神經網路的基本架構，包括卷積層、池化層、全連線層等。卷積層使用卷積核來掃描影像，提取影像中的區域性特徵。池化層使用池化函式來降低影像的維度，減少資料的量。全連線層使用全連線神經網路來進行分類或回歸等任務。

  graph LR
    A[影像] -->|掃描|> B[卷積層]
    B -->|降維|> C[池化層]
    C -->|全連線|> D[全連線層]
    D -->|輸出|> E[分類或回歸]

卷積神經網路的體積渲染

1. 初始化渲染引擎

首先，我們需要初始化渲染引擎，設定渲染的尺寸和其他引數。

render_size = 128
volume_extent_world = 3.0

2. 定義光線取樣器

接下來，我們需要定義光線取樣器，負責生成光線和取樣點。

raysampler = NDCGridRaysampler(
    image_width=render_size,
    image_height=render_size,
    n_pts_per_ray=150,
    min_depth=0.1,
    max_depth=volume_extent_world,
)

3. 建立光線遞迴器

然後，我們需要建立光線遞迴器，負責計算光線的強度和顏色。

raymarcher = EmissionAbsorptionRaymarcher()

4. 定義體積渲染器

接下來，我們需要定義體積渲染器，負責整合光線取樣器和光線遞迴器。

renderer = VolumeRenderer(
    raysampler=raysampler,
    raymarcher=raymarcher,
)

5. 定義體積模型

最後，我們需要定義體積模型，負責儲存體積的密度和顏色。

class VolumeModel(torch.nn.Module):
    def __init__(self, renderer, volume_size=[64] * 3, voxel_size=0.1):
        super().__init__()
        self.log_densities = torch.nn.Parameter(-4.0 * torch.ones(1, *volume_size))

6. 渲染體積

現在，我們可以使用體積渲染器和體積模型來渲染體積。

generate_cow_renders(num_views=40)

內容解密：

  • NDCGridRaysampler 是一個光線取樣器，負責生成光線和取樣點。
  • EmissionAbsorptionRaymarcher 是一個光線遞迴器，負責計算光線的強度和顏色。
  • VolumeRenderer 是一個體積渲染器，負責整合光線取樣器和光線遞迴器。
  • VolumeModel 是一個體積模型，負責儲存體積的密度和顏色。

圖表翻譯：

  flowchart TD
    A[體積渲染器] --> B[光線取樣器]
    B --> C[光線遞迴器]
    C --> D[體積模型]
    D --> E[渲染體積]

這個圖表展示了體積渲染的流程，從體積渲染器到光線取樣器、光線遞迴器、體積模型，最終到渲染體積。

根據PyTorch的體積渲染網路

在這個章節中，我們將探討如何使用PyTorch實作體積渲染網路。體積渲染是一種將3D體積資料轉換為2D影像的技術，廣泛應用於計算機視覺、醫學影像分析等領域。

網路架構

體積渲染網路的架構如下：

import torch
import torch.nn as nn

class VolumeRenderer(nn.Module):
    def __init__(self, volume_size, voxel_size, renderer):
        super(VolumeRenderer, self).__init__()
        self.log_densities = torch.nn.Parameter(torch.zeros(1, *volume_size))
        self.log_colors = torch.nn.Parameter(torch.zeros(3, *volume_size))
        self._voxel_size = voxel_size
        self._renderer = renderer

    def forward(self, cameras):
        batch_size = cameras.R.shape[0]
        densities = torch.sigmoid(self.log_densities)
        colors = torch.sigmoid(self.log_colors)
        volumes = Volumes(
            densities=densities[None].expand(batch_size, *self.log_densities.shape),
            features=colors[None].expand(batch_size, *self.log_colors.shape),
            voxel_size=self._voxel_size,
        )
        return self._renderer(cameras=cameras, volumes=volumes)[0]

在這個架構中，VolumeRenderer類別繼承自PyTorch的nn.Module類別。它包含了兩個可學習的引數：log_densities和log_colors，分別代表體積的密度和顏色。_voxel_size和_renderer是網路的超引數。

前向傳播

在前向傳播中，網路接收相機引數cameras作為輸入，然後計算體積的密度和顏色。體積的密度和顏色是使用sigmoid函式啟用的。

def forward(self, cameras):
    batch_size = cameras.R.shape[0]
    densities = torch.sigmoid(self.log_densities)
    colors = torch.sigmoid(self.log_colors)
    volumes = Volumes(
        densities=densities[None].expand(batch_size, *self.log_densities.shape),
        features=colors[None].expand(batch_size, *self.log_colors.shape),
        voxel_size=self._voxel_size,
    )
    return self._renderer(cameras=cameras, volumes=volumes)[0]

損失函式

在這個例子中，我們使用Huber損失函式作為網路的損失函式。Huber損失函式是一種robust的損失函式，可以抵禦異常值的影響。

def huber(x, y, scaling=0.1):
    diff_sq = (x - y) ** 2
    loss = ((1 + diff_sq / (scaling ** 2)) ** 0.5) - 1
    return loss

訓練

網路的訓練過程包括了前向傳播、計算損失、反向傳播和更新引數等步驟。

# 初始化網路和最佳化器
net = VolumeRenderer(volume_size, voxel_size, renderer)
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)

# 訓練網路
for epoch in range(100):
    # 前向傳播
    outputs = net(cameras)
    loss = huber(outputs, targets)

    # 反向傳播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

在這個例子中，我們使用Adam最佳化器來更新網路的引數。訓練過程包括了100個epoch，每個epoch都會前向傳播、計算損失、反向傳播和更新引數。

圖表翻譯：

  graph LR
    A[相機引數] --> B[體積渲染網路]
    B --> C[體積密度和顏色]
    C --> D[渲染結果]
    D --> E[損失函式]
    E --> F[最佳化器]
    F --> G[網路引數更新]

這個圖表展示了體積渲染網路的訓練過程，包括了相機引數的輸入、體積渲染網路的前向傳播、體積密度和顏色的計算、渲染結果的輸出、損失函式的計算、最佳化器的更新和網路引數的更新。

深度學習模型最佳化流程

在深度學習模型的最佳化過程中，瞭解每一步驟的重要性和實作方法是非常關鍵的。以下是最佳化流程的詳細步驟：

1. 損失函式計算

計算損失函式是模型最佳化的核心。損失函式的設計直接影響到模型的最佳化方向和效率。一個常見的損失函式是均方差損失函式，然而在這裡，我們看到了一個更為複雜的損失函式計算過程：

loss = (torch.norm(target_silhouettes - rendered_silhouettes) 
        - 1) * float(scaling)

這個損失函式計算了目標輪廓和渲染輪廓之間的差異，並根據scaling引數進行調整。

2. 儀器遷移

將所有資料和模型遷移到適當的裝置（如GPU）上，以加速計算：

target_cameras = target_cameras.to(device)
target_images = target_images.to(device)
target_silhouettes = target_silhouettes.to(device)

3. 初始化體積模型

定義一個體積模型的例項，指定渲染器、體積大小和體素大小：

volume_size = 128
volume_model = VolumeModel(
    renderer,
    volume_size=[volume_size] * 3,
    voxel_size=volume_extent_world / volume_size,
).to(device)

4. 最佳化器設定

設定最佳化器，包括學習率和最佳化器型別。在這裡，學習率設定為0.1，使用Adam最佳化器，並設定了300次迭代：

lr = 0.1
optimizer = torch.optim.Adam(volume_model.parameters(), lr=lr)
batch_size = 10
n_iter = 300

這些步驟構成了深度學習模型最佳化的基本流程，包括損失函式的計算、資料和模型的遷移、體積模型的初始化和最佳化器的設定。每一步驟都對模型的最終效能有著重要的影響。

主題標題

最佳化迴圈中的體積渲染

段落標題

體積渲染的最佳化過程

在最佳化迴圈中，體積渲染的目的是將體積的密度和顏色渲染出來，並與觀察到的多視角影像進行比較。這個過程使用了Huber損失函式來計算渲染出的影像和觀察到的影像之間的差異。

for iteration in range(n_iter):
    if iteration == round(n_iter * 0.75):
        print('減少學習率10倍...')
        optimizer = torch.optim.Adam(volume_model.parameters(), lr=lr * 0.1)
    optimizer.zero_grad()
    batch_idx = torch.randperm(len(target_cameras))[:batch_size]
    batch_cameras = FoVPerspectiveCameras(
        R=target_cameras.R[batch_idx],
        T=target_cameras.T[batch_idx],
        znear=target_cameras.znear[batch_idx],
        zfar=target_cameras.zfar[batch_idx],
        aspect_ratio=target_cameras.aspect_ratio[batch_idx],
        fov=target_cameras.fov[batch_idx],
        device=device,
    )
    rendered_images, rendered_silhouettes = volume_model(batch_cameras).split([3, 1], dim=-1)
    sil_err = huber(rendered_silhouettes[..., 0], target_silhouettes[batch_idx]).abs().mean()
    color_err = huber(rendered_images, target_images[batch_idx]).abs().mean()
    loss = color_err + sil_err

內容解密：

這段程式碼的主要目的是進行體積渲染的最佳化。首先，它會根據目前的迭代次數來調整學習率。當迭代次數達到總迭代次數的75%時，學習率會減少10倍。然後，它會選擇一批隨機的攝像機，並使用FoVPerspectiveCameras類別來建立一批攝像機物件。接下來，它會使用volume_model來渲染出體積的影像和輪廓，並計算渲染出的影像和觀察到的影像之間的差異。最後，它會計算總損失，並使用Huber損失函式來最佳化體積渲染的過程。

圖表翻譯：

  flowchart TD
    A[開始] --> B[選擇攝像機]
    B --> C[渲染體積]
    C --> D[計算差異]
    D --> E[最佳化體積渲染]
    E --> F[結束]

這個流程圖表現了體積渲染的最佳化過程。首先，選擇一批隨機的攝像機。然後，渲染出體積的影像和輪廓。接下來，計算渲染出的影像和觀察到的影像之間的差異。最後，使用Huber損失函式來最佳化體積渲染的過程。

從底層實作到高階應用的全面檢視顯示，體積渲染技術在近年來的發展中取得了顯著的進步。透過多維度效能指標的實測分析，根據PyTorch的體積渲染網路展現了其在處理3D資料和生成逼真影像方面的優勢。分析PyTorch3D提供的工具可以發現，從體積資料的表示、取樣到渲染的整個流程，都已高度模組化和可微分，這極大地簡化了開發流程，並為端到端學習和最佳化提供了堅實的基礎。然而，高解析度體積資料的處理仍然存在計算瓶頸，這也是未來研究需要重點突破的方向。同時，如何有效地整合先驗知識和領域專長到深度學習模型中，以提升渲染效率和影像品質，也是一個值得深入探索的課題。玄貓認為，隨著硬體效能的提升和演算法的持續創新，體積渲染技術將在更多領域展現其巨大的應用潛力，例如醫療影像分析、虛擬實境和增強現實等，並推動相關產業的快速發展。

玄貓

玄貓

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。