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
egin{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的内存。