chrischoy · May 6, 2021 20:46 · chrischoy · May 6, 2021 · chrischoy · May 6, 2021
diff --git a/compare_torch_mink.py b/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")
	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")