使用PolyGen和PyTorch生成3D模型

介紹

深度學習研究的一個新興領域是致力於將DL技術應用於3D幾何和計算機圖形應用程式， 對於希望自己嘗試3D深度學習的PyTorch使用者而言，一個叫Kaolin 庫值得研究。 對於TensorFlow使用者，還有TensorFlow Graphics庫。 3D技術中一個特別熱門的子領域是3D模型的生成。 創造性地組合3D模型，從影像快速生成3D模型，以及為其他機器學習應用程式和模擬建立綜合資料，這只是3D模型生成的眾多用例中的少數幾個。

使用top-p = 0.9的核取樣和地面真實網格（藍色）生成的影像條件樣本（黃色）。

但是，在3D深度學習研究領域，為資料選擇合適的表示是成功的一半。在計算機視覺中，資料的結構非常簡單：由密集畫素組成的影像，這些畫素整齊均勻地排列在精確的網格中。 3D資料的世界沒有這種一致性。 3D模型可以表示為體素，點雲，網格，多檢視影像集等。這些輸入表示形式也各有其缺點。例如，體素儘管計算成本高，但輸出解析度低。點雲沒有編碼表面或其法線的概念，因此不能僅從點雲中唯一地推斷拓撲。網格也不會對拓撲進行唯一編碼，因為可以細分任何網格以生成相似的曲面。這些缺點促使DeepMind的研究人員建立了PolyGen，這是一種用於網格的神經生成模型，可以共同估計模型的面和頂點以直接生成網格。官方實現可在DeepMind GitHub上獲得。https://github.com/deepmind/deepmind-research/tree/master/polygen

研究

3D重建問題和3D-R2N2方法

當今非常經典的PointNet論文為建模點雲資料（例如3D模型的尖端）提供了藍圖。它是一種通用演算法，不會對3D模型的面或佔用進行建模，因此無法僅使用PointNet來生成3D-R2N2採用的體素方法將我們都熟悉的2D卷積擴充套件到3D，並通過自然地從RGB影像生成水密網格。但是，體素表示在更高的空間解析度下在計算上變得昂貴，從而有效地限制了它可以生成的網格的大小。

通過變形模板網格（通常是橢圓形），Pixel2Mesh可以從間隙影像預測3D模型的尖端和麵。目標模型必須與模板網格同胚，因此使用橢圓形之類的凸形模板網格會在高度不凸的物件（例如椅子和燈具）上個月多個假物體。拓撲修改網路（TMN）通過另一個兩個新階段在Pixel2Mesh上進行迭代：變形修改階段（用於補償會增加模型重建誤差的錯誤面孔） ）和邊界優化階段。

同胚的經典例子

儘管變形和改進模板網格的常用方法效果很好，但它始於有關模型拓撲的主要假設。 3D模型的核心只是一個3D空間中的頂點集合，通過各個面進行分組和連線在一起。 是否可以避開中間表示並直接預測這些頂點和麵？

PolyGen

PolyGen 架構

PolyGen通過將3D模型表示為頂點和麵的嚴格有序序列，而不是影像，體素或點雲，對模型生成任務採取了一種非常獨特的方法。 這種嚴格的排序使他們能夠將基於注意力的序列建模方法應用於生成3D網格，就像BERT或GPT模型對文字所做的一樣。

PolyGen的總體目標是雙重的：首先為3D模型生成一組可能的頂點（可能由影像，體素或類標籤來限制），然後生成一系列的面，一個接一個地連線 頂點在一起，為該模型提供了一個合理的表面。 組合模型將網格p（M）上的分佈表示為兩個模型之間的聯合分佈：代表頂點的頂點模型p（V）和代表以頂點為條件的面的模型p（F | V）。

頂點模型是一種解碼器，它嘗試預測以先前標記為條件的序列中的下一個標記（並可選地以影像，體素欄位或類標籤為條件）。 表面模型由一個編碼器和一個解碼器指標網路組成，該網路表示頂點序列的分佈。 該指標網路一次有效地“選擇”一個頂點以新增到當前面序列並構建模型的面。 此模型均以先前的面序列和整個頂點序列為條件。 由於PolyGen架構非常複雜，並且依賴於各種概念，因此本文僅限於頂點模型。我將在後續文章中介紹表面模型。

預處理頂點

流行的ShapeNetCore資料集中的每個模型都可以表示為頂點和麵的集合。每個頂點都包含一個（x，y，z）座標，該座標描述了3D網格中的一個點。每個面都是指向組成該面角的頂點的索引列表。對於三角形面，此列表的長度為3個索引。對於n形面，此列表的長度是可變的。原始資料集非常大，因此為了節省時間，我在此處為您的實驗提供了資料集的一個更輕量級，經過預處理的子集。該子集僅包含來自5個形狀類別的模型，並且轉換為n形後的頂點少於800個（如下所述）。

為了使序列建模方法起作用，必須以受限的確定性方式表示資料，以消除儘可能多的可變性。因此，作者對資料集進行了許多簡化。首先，他們將所有輸入模型從三角形（連線3個頂點的面）轉換為n角（連線n個頂點的面），並使用Blender的平面抽取修改器合併面。由於大型網格並不總是具有唯一的三角剖分，因此可以更緊湊地表示相同的拓撲，並減少三角剖分中的歧義。為了篇幅所限，我不會在本文中介紹Blender指令碼，但是很多資源都很好地涵蓋了這一主題。我提供的資料集已被預先抽取。

在Blender的“平面”模式下應用“ Decimate”修改器前後，角度限制為1.0度的3D模型。

要繼續學習，請隨時下載此示例cube.obj（https://masonmcgough-data-bucket.s3-us-west-2.amazonaws.com/cube.obj）檔案。 此模型是具有8個頂點和6個面的基本立方體。 下面的簡單程式碼段從單個.obj檔案讀取所有頂點。

def load_obj(filename):
  """Load vertices from .obj wavefront format file."""
  vertices = []
  with open(filename, 'r') as mesh:
    for line in mesh:
      data = line.split()
      if len(data) > 0 and data[0] == 'v':
        vertices.append(data[1:])
  return np.array(vertices, dtype=np.float32)verts = load_obj(cube_path)
print('Cube Vertices')
print(verts)

其次，頂點從其z軸（在這種情況下為垂直軸）升序排列，然後是y軸，最後是x軸。 這樣，模型的頂點從下至上表示。 然後，在經典PolyGen模型中，將頂點連線成一維序列向量，對於較大的模型，該序列可以以非常長的序列向量結束。 作者在本文附錄E中描述了幾種減輕此負擔的修改方法。

要對一系列頂點進行排序，我們可以使用字典排序。 這與對字典中的單詞進行排序時所採用的方法相同。 要對兩個單詞進行排序，請檢視第一個字母，如果有平局，則檢視第二個字母，依此類推。 對於單詞“ aardvark”和“ apple”，第一個字母是“ a”和“ a”，因此我們移到第二個字母“ a”和“ p”，以告訴我們“ aardvark”在“ apple”之前。 在這種情況下，我們的“字母”依次是z，y和x座標。

verts_keys = [verts[..., i] for i in range(verts.shape[-1])]
sort_idxs = np.lexsort(verts_keys)
verts_sorted = verts[sort_idxs]

最後，將頂點座標標準化，然後進行量化，以將其轉換為離散的8位值。 該方法已在Pixel Recurrent Neural Networks和WaveNet中用於對音訊訊號進行建模，使它們能夠對頂點值施加分類分佈。 在WaveNet的原始論文中，作者指出：“分類分佈更靈活，可以更輕鬆地對任意分佈建模，因為它不對形狀進行任何假設。” 這種質量對於建模複雜的依存關係非常重要，例如3D模型中頂點之間的對稱性。

# normalize vertices to range [0.0, 1.0]
lims = [-1.0, 1.0]
norm_verts = (verts - lims[0]) / (lims[1] - lims[0])

# quantize vertices to integers in range [0, 255]
n_vals = 2 ** 8
delta = 1. / n_vals
quant_verts = np.maximum(np.minimum((norm_verts // delta), n_vals - 1), 0).astype(np.int32)

頂點模型

頂點模型由一個解碼器網路組成，該網路具有轉換器模型的所有標準特徵：輸入嵌入，18個轉換器解碼器層的堆疊，層歸一化以及最後在所有可能的序列標記上表示的softmax分佈。 給定長度N的扁平頂點序列Vseq，其目標是在給定模型引數的情況下最大化資料序列的對數似然性：

與LSTM不同，transformer 模型能夠以並行方式處理順序輸入，同時仍允許來自序列一部分的資訊為另一部分提供上下文。 這全都歸功於他們的注意力模組。 3D模型的頂點包含各種對稱性和遠點之間的複雜依存關係。 例如，考慮一個典型的表，其中模型相對角的邊是彼此的映象版本。 注意模組允許對這些型別的模式進行建模。

輸入嵌入

嵌入層是序列建模中用於將有限數量的標記轉換為特徵集的常用技術。在語言模型中，“國家”和“民族”一詞的含義可能非常相似，但與“蘋果”一詞卻相距甚遠。當單詞用唯一標記表示時，就沒有類似或不同的固有概念。嵌入層將這些標記轉換為向量表示，可以在其中模擬有意義的距離感。

PolyGen將相同的原理應用於頂點。該模型利用三種型別的嵌入層：座標（指示輸入令牌是x，y或z座標），值（指示令牌的值）以及位置（對頂點的順序進行編碼）。每個人都向模型傳達有關令牌的一條資訊。由於我們的頂點一次沿一個軸進給，因此座標嵌入為模型提供了基本的座標資訊，以使其知道給定值對應於哪種座標型別。

coord_tokens = np.concatenate(([0], np.arange(len(quant_verts)) % 3 + 1, (n_padding + 1) * [0]))

值嵌入對我們先前建立的量化頂點值進行編碼。 我們還需要一些序列控制點：額外的開始標記和停止標記，分別標記序列的開始和結束，以及填充標記，直到最大序列長度。

TOKENS = {
  '<pad>': 0,
  '<sos>': 1,
  '<eos>': 2
}max_verts = 12 # set low for prototyping
max_seq_len = 3 * max_verts + 2 

# num coords + start & stop tokens
n_tokens = len(TOKENS)
seq_len = len(quant_verts) + 2
n_padding = max_seq_len - seq_lenval_tokens = np.concatenate((
  [TOKENS['<sos>']],
  quant_verts + n_tokens,
  [TOKENS['<eos>']],
  n_padding * [TOKENS['<pad>']]
))

通過位置嵌入恢復由於並行化而丟失的給定序列位置n的位置資訊。 人們還可以使用位置編碼，這是一種無需學習的封閉形式的表示式。 在經典的transformer 論文“Attention Is All You Need”中，作者定義了一個位置編碼，該位置編碼由不同頻率的正弦和餘弦函式組成。 他們通過實驗確定了位置嵌入與位置編碼一樣好，但是編碼的優點是可以推斷出比訓練中遇到的序列更長的序列。

pos_tokens = np.arange(len(quant_tokens), dtype=np.int32)

生成所有這些令牌序列後，最後要做的就是建立一些嵌入層並將其組合。 每個嵌入層都需要知道期望的輸入字典的大小和要輸出的嵌入尺寸。 每層的嵌入維數為256，這意味著我們可以將它們與加法結合起來。 字典的大小取決於輸入可以具有的唯一值的數量。 對於值嵌入，它是量化值的數量加上控制令牌的數量。 對於座標嵌入，x，y和z的每個座標為1，以上都不為（控制標記）。 最後，對於每個可能的位置或最大序列長度，位置嵌入都需要一個。

n_embedding_channels = 256

# initialize value embedding layer
n_embeddings_value = 2 ** n_bits + n_tokens
value_embedding = torch.nn.Embedding(n_embeddings_value,
  n_embedding_channels, padding_idx=TOKENS['<pad>'])
  
# initialize coordinate embedding layer
n_embeddings_coord = 4
coord_embedding = torch.nn.Embedding(n_embeddings_coord,
  n_embedding_channels)
  
# initialize position embedding layer
n_embeddings_pos = max_seq_len
pos_embedding = torch.nn.Embedding(n_embeddings_pos,
  n_embedding_channels)

# pass through layers
value_embed = self.value_embedding(val_tokens)
coord_embed = self.coord_embedding(coord_tokens)
pos_embed = self.pos_embedding(pos_tokens)

# merge
x = value_embed + coord_embed + pos_embed

序列掩蔽

transformer 模型如此並行化的另一個結果是？ 對於在時間n的給定輸入令牌，模型實際上可以在序列的後面“看到”目標值，當您嘗試僅根據先前的序列值對模型進行條件調整時，這將成為一個問題。 為了防止模型使用無效的未來目標值，可以在自注意層的softmax步驟之前用-Inf遮蔽未來位置。

n_seq = len(val_tokens)
mask_dims = (n_seq, n_seq)
target_mask = torch.from_numpy(
  (val_tokens != TOKENS['<pad>'])[..., np.newaxis] \
  & (np.triu(np.ones(mask_dims), k=1).astype('uint8') == 0))

PolyGen還廣泛使用了無效的預測遮罩，以確保其生成的頂點和麵序列編碼有效的3D模型。 例如，必須執行諸如“ z座標不變小”和“只有在完整的頂點（z，y和x標記的三元組）之後才能出現停止標記”之類的規則，以防止模型產生無效的網格 。 這些約束僅在預測時強制執行，因為它們實際上會損害訓練效果。

核取樣

像許多序列預測模型一樣，該模型是自迴歸的，這意味著給定時間步長的輸出是下一時間步長可能值的分佈。整個序列一次被預測為一個令牌，模型在每個步驟中都會瀏覽先前時間步中的所有令牌，以選擇下一個令牌。解碼策略規定了如何從該分佈中選擇下一個令牌。

如果使用了次優的解碼策略，生成模型有時會陷入重複迴圈，或者產生質量差的序列。我們都看到過看起來像胡說八道的文字。 PolyGen採用一種稱為核取樣的解碼策略來生成高質量序列。原始論文在文字生成上下文中應用了此方法，但也可以將其應用於頂點。前提很簡單：僅從softmax分佈中共享top-p概率質量的標記中隨機抽取下一個標記。在推理時將其應用於生成網格，同時避免序列退化。有關核取樣的PyTorch實現，請參考此gist（https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317）

條件輸入

除了無條件生成模型外，PolyGen還支援使用類標籤，影像和體素進行輸入條件處理。 這些可以指導具有特定型別，外觀或形狀的網格的生成。 類標籤通過嵌入進行投影，然後在每個注意塊中的自注意層之後新增。 對於影像和體素，編碼器會建立一組嵌入，然後將其與transformer 解碼器進行交叉注意。

結論

PolyGen模型描述了用於有條件生成3D網格的強大，高效且靈活的框架。 序列生成可以在各種條件和輸入型別下完成，範圍從影像到體素到簡單的類標籤，甚至除了起始標記外什麼都不做。 表示網格網格頂點上的分佈的頂點模型只是關節分佈難題的一部分。

作者：Mason McGough

deephub翻譯組