ModelNet40 分类¶

在此页面中，我们将介绍一个简单的演示示例，该示例使用 3D 卷积神经网络进行分类训练。输入是一个稀疏张量，卷积也定义在稀疏张量上。该网络是以下架构的扩展，但具有残差块和更多层。

制作 ModelNet40 数据加载器¶

首先，我们需要创建一个数据加载器，该加载器返回网格的稀疏张量表示。但是，如果我们仅使用顶点，则 3D 模型的网格表示可以是稀疏的。

我们将首先以相同的密度采样点。

def resample_mesh(mesh_cad, density=1):
    '''
    https://chrischoy.github.io/research/barycentric-coordinate-for-mesh-sampling/
    Samples point cloud on the surface of the model defined as vectices and
    faces. This function uses vectorized operations so fast at the cost of some
    memory.

    param mesh_cad: low-polygon triangle mesh in o3d.geometry.TriangleMesh
    param density: density of the point cloud per unit area
    param return_numpy: return numpy format or open3d pointcloud format
    return resampled point cloud

    Reference :
      [1] Barycentric coordinate system
      \begin{align}
        P = (1 - \sqrt{r_1})A + \sqrt{r_1} (1 - r_2) B + \sqrt{r_1} r_2 C
      \end{align}
    '''
    faces = np.array(mesh_cad.triangles).astype(int)
    vertices = np.array(mesh_cad.vertices)

    vec_cross = np.cross(vertices[faces[:, 0], :] - vertices[faces[:, 2], :],
                         vertices[faces[:, 1], :] - vertices[faces[:, 2], :])
    face_areas = np.sqrt(np.sum(vec_cross**2, 1))

    n_samples = (np.sum(face_areas) * density).astype(int)

    n_samples_per_face = np.ceil(density * face_areas).astype(int)
    floor_num = np.sum(n_samples_per_face) - n_samples
    if floor_num > 0:
        indices = np.where(n_samples_per_face > 0)[0]
        floor_indices = np.random.choice(indices, floor_num, replace=True)
        n_samples_per_face[floor_indices] -= 1

    n_samples = np.sum(n_samples_per_face)

    # Create a vector that contains the face indices
    sample_face_idx = np.zeros((n_samples,), dtype=int)
    acc = 0
    for face_idx, _n_sample in enumerate(n_samples_per_face):
        sample_face_idx[acc:acc + _n_sample] = face_idx
        acc += _n_sample

    r = np.random.rand(n_samples, 2)
    A = vertices[faces[sample_face_idx, 0], :]
    B = vertices[faces[sample_face_idx, 1], :]
    C = vertices[faces[sample_face_idx, 2], :]

    P = (1 - np.sqrt(r[:, 0:1])) * A + \
        np.sqrt(r[:, 0:1]) * (1 - r[:, 1:]) * B + \
        np.sqrt(r[:, 0:1]) * r[:, 1:] * C

    return P

上述函数将以相同的密度对网格上的点进行采样。接下来，我们在量化步骤之前执行一系列数据增强步骤。

数据增强¶

稀疏张量由两个组成部分构成：1）坐标，以及 2）与这些坐标关联的特征。我们必须将数据增强应用于这两个组成部分，以最大程度地利用固定数据集，并使网络对噪声具有鲁棒性。

这在图像数据增强中并不是什么新鲜事。我们将随机平移、剪切、缩放应用于图像，所有这些都是坐标数据增强。颜色失真（例如色度平移、颜色通道上的高斯噪声、色相饱和度增强）都是特征数据增强。

但是，由于我们在 ModelNet40 数据集中仅将坐标作为数据，因此我们将仅应用坐标数据增强。

class RandomRotation:

    def _M(self, axis, theta):
        return expm(np.cross(np.eye(3), axis / norm(axis) * theta))

    def __call__(self, coords, feats):
        R = self._M(
            np.random.rand(3) - 0.5, 2 * np.pi * (np.random.rand(1) - 0.5))
        return coords @ R, feats


class RandomScale:

    def __init__(self, min, max):
        self.scale = max - min
        self.bias = min

    def __call__(self, coords, feats):
        s = self.scale * np.random.rand(1) + self.bias
        return coords * s, feats


class RandomShear:

    def __call__(self, coords, feats):
        T = np.eye(3) + np.random.randn(3, 3)
        return coords @ T, feats


class RandomTranslation:

    def __call__(self, coords, feats):
        trans = 0.05 * np.random.randn(1, 3)
        return coords + trans, feats

训练用于 ModelNet40 分类的 ResNet¶

主要的训练功能很简单。但是，我使用基于迭代的训练而不是基于 epoch 的训练。基于迭代的训练优于基于 epoch 的训练的一个优点是，训练逻辑与批处理大小无关。

def train(net, device, config):
    optimizer = optim.SGD(
        net.parameters(),
        lr=config.lr,
        momentum=config.momentum,
        weight_decay=config.weight_decay)
    scheduler = optim.lr_scheduler.ExponentialLR(optimizer, 0.95)

    crit = torch.nn.CrossEntropyLoss()

   ...

    net.train()
    train_iter = iter(train_dataloader)
    val_iter = iter(val_dataloader)
    logging.info(f'LR: {scheduler.get_lr()}')
    for i in range(curr_iter, config.max_iter):

        s = time()
        data_dict = train_iter.next()
        d = time() - s

        optimizer.zero_grad()
        sin = ME.SparseTensor(data_dict['feats'],
                              data_dict['coords'].int()).to(device)
        sout = net(sin)
        loss = crit(sout.F, data_dict['labels'].to(device))
        loss.backward()
        optimizer.step()
        t = time() - s

        ...

运行示例¶

组装所有代码块后，您可以运行自己的 ModelNet40 分类网络。

python -m examples.modelnet40 --batch_size 128 --stat_freq 100

完整的代码可以在example/modelnet40.py中找到。

警告

ModelNet40 数据加载和体素化是训练中最耗时的部分。因此，该示例将所有 ModelNet40 数据缓存到内存中，这会占用大约 10G 的内存。