learning3d 0.0.1__py3-none-any.whl

learning3d/examples/train_rpmnet.py

@@ -0,0 +1,228 @@
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import RPMNet, PPFNet
+from learning3d.losses import FrobeniusNormLoss, RMSEFeaturesLoss
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def _init_(args):
+	if not os.path.exists('checkpoints'):
+		os.makedirs('checkpoints')
+	if not os.path.exists('checkpoints/' + args.exp_name):
+		os.makedirs('checkpoints/' + args.exp_name)
+	if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
+		os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
+	os.system('cp main.py checkpoints' + '/' + args.exp_name + '/' + 'main.py.backup')
+	os.system('cp model.py checkpoints' + '/' + args.exp_name + '/' + 'model.py.backup')
+
+
+class IOStream:
+	def __init__(self, path):
+		self.f = open(path, 'a')
+
+	def cprint(self, text):
+		print(text)
+		self.f.write(text + '\n')
+		self.f.flush()
+
+	def close(self):
+		self.f.close()
+
+def test_one_epoch(device, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		loss_val = FrobeniusNormLoss()(output['est_T'], igt) + RMSEFeaturesLoss()(output['r'])
+
+		test_loss += loss_val.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader, textio):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, test_loader)
+	textio.cprint('Validation Loss: %f & Validation Accuracy: %f'%(test_loss, test_accuracy))
+
+def train_one_epoch(device, model, train_loader, optimizer):
+	model.train()
+	train_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(train_loader)):
+		template, source, igt = data
+
+		template = template.to(device)
+		source = source.to(device)
+		igt = igt.to(device)
+
+		output = model(template, source)
+		loss_val = FrobeniusNormLoss()(output['est_T'], igt) + RMSEFeaturesLoss()(output['r'])
+		# print(loss_val.item())
+
+		# forward + backward + optimize
+		optimizer.zero_grad()
+		loss_val.backward()
+		optimizer.step()
+
+		train_loss += loss_val.item()
+		count += 1
+
+	train_loss = float(train_loss)/count
+	return train_loss
+
+def train(args, model, train_loader, test_loader, boardio, textio, checkpoint):
+	learnable_params = filter(lambda p: p.requires_grad, model.parameters())
+	if args.optimizer == 'Adam':
+		optimizer = torch.optim.Adam(learnable_params)
+	else:
+		optimizer = torch.optim.SGD(learnable_params, lr=0.1)
+
+	if checkpoint is not None:
+		min_loss = checkpoint['min_loss']
+		optimizer.load_state_dict(checkpoint['optimizer'])
+
+	best_test_loss = np.inf
+
+	for epoch in range(args.start_epoch, args.epochs):
+		train_loss = train_one_epoch(args.device, model, train_loader, optimizer)
+		test_loss = test_one_epoch(args.device, model, test_loader)
+
+		if test_loss<best_test_loss:
+			best_test_loss = test_loss
+			snap = {'epoch': epoch + 1,
+					'model': model.state_dict(),
+					'min_loss': best_test_loss,
+					'optimizer' : optimizer.state_dict(),}
+			torch.save(snap, 'checkpoints/%s/models/best_model_snap.t7' % (args.exp_name))
+			torch.save(model.state_dict(), 'checkpoints/%s/models/best_model.t7' % (args.exp_name))
+			torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/best_ptnet_model.t7' % (args.exp_name))
+
+		torch.save(snap, 'checkpoints/%s/models/model_snap.t7' % (args.exp_name))
+		torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % (args.exp_name))
+		torch.save(model.feature_model.state_dict(), 'checkpoints/%s/models/ptnet_model.t7' % (args.exp_name))
+		
+		boardio.add_scalar('Train Loss', train_loss, epoch+1)
+		boardio.add_scalar('Test Loss', test_loss, epoch+1)
+		boardio.add_scalar('Best Test Loss', best_test_loss, epoch+1)
+
+		textio.cprint('EPOCH:: %d, Traininig Loss: %f, Testing Loss: %f, Best Loss: %f'%(epoch+1, train_loss, test_loss, best_test_loss))
+
+def options():
+	parser = argparse.ArgumentParser(description='Point Cloud Registration')
+	parser.add_argument('--exp_name', type=str, default='exp_rpmnet', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--dataset_path', type=str, default='ModelNet40',
+						metavar='PATH', help='path to the input dataset') # like '/path/to/ModelNet40'
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+
+	# settings for input data
+	parser.add_argument('--dataset_type', default='modelnet', choices=['modelnet', 'shapenet2'],
+						metavar='DATASET', help='dataset type (default: modelnet)')
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+
+	# settings for PointNet
+	parser.add_argument('--fine_tune_pointnet', default='tune', type=str, choices=['fixed', 'tune'],
+						help='train pointnet (default: tune)')
+	parser.add_argument('--transfer_ptnet_weights', default='./checkpoints/exp_classifier/models/best_ptnet_model.t7', type=str,
+						metavar='PATH', help='path to pointnet features file')
+	parser.add_argument('--emb_dims', default=1024, type=int,
+						metavar='K', help='dim. of the feature vector (default: 1024)')
+	parser.add_argument('--symfn', default='max', choices=['max', 'avg'],
+						help='symmetric function (default: max)')
+
+	# settings for on training
+	parser.add_argument('--seed', type=int, default=1234)
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=10, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--epochs', default=200, type=int,
+						metavar='N', help='number of total epochs to run')
+	parser.add_argument('--start_epoch', default=0, type=int,
+						metavar='N', help='manual epoch number (useful on restarts)')
+	parser.add_argument('--optimizer', default='Adam', choices=['Adam', 'SGD'],
+						metavar='METHOD', help='name of an optimizer (default: Adam)')
+	parser.add_argument('--resume', default='', type=str,
+						metavar='PATH', help='path to latest checkpoint (default: null (no-use))')
+	parser.add_argument('--pretrained', default='', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+
+	torch.backends.cudnn.deterministic = True
+	torch.manual_seed(args.seed)
+	torch.cuda.manual_seed_all(args.seed)
+	np.random.seed(args.seed)
+
+	boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)
+	_init_(args)
+
+	textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
+	textio.cprint(str(args))
+
+	
+	trainset = RegistrationData('RPMNet', ModelNet40Data(train=True, num_points=args.num_points, use_normals=True), partial_source=True, partial_template=True)
+	testset = RegistrationData('RPMNet', ModelNet40Data(train=False, num_points=args.num_points, use_normals=True), partial_source=True, partial_template=True)
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=8, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	# Create RPMNet Model.
+	model = RPMNet(feature_model=PPFNet())
+	model = model.to(args.device)
+
+	checkpoint = None
+	if args.resume:
+		assert os.path.isfile(args.resume)
+		checkpoint = torch.load(args.resume)
+		args.start_epoch = checkpoint['epoch']
+		model.load_state_dict(checkpoint['model'])
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	if args.eval:
+		test(args, model, test_loader, textio)
+	else:
+		train(args, model, train_loader, test_loader, boardio, textio, checkpoint)
+
+if __name__ == '__main__':
+	main()

learning3d/losses/__init__.py

@@ -0,0 +1,12 @@
+from .rmse_features import RMSEFeaturesLoss
+from .frobenius_norm import FrobeniusNormLoss
+from .classification import ClassificationLoss
+from .correspondence_loss import CorrespondenceLoss
+try:
+	from .emd import EMDLoss
+except:
+	print("Sorry EMD loss is not compatible with your system!")
+try:
+	from .chamfer_distance import ChamferDistanceLoss
+except:
+	print("Sorry ChamferDistance loss is not compatible with your system!")

learning3d/losses/chamfer_distance.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def pairwise_distances(a: torch.Tensor, b: torch.Tensor, p=2):
+    """
+    Compute the pairwise distance_tensor matrix between a and b which both have size [m, n, d]. The result is a tensor of
+    size [m, n, n] whose entry [m, i, j] contains the distance_tensor between a[m, i, :] and b[m, j, :].
+    :param a: A tensor containing m batches of n points of dimension d. i.e. of size [m, n, d]
+    :param b: A tensor containing m batches of n points of dimension d. i.e. of size [m, n, d]
+    :param p: Norm to use for the distance_tensor
+    :return: A tensor containing the pairwise distance_tensor between each pair of inputs in a batch.
+    """
+
+    if len(a.shape) != 3:
+        raise ValueError("Invalid shape for a. Must be [m, n, d] but got", a.shape)
+    if len(b.shape) != 3:
+        raise ValueError("Invalid shape for a. Must be [m, n, d] but got", b.shape)
+    return (a.unsqueeze(2) - b.unsqueeze(1)).abs().pow(p).sum(3)
+
+def chamfer(a, b):
+    """
+    Compute the chamfer distance between two sets of vectors, a, and b
+    :param a: A m-sized minibatch of point sets in R^d. i.e. shape [m, n_a, d]
+    :param b: A m-sized minibatch of point sets in R^d. i.e. shape [m, n_b, d]
+    :return: A [m] shaped tensor storing the Chamfer distance between each minibatch entry
+    """
+    M = pairwise_distances(a, b)
+    dist1 = torch.mean(torch.sqrt(M.min(1)[0]))
+    dist2 = torch.mean(torch.sqrt(M.min(2)[0]))
+    return (dist1 + dist2) / 2.0
+
+
+def chamfer_distance(template: torch.Tensor, source: torch.Tensor):
+	try:
+		from .cuda.chamfer_distance import ChamferDistance
+		cost_p0_p1, cost_p1_p0 = ChamferDistance()(template, source)
+		cost_p0_p1 = torch.mean(torch.sqrt(cost_p0_p1))
+		cost_p1_p0 = torch.mean(torch.sqrt(cost_p1_p0))
+		chamfer_loss = (cost_p0_p1 + cost_p1_p0)/2.0
+	except:
+		chamfer_loss = chamfer(template, source)
+	return chamfer_loss
+
+
+class ChamferDistanceLoss(nn.Module):
+	def __init__(self):
+		super(ChamferDistanceLoss, self).__init__()
+
+	def forward(self, template, source):
+		return chamfer_distance(template, source)

learning3d/losses/classification.py

@@ -0,0 +1,14 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def classification_loss(prediction: torch.Tensor, target: torch.Tensor):
+	return F.nll_loss(prediction, target)
+
+
+class ClassificationLoss(nn.Module):
+	def __init__(self):
+		super(ClassificationLoss, self).__init__()
+
+	def forward(self, prediction, target):
+		return classification_loss(prediction, target)

learning3d/losses/correspondence_loss.py

@@ -0,0 +1,10 @@
+import torch
+
+class CorrespondenceLoss(torch.nn.Module):
+	def forward(self, template, source, corr_mat_pred, corr_mat):
+		# corr_mat:			batch_size x num_template x num_source (ground truth correspondence matrix)
+		# corr_mat_pred:	batch_size x num_source x num_template (predicted correspondence matrix)
+		batch_size, _, num_points_template = template.shape
+		_, _, num_points = source.shape
+		return torch.nn.functional.cross_entropy(corr_mat_pred.view(batch_size*num_points, num_points_template), 
+				torch.argmax(corr_mat.transpose(1,2).reshape(-1, num_points_template), axis=1))

learning3d/losses/cuda/chamfer_distance/__init__.py

@@ -0,0 +1 @@
+from .chamfer_distance import ChamferDistance

learning3d/losses/cuda/chamfer_distance/chamfer_distance.cpp

@@ -0,0 +1,185 @@
+#include <torch/torch.h>
+
+// CUDA forward declarations
+int ChamferDistanceKernelLauncher(
+    const int b, const int n,
+    const float* xyz,
+    const int m,
+    const float* xyz2,
+    float* result,
+    int* result_i,
+    float* result2,
+    int* result2_i);
+
+int ChamferDistanceGradKernelLauncher(
+    const int b, const int n,
+    const float* xyz1,
+    const int m,
+    const float* xyz2,
+    const float* grad_dist1,
+    const int* idx1,
+    const float* grad_dist2,
+    const int* idx2,
+    float* grad_xyz1,
+    float* grad_xyz2);
+
+
+void chamfer_distance_forward_cuda(
+    const at::Tensor xyz1, 
+    const at::Tensor xyz2, 
+    const at::Tensor dist1, 
+    const at::Tensor dist2, 
+    const at::Tensor idx1, 
+    const at::Tensor idx2) 
+{
+    ChamferDistanceKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data<float>(),
+                                            xyz2.size(1), xyz2.data<float>(),
+                                            dist1.data<float>(), idx1.data<int>(),
+                                            dist2.data<float>(), idx2.data<int>());
+}
+
+void chamfer_distance_backward_cuda(
+    const at::Tensor xyz1,
+    const at::Tensor xyz2, 
+    at::Tensor gradxyz1, 
+    at::Tensor gradxyz2, 
+    at::Tensor graddist1, 
+    at::Tensor graddist2, 
+    at::Tensor idx1, 
+    at::Tensor idx2)
+{
+    ChamferDistanceGradKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data<float>(),
+                                           xyz2.size(1), xyz2.data<float>(),
+                                           graddist1.data<float>(), idx1.data<int>(),
+                                           graddist2.data<float>(), idx2.data<int>(),
+                                           gradxyz1.data<float>(), gradxyz2.data<float>());
+}
+
+
+void nnsearch(
+    const int b, const int n, const int m,
+    const float* xyz1,
+    const float* xyz2,
+    float* dist,
+    int* idx)
+{
+    for (int i = 0; i < b; i++) {
+        for (int j = 0; j < n; j++) {
+            const float x1 = xyz1[(i*n+j)*3+0];
+            const float y1 = xyz1[(i*n+j)*3+1];
+            const float z1 = xyz1[(i*n+j)*3+2];
+            double best = 0;
+            int besti = 0;
+            for (int k = 0; k < m; k++) {
+                const float x2 = xyz2[(i*m+k)*3+0] - x1;
+                const float y2 = xyz2[(i*m+k)*3+1] - y1;
+                const float z2 = xyz2[(i*m+k)*3+2] - z1;
+                const double d=x2*x2+y2*y2+z2*z2;
+                if (k==0 || d < best){
+                    best = d;
+                    besti = k;
+                }
+            }
+            dist[i*n+j] = best;
+            idx[i*n+j] = besti;
+        }
+    }
+}
+
+
+void chamfer_distance_forward(
+    const at::Tensor xyz1, 
+    const at::Tensor xyz2, 
+    const at::Tensor dist1, 
+    const at::Tensor dist2, 
+    const at::Tensor idx1, 
+    const at::Tensor idx2) 
+{
+    const int batchsize = xyz1.size(0);
+    const int n = xyz1.size(1);
+    const int m = xyz2.size(1);
+
+    const float* xyz1_data = xyz1.data<float>();
+    const float* xyz2_data = xyz2.data<float>();
+    float* dist1_data = dist1.data<float>();
+    float* dist2_data = dist2.data<float>();
+    int* idx1_data = idx1.data<int>();
+    int* idx2_data = idx2.data<int>();
+
+    nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data);
+    nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data);
+}
+
+
+void chamfer_distance_backward(
+    const at::Tensor xyz1, 
+    const at::Tensor xyz2, 
+    at::Tensor gradxyz1, 
+    at::Tensor gradxyz2, 
+    at::Tensor graddist1, 
+    at::Tensor graddist2, 
+    at::Tensor idx1, 
+    at::Tensor idx2) 
+{
+    const int b = xyz1.size(0);
+    const int n = xyz1.size(1);
+    const int m = xyz2.size(1);
+
+    const float* xyz1_data = xyz1.data<float>();
+    const float* xyz2_data = xyz2.data<float>();
+    float* gradxyz1_data = gradxyz1.data<float>();
+    float* gradxyz2_data = gradxyz2.data<float>();
+    float* graddist1_data = graddist1.data<float>();
+    float* graddist2_data = graddist2.data<float>();
+    const int* idx1_data = idx1.data<int>();
+    const int* idx2_data = idx2.data<int>();
+
+    for (int i = 0; i < b*n*3; i++)
+        gradxyz1_data[i] = 0;
+    for (int i = 0; i < b*m*3; i++)
+        gradxyz2_data[i] = 0;
+    for (int i = 0;i < b; i++) {
+        for (int j = 0; j < n; j++) {
+            const float x1 = xyz1_data[(i*n+j)*3+0];
+            const float y1 = xyz1_data[(i*n+j)*3+1];
+            const float z1 = xyz1_data[(i*n+j)*3+2];
+            const int j2 = idx1_data[i*n+j];
+
+            const float x2 = xyz2_data[(i*m+j2)*3+0];
+            const float y2 = xyz2_data[(i*m+j2)*3+1];
+            const float z2 = xyz2_data[(i*m+j2)*3+2];
+            const float g = graddist1_data[i*n+j]*2;
+
+            gradxyz1_data[(i*n+j)*3+0] += g*(x1-x2);
+            gradxyz1_data[(i*n+j)*3+1] += g*(y1-y2);
+            gradxyz1_data[(i*n+j)*3+2] += g*(z1-z2);
+            gradxyz2_data[(i*m+j2)*3+0] -= (g*(x1-x2));
+            gradxyz2_data[(i*m+j2)*3+1] -= (g*(y1-y2));
+            gradxyz2_data[(i*m+j2)*3+2] -= (g*(z1-z2));
+        }
+        for (int j = 0; j < m; j++) {
+            const float x1 = xyz2_data[(i*m+j)*3+0];
+            const float y1 = xyz2_data[(i*m+j)*3+1];
+            const float z1 = xyz2_data[(i*m+j)*3+2];
+            const int j2 = idx2_data[i*m+j];
+            const float x2 = xyz1_data[(i*n+j2)*3+0];
+            const float y2 = xyz1_data[(i*n+j2)*3+1];
+            const float z2 = xyz1_data[(i*n+j2)*3+2];
+            const float g = graddist2_data[i*m+j]*2;
+            gradxyz2_data[(i*m+j)*3+0] += g*(x1-x2);
+            gradxyz2_data[(i*m+j)*3+1] += g*(y1-y2);
+            gradxyz2_data[(i*m+j)*3+2] += g*(z1-z2);
+            gradxyz1_data[(i*n+j2)*3+0] -= (g*(x1-x2));
+            gradxyz1_data[(i*n+j2)*3+1] -= (g*(y1-y2));
+            gradxyz1_data[(i*n+j2)*3+2] -= (g*(z1-z2));
+        }
+    }
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &chamfer_distance_forward, "ChamferDistance forward");
+    m.def("forward_cuda", &chamfer_distance_forward_cuda, "ChamferDistance forward (CUDA)");
+    m.def("backward", &chamfer_distance_backward, "ChamferDistance backward");
+    m.def("backward_cuda", &chamfer_distance_backward_cuda, "ChamferDistance backward (CUDA)");
+}