使用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翻译组