MinkowskiEngine demo ModelNet40分类

MinkowskiEngine demo 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

训练ResNet进行ModelNet40分类

主要训练功能很简单。没有使用基于时间的训练，而是使用了基于迭代的训练。与基于时间的训练相比，基于迭代的训练的一个优势在于，训练逻辑独立于批处理大小。

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找到。

https://github.com/NVIDIA/MinkowskiEngine/blob/master/examples/modelnet40.py

警告

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