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 cloudReference :[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] -= 1n_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_idxacc += _n_sampler = 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:] * Creturn 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 - minself.bias = mindef __call__(self, coords, feats):s = self.scale * np.random.rand(1) + self.biasreturn 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() - soptimizer.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的內存。
總結
以上是生活随笔為你收集整理的MinkowskiEngine demo ModelNet40分类的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: MinkowskiEngine基准测试
- 下一篇: MinkowskiEngine语义分割