Hands-On Guide To Torch-Points3D: A Modular Deep Learning Framework For 3D Data
There has been a surge of advancements in automated analysis of 3D data caused by affordable LiDAR sensors, more efficient photogrammetry algorithms, and new neural network architectures. So much that the number of papers related to 3D data being presented at vision conferences is now on par with images, although this rapid methodological development is beneficial to the young field of deep learning for 3D, with its fast pace come several shortcomings: Adding new datasets, tasks, or neural architectures to existing approaches is a complicated endeavour, sometimes equivalent to reimplementing from scratch. Handling large 3D datasets requires a significant time investment and is prone to many implementation pitfalls. There is no standard approach for inference schemes and performance metrics, which makes assessing and reproducing new algorithms’ intrinsic performance difficult. Torch-Points3D aims to solve these issues. It is an open-source framework designed to facilitate deep neural networks on point cloud-based computer vision. It provides an intuitive interface with most open-access 3D datasets, implementations of many state-of-the-art networks, data augmentation schemes, and validated performance metrics. Torch-Points3D has a modular design and its components are highly customizable, they can be plugged into one another using a unified system of configuration files. It aims to make it easy to standardize experiments to ensure reproducibility and to help evaluate the performances of different approaches fairly. As the developers put it, “the purpose of our framework is to become for 3D point clouds what torchvision or PyTorch-geometric have become for images and graphs respectively“. The framework is built upon Pytorch Geometric and Facebook Hydra. Like PyTorch, Torch-Points3D uses the background processes to help increase the data processing speed. It off-loads the radius search and subsampling operations to background processes operating on CPUs. Training speed of the KPconv model, in thousandsof points processed per second (kpts/s) Functionalities/operations supported by Torch-Points3D You can check out all supported tasks and algorithms here. Supported datasets Torch-Points3D supports multiple 3D datasets with the data download, pre-processing, as well as automatic result submission. ScanNetS3DISModelNet10/40Shapenet3DMactchKITTI Odometry You can find a comprehensive list of all supported datasets here. Installation and Requirements Requirements CUDA 10 or higher (if you want GPU version)Python 3.7 or higher + headers (python-dev)PyTorch 1.7 or higherA Sparse convolution backend (optional) like torchsparse Run the following code before installing Torch-Points3D to ensure that you don’t run into a CUDA version mismatch error. import torch def format_pytorch_version(version): return version.split('+')[0] TORCH_version = torch.__version__ TORCH = format_pytorch_version(TORCH_version) def format_cuda_version(version): return 'cu' + version.replace('.', '') CUDA_version = torch.version.cuda CUDA = format_cuda_version(CUDA_version) !pip install torch-scatter    -f https://pytorch-geometric.com/whl/torch-${TORCH}+${CUDA}.html !pip install torch-sparse      -f https://pytorch-geometric.com/whl/torch-${TORCH}+${CUDA}.html !pip install torch-cluster     -f https://pytorch-geometric.com/whl/torch-${TORCH}+${CUDA}.html !pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-${TORCH}+${CUDA}.html !pip install torch-geometric Install Torch-Points3D from PyPI !pip install torch-points3d For instructions on how to install using other methods see this. Install PyVista for visualizing point clouds !pip install pyvista Creating a KP Conv Segmentation model with Torch-Points3D Import necessary libraries import os #omegaconf config is used for dealing with config files. from omegaconf import OmegaConf import pyvista as pv import torch import numpy as np We are going to use the Torch-Points3D version of ShapeNet. Create the config file for the dataset and download it using the torch_points3d.datasets.segmentation.ShapeNet class. CATEGORY = "All" USE_NORMALS = True shapenet_yaml = """ class: shapenet.ShapeNetDataset task: segmentation dataroot: %s normal: %r                                    # Use normal vectors as features first_subsampling: 0.02                       # Grid size of the input data pre_transforms:                               # Offline transforms, done only once - transform: NormalizeScale - transform: GridSampling3D params: size: ${first_subsampling} train_transforms:                             # Data augmentation pipeline - transform: RandomNoise params: sigma: 0.01 clip: 0.05 - transform: RandomScaleAnisotropic params: scales: [0.9,1.1] - transform: AddOnes - transform: AddFeatsByKeys params: list_add_to_x: [True] feat_names: ["ones"] delete_feats: [True] test_transforms: - transform: AddOnes - transform: AddFeatsByKeys params: list_add_to_x: [True] feat_names: ["ones"] delete_feats: [True] """ % (os.path.join(DIR,"data"), USE_NORMALS) params = OmegaConf.create(shapenet_yaml) if CATEGORY != "All": params.category = CATEGORY from torch_points3d.datasets.segmentation import ShapeNetDataset dataset = ShapeNetDataset(params) Visualize some random point clouds from the dataset using pyvista. objectid_1 = 9 objectid_2 = 82 objectid_3 = 95 samples = [objectid_1,objectid_2,objectid_3] p = pv.Plotter(notebook=True,shape=(1, len(samples)),window_size=[1024,412]) for i in range(len(samples)): p.subplot(0, i) sample = dataset.train_dataset[samples[i]] point_cloud = pv.PolyData(sample.pos.numpy()) point_cloud['y'] = sample.y.numpy() p.add_points(point_cloud,  show_scalar_bar=False, point_size=3) p.camera_position = [-1,5, -10] p.show() Create a multi-headed segmentation module to use with the KP Convolution network. from torch_points3d.core.common_modules import MLP, UnaryConv class MultiHeadClassifier(torch.nn.Module): """ Allows segregated segmentation in case the category of an object is known. This is the case in ShapeNet for example. Parameters ---------- in_features - size of the input channel cat_to_seg category to segment maps for example: { 'Airplane': [0,1,2], 'Table': [3,4] } """ def __init__(self, in_features, cat_to_seg, dropout_proba=0.5, bn_momentum=0.1): super().__init__() self._cat_to_seg = {} self._num_categories = len(cat_to_seg) self._max_seg_count = 0 self._max_seg = 0 self._shifts = torch.zeros((self._num_categories,), dtype=torch.long) for i, seg in enumerate(cat_to_seg.values()): self._max_seg_count = max(self._max_seg_count, len(seg)) self._max_seg = max(self._max_seg, max(seg)) self._shifts[i] = min(seg) self._cat_to_seg[i] = seg self.channel_rasing = MLP( [in_features, self._num_categories * in_features], bn_momentum=bn_momentum, bias=False ) if dropout_proba: self.channel_rasing.add_module("Dropout", torch.nn.Dropout(p=dropout_proba)) self.classifier = UnaryConv((self._num_categories, in_features, self._max_seg_count)) self._bias = torch.nn.Parameter(torch.zeros(self._max_seg_count,)) def forward(self, features, category_labels, **kwargs): assert features.dim() == 2 self._shifts = self._shifts.to(features.device) in_dim = features.shape[-1] features = self.channel_rasing(features) features = features.reshape((-1, self._num_categories, in_dim)) features = features.transpose(0, 1)  # [num_categories, num_points, in_dim] features = self.classifier(features) + self._bias  # [num_categories, num_points, max_seg] ind = category_labels.unsqueeze(-1).repeat(1, 1, features.shape[-1]).long() logits = features.gather(0, ind).squeeze(0) softmax = torch.nn.functional.log_softmax(logits, dim=-1) output = torch.zeros(logits.shape[0], self._max_seg + 1).to(features.device) cats_in_batch = torch.unique(category_labels) for cat in cats_in_batch: cat_mask = category_labels == cat seg_indices = self._cat_to_seg[cat.item()] probs = softmax[cat_mask, : len(seg_indices)] output[cat_mask, seg_indices[0] : seg_indices[-1] + 1] = probs return output Create a KPConv backbone model using the KPCONV method, you learn more about available models here. from torch_points3d.applications.kpconv import KPConv class PartSegKPConv(torch.nn.Module): def __init__(self, cat_to_seg): super().__init__() self.unet = KPConv( architecture="unet", input_nc=USE_NORMALS * 3, num_layers=4, in_grid_size=0.02 ) self.classifier = MultiHeadClassifier(self.unet.output_nc, cat_to_seg) @property def conv_type(self): """ This is needed by the dataset to infer which batch collate should be used""" return self.unet.conv_type def get_batch(self): return self.batch def get_output(self): """ This is needed by the tracker to get access to the ouputs of the network""" return self.output def get_labels(self): """ Needed by the tracker in order to access ground truth labels""" return self.labels def get_current_losses(self): """ Entry point for the tracker to grab the loss """ return {"loss_seg": float(self.loss_seg)} def forward(self, data): self.labels = data.y self.batch = data.batch # Forward through unet and classifier data_features = self.unet(data) self.output = self.classifier(data_features.x, data.category) # Set loss for the backward pass self.loss_seg = torch.nn.functional.nll_loss(self.output, self.labels) return self.output def get_spatial_ops(self): return self.unet.get_spatial_ops() def backward(self): self.loss_seg.backward() model = PartSegKPConv(dataset.class_to_segments) Create the data loaders and toggle the CPU operation precompute by setting the precompute_multi_scale parameter to True NUM_WORKERS = 4 BATCH_SIZE = 16 dataset.create_dataloaders( model, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS, shuffle=True, precompute_multi_scale=True ) sample = next(iter(dataset.train_dataloader)) sample.keys The sample contains the pre-computed spatial information in the multiscale (encoder side) and upsample (decoder) attributes. sample.multiscale contains 10 different versions of the input batch, each one of these versions contains the location of the points in pos as well as the neighbors of these points in the previous point cloud. Let’s take a look at the points coming out of each downsampling layer. sample_in_batch = 0 ms_data = sample.multiscale num_downsize = int(len(ms_data) / 2) p = pv.Plotter(notebook=True,shape=(1, num_downsize),window_size=[1024,256]) for i in range(0,num_downsize): p.subplot(0, i) pos = ms_data[2*i].pos[ms_data[2*i].batch == sample_in_batch].numpy() point_cloud = pv.PolyData(pos) point_cloud['y'] = pos[:,1] p.add_points(point_cloud,  show_scalar_bar=False, point_size=3) p.add_text("Layer {}".format(i+1),font_size=10) p.camera_position = [-1,5, -10] p.show() Train the model from tqdm.auto import tqdm import time class Trainer: def __init__(self,model, dataset, num_epoch = 50, device=torch.device('cuda')): self.num_epoch = num_epoch self._model = model self._dataset=dataset self.device = device def fit(self): self.optimizer = torch.optim.Adam(self._model.parameters(), lr=0.001) self.tracker = self._dataset.get_tracker(False, True) for i in range(self.num_epoch): print("=========== EPOCH %i ===========" % i) time.sleep(0.5) self.train_epoch() self.tracker.publish(i) self.test_epoch() self.tracker.publish(i) def train_epoch(self): self._model.to(self.device) self._model.train() self.tracker.reset("train") train_loader = self._dataset.train_dataloader iter_data_time = time.time() with tqdm(train_loader) as tq_train_loader: for i, data in enumerate(tq_train_loader): t_data = time.time() - iter_data_time iter_start_time = time.time() self.optimizer.zero_grad() data.to(self.device) self._model.forward(data) self._model.backward() self.optimizer.step() if i % 10 == 0: self.tracker.track(self._model) tq_train_loader.set_postfix( **self.tracker.get_metrics(), data_loading=float(t_data), iteration=float(time.time() - iter_start_time), ) iter_data_time = time.time() def test_epoch(self): self._model.to(self.device) self._model.eval() self.tracker.reset("test") test_loader = self._dataset.test_dataloaders[0] iter_data_time = time.time() with tqdm(test_loader) as tq_test_loader: for i, data in enumerate(tq_test_loader): t_data = time.time() - iter_data_time iter_start_time = time.time() data.to(self.device) self._model.forward(data) self.tracker.track(self._model) tq_test_loader.set_postfix( **self.tracker.get_metrics(), data_loading=float(t_data), iteration=float(time.time() - iter_start_time), ) iter_data_time = time.time() trainer = Trainer(model, dataset) trainer.fit() Last Epoch (Endnote) In this article, we discussed Torch-Points3D, a flexible and powerful framework that aims to make deep learning on 3D data both more accessible and reproducible. It’s built on Pytorch Geometric and Facebook Hydra. It has a modular design to facilitate easy experimentation and comes with many datasets and models built-in. As per the paper, the developers are currently working on a high-level API for pre-trained, self-supervised, self-trained, and unsupervised deep learning approaches operating on 3D point clouds. For the official code, documentation, papers, and tutorials, see: GitHubDocsBlogCollab NotebookExamplesTorch-Point 3D PaperKP Conv  Paper
Torch-Points3D is a flexible and powerful framework that aims to make deep learning on 3D data both more accessible and reproducible.
Torch-Points3D is a flexible and powerful framework that aims to make deep learning on 3D data both more accessible and reproducible.
["Deep Learning", "deep learning frameworks", "Guide", "point cloud data", "Python Libraries", "Visualize Spatial Data"]
["CUDA", "NumPy", "TPU", "deep learning frameworks", "Python Libraries", "AI", "ML", "neural network", "point cloud data", "PyTorch", "computer vision", "Visualize Spatial Data", "Aim", "deep learning", "Deep Learning", "Guide"]
["AI", "ML", "deep learning", "neural network", "computer vision", "Aim", "PyTorch", "NumPy", "TPU", "CUDA"]
A machine learning enthusiast with a knack for finding patterns. In my free time, I like to delve into the world of non-fiction books and video essays.