compare_torch_mink.py

import os
import time
import argparse
import numpy as np
from urllib.request import urlretrieve

try:
    import open3d as o3d
except ImportError:
    raise ImportError("Please install open3d with `pip install open3d`.")

import torch

import MinkowskiEngine as ME

import torchsparse.nn as spnn
import torchsparse.utils as sputil
from torchsparse.sparse_tensor import SparseTensor as TorchSparseTensor


parser = argparse.ArgumentParser()
parser.add_argument("--batch_size", type=int, default=10)
parser.add_argument("--out_channels", type=int, default=16)
parser.add_argument("--stride", type=int, default=1)
parser.add_argument("--kernel_size", type=int, default=3)
parser.add_argument("--mode", type=str, default='mink')


def get_conv(inc, outc, kernel_size, stride, mode="mink"):
    if mode == "mink":
        return ME.MinkowskiConvolution(
            inc, outc, kernel_size=kernel_size, stride=stride, dimension=3
        )
    else:
        return spnn.Conv3d(inc, outc, kernel_size=kernel_size, stride=stride)


def get_sparse_tensor(coords, feats, device="cuda", mode="mink", ic=3, oc=16, kernel_size=3):
    if mode == "mink":
        return ME.SparseTensor(
            feats,
            coords,
            device=device,
            minkowski_algorithm=ME.MinkowskiAlgorithm.SPEED_OPTIMIZED,
        )
    else:
        return TorchSparseTensor(feats.to(device), bcoords.to(device), stride=1)


# Check if the weights and file exist and download
if not os.path.isfile("1.ply"):
    print("Downloading an example pointcloud...")
    urlretrieve("https://bit.ly/3c2iLhg", "1.ply")


def load_file(file_name):
    pcd = o3d.io.read_point_cloud(file_name)
    coords = np.array(pcd.points)
    colors = np.array(pcd.colors)
    return coords, colors, pcd


if __name__ == "__main__":
    config = parser.parse_args()

    coords, colors, pcd = load_file("1.ply")
    batch_size = config.batch_size
    device = "cuda:0"
    mode = config.mode
    dcoords, dcolors = ME.utils.sparse_quantize(coords, colors, quantization_size=0.02)
    bcoords = ME.utils.batched_coordinates([dcoords for i in range(batch_size)]).to(
        device
    )
    bcolors = torch.cat(
        [torch.from_numpy(dcolors).float() for i in range(batch_size)], 0
    ).to(device)

    conv = get_conv(3, config.out_channels, kernel_size=config.kernel_size, stride=config.stride, mode=mode).to(device)

    ts = np.zeros(10)
    for i in range(10):
        c = time.time()
        stensor = get_sparse_tensor(bcoords, bcolors, mode=mode)
        output = conv(stensor)
        ts[i] = time.time() - c
    print(f"Forward Min time for {mode} {conv}: {np.min(ts)} for size {len(bcoords)} sparse tensor")

    loss = output.F.sum()
    for i in range(10):
        c = time.time()
        loss.backward()
        ts[i] = time.time() - c
    print(f"Backward Min time for {mode} {conv}: {np.min(ts)} for size {len(bcoords)} sparse tensor")

(py3.8-mink-dev) ➜  MinkowskiEngine git:(master) ✗ python compare_torch_sparse.py --mode mink --stride 2
Forward Min time for mink MinkowskiConvolution(in=3, out=16, kernel_size=[3, 3, 3], stride=[2, 2, 2], dilation=[1, 1, 1]): 0.01503896713256836 for size 1618900 sparse tensor
Backward Min time for mink MinkowskiConvolution(in=3, out=16, kernel_size=[3, 3, 3], stride=[2, 2, 2], dilation=[1, 1, 1]): 0.0053861141204833984 for size 1618900 sparse tensor
(py3.8-mink-dev) ➜  MinkowskiEngine git:(master) ✗ python compare_torch_sparse.py --mode sparse --stride 2
Forward Min time for sparse Conv3d(in_channels=3, out_channels=16, kernel_size=3, stride=2, dilation=1): 0.025519609451293945 for size 1618900 sparse tensor
Backward Min time for sparse Conv3d(in_channels=3, out_channels=16, kernel_size=3, stride=2, dilation=1): 0.0011670589447021484 for size 1618900 sparse tensor