Skip to content

Instantly share code, notes, and snippets.

@chrischoy

chrischoy/compare_torch_mink.py

Created May 6, 2021 20:46
Show Gist options
  • Save chrischoy/d8e971daf8308aa1dcba1524bf1fd91a to your computer and use it in GitHub Desktop.
Save chrischoy/d8e971daf8308aa1dcba1524bf1fd91a to your computer and use it in GitHub Desktop.
Download ZIP
Raw
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")
@chrischoy
Copy link
Author

chrischoy commented May 6, 2021 •
edited
Loading

Tested on AMD Ryzen 7 3700 + Titan RTX

>>> torchsparse.__version__
'1.2.0'
>>> import MinkowskiEngine; MinkowskiEngine.__version__
'0.5.3'

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

@chrischoy
Copy link
Author 

(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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment