learning3d 0.0.1__py3-none-any.whl

learning3d/models/rpmnet.py

@@ -0,0 +1,359 @@
+import argparse
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .. utils import square_distance, angle_difference
+from .. ops.transform_functions import convert2transformation
+from .ppfnet import PPFNet
+_EPS = 1e-5  # To prevent division by zero
+
+
+class ParameterPredictionNet(nn.Module):
+	def __init__(self, weights_dim):
+		"""PointNet based Parameter prediction network
+
+		Args:
+			weights_dim: Number of weights to predict (excluding beta), should be something like
+						 [3], or [64, 3], for 3 types of features
+		"""
+
+		super().__init__()
+
+		self._logger = logging.getLogger(self.__class__.__name__)
+
+		self.weights_dim = weights_dim
+
+		# Pointnet
+		self.prepool = nn.Sequential(
+			nn.Conv1d(4, 64, 1),
+			nn.GroupNorm(8, 64),
+			nn.ReLU(),
+
+			nn.Conv1d(64, 64, 1),
+			nn.GroupNorm(8, 64),
+			nn.ReLU(),
+
+			nn.Conv1d(64, 64, 1),
+			nn.GroupNorm(8, 64),
+			nn.ReLU(),
+
+			nn.Conv1d(64, 128, 1),
+			nn.GroupNorm(8, 128),
+			nn.ReLU(),
+
+			nn.Conv1d(128, 1024, 1),
+			nn.GroupNorm(16, 1024),
+			nn.ReLU(),
+		)
+		self.pooling = nn.AdaptiveMaxPool1d(1)
+		self.postpool = nn.Sequential(
+			nn.Linear(1024, 512),
+			nn.GroupNorm(16, 512),
+			nn.ReLU(),
+
+			nn.Linear(512, 256),
+			nn.GroupNorm(16, 256),
+			nn.ReLU(),
+
+			nn.Linear(256, 2 + np.prod(weights_dim)),
+		)
+
+		self._logger.info('Predicting weights with dim {}.'.format(self.weights_dim))
+
+	def forward(self, x):
+		""" Returns alpha, beta, and gating_weights (if needed)
+
+		Args:
+			x: List containing two point clouds, x[0] = src (B, J, 3), x[1] = ref (B, K, 3)
+
+		Returns:
+			beta, alpha, weightings
+		"""
+
+		src_padded = F.pad(x[0], (0, 1), mode='constant', value=0)
+		ref_padded = F.pad(x[1], (0, 1), mode='constant', value=1)
+		concatenated = torch.cat([src_padded, ref_padded], dim=1)
+
+		prepool_feat = self.prepool(concatenated.permute(0, 2, 1))
+		pooled = torch.flatten(self.pooling(prepool_feat), start_dim=-2)
+		raw_weights = self.postpool(pooled)
+
+		beta = F.softplus(raw_weights[:, 0])
+		alpha = F.softplus(raw_weights[:, 1])
+
+		return beta, alpha
+
+
+
+def to_numpy(tensor):
+	"""Wrapper around .detach().cpu().numpy() """
+	if isinstance(tensor, torch.Tensor):
+		return tensor.detach().cpu().numpy()
+	elif isinstance(tensor, np.ndarray):
+		return tensor
+	else:
+		raise NotImplementedError
+
+
+def se3_transform(g, a, normals=None):
+	""" Applies the SE3 transform
+
+	Args:
+		g: SE3 transformation matrix of size ([1,] 3/4, 4) or (B, 3/4, 4)
+		a: Points to be transformed (N, 3) or (B, N, 3)
+		normals: (Optional). If provided, normals will be transformed
+
+	Returns:
+		transformed points of size (N, 3) or (B, N, 3)
+
+	"""
+	R = g[..., :3, :3]  # (B, 3, 3)
+	p = g[..., :3, 3]  # (B, 3)
+
+	if len(g.size()) == len(a.size()):
+		b = torch.matmul(a, R.transpose(-1, -2)) + p[..., None, :]
+	else:
+		raise NotImplementedError
+		b = R.matmul(a.unsqueeze(-1)).squeeze(-1) + p  # No batch. Not checked
+
+	if normals is not None:
+		rotated_normals = normals @ R.transpose(-1, -2)
+		return b, rotated_normals
+
+	else:
+		return b
+
+
+def match_features(feat_src, feat_ref, metric='l2'):
+	""" Compute pairwise distance between features
+
+	Args:
+		feat_src: (B, J, C)
+		feat_ref: (B, K, C)
+		metric: either 'angle' or 'l2' (squared euclidean)
+
+	Returns:
+		Matching matrix (B, J, K). i'th row describes how well the i'th point
+		 in the src agrees with every point in the ref.
+	"""
+	assert feat_src.shape[-1] == feat_ref.shape[-1]
+
+	if metric == 'l2':
+		dist_matrix = square_distance(feat_src, feat_ref)
+	elif metric == 'angle':
+		feat_src_norm = feat_src / (torch.norm(feat_src, dim=-1, keepdim=True) + _EPS)
+		feat_ref_norm = feat_ref / (torch.norm(feat_ref, dim=-1, keepdim=True) + _EPS)
+
+		dist_matrix = angle_difference(feat_src_norm, feat_ref_norm)
+	else:
+		raise NotImplementedError
+
+	return dist_matrix
+
+
+def sinkhorn(log_alpha, n_iters: int = 5, slack: bool = True, eps: float = -1) -> torch.Tensor:
+	""" Run sinkhorn iterations to generate a near doubly stochastic matrix, where each row or column sum to <=1
+
+	Args:
+		log_alpha: log of positive matrix to apply sinkhorn normalization (B, J, K)
+		n_iters (int): Number of normalization iterations
+		slack (bool): Whether to include slack row and column
+		eps: eps for early termination (Used only for handcrafted RPM). Set to negative to disable.
+
+	Returns:
+		log(perm_matrix): Doubly stochastic matrix (B, J, K)
+
+	Modified from original source taken from:
+		Learning Latent Permutations with Gumbel-Sinkhorn Networks
+		https://github.com/HeddaCohenIndelman/Learning-Gumbel-Sinkhorn-Permutations-w-Pytorch
+	"""
+
+	# Sinkhorn iterations
+	prev_alpha = None
+	if slack:
+		zero_pad = nn.ZeroPad2d((0, 1, 0, 1))
+		log_alpha_padded = zero_pad(log_alpha[:, None, :, :])
+
+		log_alpha_padded = torch.squeeze(log_alpha_padded, dim=1)
+
+		for i in range(n_iters):
+			# Row normalization
+			log_alpha_padded = torch.cat((
+					log_alpha_padded[:, :-1, :] - (torch.logsumexp(log_alpha_padded[:, :-1, :], dim=2, keepdim=True)),
+					log_alpha_padded[:, -1, None, :]),  # Don't normalize last row
+				dim=1)
+
+			# Column normalization
+			log_alpha_padded = torch.cat((
+					log_alpha_padded[:, :, :-1] - (torch.logsumexp(log_alpha_padded[:, :, :-1], dim=1, keepdim=True)),
+					log_alpha_padded[:, :, -1, None]),  # Don't normalize last column
+				dim=2)
+
+			if eps > 0:
+				if prev_alpha is not None:
+					abs_dev = torch.abs(torch.exp(log_alpha_padded[:, :-1, :-1]) - prev_alpha)
+					if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
+						break
+				prev_alpha = torch.exp(log_alpha_padded[:, :-1, :-1]).clone()
+
+		log_alpha = log_alpha_padded[:, :-1, :-1]
+	else:
+		for i in range(n_iters):
+			# Row normalization (i.e. each row sum to 1)
+			log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=2, keepdim=True))
+
+			# Column normalization (i.e. each column sum to 1)
+			log_alpha = log_alpha - (torch.logsumexp(log_alpha, dim=1, keepdim=True))
+
+			if eps > 0:
+				if prev_alpha is not None:
+					abs_dev = torch.abs(torch.exp(log_alpha) - prev_alpha)
+					if torch.max(torch.sum(abs_dev, dim=[1, 2])) < eps:
+						break
+				prev_alpha = torch.exp(log_alpha).clone()
+
+	return log_alpha
+
+
+def compute_rigid_transform(a: torch.Tensor, b: torch.Tensor, weights: torch.Tensor):
+	"""Compute rigid transforms between two point sets
+
+	Args:
+		a (torch.Tensor): (B, M, 3) points
+		b (torch.Tensor): (B, N, 3) points
+		weights (torch.Tensor): (B, M)
+
+	Returns:
+		Transform T (B, 3, 4) to get from a to b, i.e. T*a = b
+	"""
+
+	weights_normalized = weights[..., None] / (torch.sum(weights[..., None], dim=1, keepdim=True) + _EPS)
+	centroid_a = torch.sum(a * weights_normalized, dim=1)
+	centroid_b = torch.sum(b * weights_normalized, dim=1)
+	a_centered = a - centroid_a[:, None, :]
+	b_centered = b - centroid_b[:, None, :]
+	cov = a_centered.transpose(-2, -1) @ (b_centered * weights_normalized)
+
+	# Compute rotation using Kabsch algorithm. Will compute two copies with +/-V[:,:3]
+	# and choose based on determinant to avoid flips
+	u, s, v = torch.svd(cov, some=False, compute_uv=True)
+	rot_mat_pos = v @ u.transpose(-1, -2)
+	v_neg = v.clone()
+	v_neg[:, :, 2] *= -1
+	rot_mat_neg = v_neg @ u.transpose(-1, -2)
+	rot_mat = torch.where(torch.det(rot_mat_pos)[:, None, None] > 0, rot_mat_pos, rot_mat_neg)
+	assert torch.all(torch.det(rot_mat) > 0)
+
+	# Compute translation (uncenter centroid)
+	translation = -rot_mat @ centroid_a[:, :, None] + centroid_b[:, :, None]
+
+	transform = torch.cat((rot_mat, translation), dim=2)
+	return transform
+
+
+class RPMNet(nn.Module):
+	def __init__(self, feature_model=PPFNet()):
+		super().__init__()
+
+		self.add_slack = True
+		self.num_sk_iter = 5
+
+		self.weights_net = ParameterPredictionNet(weights_dim=[0])
+		self.feat_extractor = feature_model
+
+	def compute_affinity(self, beta, feat_distance, alpha=0.5):
+		"""Compute logarithm of Initial match matrix values, i.e. log(m_jk)"""
+		if isinstance(alpha, float):
+			hybrid_affinity = -beta[:, None, None] * (feat_distance - alpha)
+		else:
+			hybrid_affinity = -beta[:, None, None] * (feat_distance - alpha[:, None, None])
+		return hybrid_affinity
+
+	@staticmethod
+	def split_normals(data):
+		if data.shape[2] == 6:
+			xyz, normals = data[:, :, :3], data[:, :, 3:6]
+		elif data.shape[2] == 3:
+			xyz, normals = data, torch.zeros(data.shape).to(data.device)
+		return xyz, normals
+
+	def spam(self, xyz_template, norm_template, xyz_source, norm_source):
+		self.beta, self.alpha = self.weights_net([xyz_source, xyz_template])
+		self.feat_source = self.feat_extractor(xyz_source, norm_source)
+		self.feat_template = self.feat_extractor(xyz_template, norm_template)
+
+		feat_distance = match_features(self.feat_source, self.feat_template)
+		self.affinity = self.compute_affinity(self.beta, feat_distance, alpha=self.alpha)
+
+		# Compute weighted coordinates
+		log_perm_matrix = sinkhorn(self.affinity, n_iters=self.num_sk_iter, slack=self.add_slack)
+		self.perm_matrix = torch.exp(log_perm_matrix)
+		weighted_template = self.perm_matrix @ xyz_template / (torch.sum(self.perm_matrix, dim=2, keepdim=True) + _EPS)
+
+		return weighted_template
+
+	def forward(self, template, source, max_iterations: int = 1):
+		"""Forward pass for RPMNet
+
+		Args:
+			data: Dict containing the following fields:
+					'points_src': Source points (B, J, 6)
+					'points_ref': Reference points (B, K, 6)
+			num_iter (int): Number of iterations. Recommended to be 2 for training
+
+		Returns:
+			transform: Transform to apply to source points such that they align to reference
+			src_transformed: Transformed source points
+		"""
+
+		xyz_template, norm_template = self.split_normals(template)
+		xyz_source, norm_source = self.split_normals(source)
+
+		xyz_source_t, norm_source_t = xyz_source, norm_source
+
+		transforms = []
+		all_gamma, all_perm_matrices, all_weighted_template = [], [], []
+		all_beta, all_alpha = [], []
+
+		for i in range(max_iterations):
+			weighted_template = self.spam(xyz_template, norm_template, xyz_source_t, norm_source_t)			# Finding better correspondences after each iteration.
+			
+			# Compute transform and transform points
+			transform = compute_rigid_transform(xyz_source, weighted_template, weights=torch.sum(self.perm_matrix, dim=2))
+			xyz_source_t, norm_source_t = se3_transform(transform.detach(), xyz_source, norm_source)			# Apply transformation to original source.
+
+			transforms.append(transform)
+			all_gamma.append(torch.exp(self.affinity))
+			all_perm_matrices.append(self.perm_matrix)
+			all_weighted_template.append(weighted_template)
+			all_beta.append(to_numpy(self.beta))
+			all_alpha.append(to_numpy(self.alpha))
+
+		est_T = convert2transformation(transforms[max_iterations-1][:, :3, :3], transforms[max_iterations-1][:, :3, 3])
+		transformed_source = torch.bmm(est_T[:, :3, :3], source[:,:,:3].permute(0, 2, 1)).permute(0, 2, 1) + est_T[:, :3, 3].unsqueeze(1)
+
+		result = {'est_R': est_T[:, :3, :3],				# source -> template
+				  'est_t': est_T[:, :3,  3],				# source -> template
+				  'est_T': est_T,			# source -> template
+				  # 'r': self.feat_template - self.feat_source,
+				  'transformed_source': transformed_source}
+		
+		result['perm_matrices_init'] = all_gamma
+		result['perm_matrices'] = all_perm_matrices
+		result['weighted_template'] = all_weighted_template
+		result['beta'] = np.stack(all_beta, axis=0)
+		result['alpha'] = np.stack(all_alpha, axis=0)
+		result['transforms'] = transforms
+
+		return result
+
+
+if __name__ == '__main__':
+	template, source = torch.rand(10,1024,6), torch.rand(10,1024,6)
+	
+	net = RPMNet()
+	result = net(template, source)
+	import ipdb; ipdb.set_trace()

learning3d/models/segmentation.py

@@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Segmentation(nn.Module):
+	def __init__(self, feature_model, num_classes=40):
+		super(Segmentation, self).__init__()
+		self.feature_model = feature_model
+		self.num_classes = num_classes
+
+		self.conv1 = torch.nn.Conv1d(self.feature_model.emb_dims+64, 512, 1)
+		self.conv2 = torch.nn.Conv1d(512, 256, 1)
+		self.conv3 = torch.nn.Conv1d(256, 128, 1)
+		self.conv4 = torch.nn.Conv1d(128, self.num_classes, 1)
+		self.bn1 = nn.BatchNorm1d(512)
+		self.bn2 = nn.BatchNorm1d(256)
+		self.bn3 = nn.BatchNorm1d(128)
+
+	def forward(self, input_data):
+		output = self.feature_model(input_data)
+		output = F.relu(self.bn1(self.conv1(output)))
+		output = F.relu(self.bn2(self.conv2(output)))
+		output = F.relu(self.bn3(self.conv3(output)))
+		output = self.conv4(output)
+		output = output.permute(0, 2, 1)				# B x N x num_classes
+		return output
+
+if __name__ == '__main__':
+	from pointnet import PointNet
+	x = torch.rand(10,1024,3)
+
+	pn = PointNet(global_feat=False)
+	seg = Segmentation(pn)
+	seg_result = seg(x)
+	
+	print('Input Shape: {}\n Segmentation Output Shape: {}'
+		  .format(x.shape, seg_result.shape))

learning3d/ops/data_utils.py

@@ -0,0 +1,45 @@
+import torch
+
+def mean_shift(template, source, p0_zero_mean, p1_zero_mean):
+	template_mean = torch.eye(3).view(1, 3, 3).expand(template.size(0), 3, 3).to(template) 		# [B, 3, 3]
+	source_mean = torch.eye(3).view(1, 3, 3).expand(source.size(0), 3, 3).to(source) 			# [B, 3, 3]
+	
+	if p0_zero_mean:
+		p0_m = template.mean(dim=1) # [B, N, 3] -> [B, 3]
+		template_mean = torch.cat([template_mean, p0_m.unsqueeze(-1)], dim=2)
+		one_ = torch.tensor([[[0.0, 0.0, 0.0, 1.0]]]).repeat(template_mean.shape[0], 1, 1).to(template_mean)    # (Bx1x4)
+		template_mean = torch.cat([template_mean, one_], dim=1)
+		template = template - p0_m.unsqueeze(1)
+	# else:
+		# q0 = template
+
+	if p1_zero_mean:
+		#print(numpy.any(numpy.isnan(p1.numpy())))
+		p1_m = source.mean(dim=1) # [B, N, 3] -> [B, 3]
+		source_mean = torch.cat([source_mean, -p0_m.unsqueeze(-1)], dim=2)
+		one_ = torch.tensor([[[0.0, 0.0, 0.0, 1.0]]]).repeat(source_mean.shape[0], 1, 1).to(source_mean)    # (Bx1x4)
+		source_mean = torch.cat([source_mean, one_], dim=1)
+		source = source - p1_m.unsqueeze(1)
+	# else:
+		# q1 = source
+	return template, source, template_mean, source_mean
+
+def postprocess_data(result, p0, p1, a0, a1, p0_zero_mean, p1_zero_mean):
+	#output' = trans(p0_m) * output * trans(-p1_m)
+	#        = [I, p0_m;] * [R, t;] * [I, -p1_m;]
+	#          [0, 1    ]   [0, 1 ]   [0,  1    ]
+	est_g = result['est_T']
+	if p0_zero_mean:
+		est_g = a0.to(est_g).bmm(est_g)
+	if p1_zero_mean:
+		est_g = est_g.bmm(a1.to(est_g))
+	result['est_T'] = est_g
+
+	est_gs = result['est_T_series'] # [M, B, 4, 4]
+	if p0_zero_mean:
+		est_gs = a0.unsqueeze(0).contiguous().to(est_gs).matmul(est_gs)
+	if p1_zero_mean:
+		est_gs = est_gs.matmul(a1.unsqueeze(0).contiguous().to(est_gs))
+	result['est_T_series'] = est_gs
+
+	return result

learning3d/ops/invmat.py

@@ -0,0 +1,134 @@
+""" inverse matrix """
+
+import torch
+
+
+def batch_inverse(x):
+    """ M(n) -> M(n); x -> x^-1 """
+    batch_size, h, w = x.size()
+    assert h == w
+    y = torch.zeros_like(x)
+    for i in range(batch_size):
+        y[i, :, :] = x[i, :, :].inverse()
+    return y
+
+def batch_inverse_dx(y):
+    """ backward """
+    # Let y(x) = x^-1.
+    # compute dy
+    #   dy = dy(j,k)
+    #      = - y(j,m) * dx(m,n) * y(n,k)
+    #      = - y(j,m) * y(n,k) * dx(m,n)
+    # therefore,
+    #   dy(j,k)/dx(m,n) = - y(j,m) * y(n,k)
+    batch_size, h, w = y.size()
+    assert h == w
+    # compute dy(j,k,m,n) = dy(j,k)/dx(m,n) = - y(j,m) * y(n,k)
+    #   = - (y(j,:))' * y'(k,:)
+    yl = y.repeat(1, 1, h).view(batch_size*h*h, h, 1)
+    yr = y.transpose(1, 2).repeat(1, h, 1).view(batch_size*h*h, 1, h)
+    dy = - yl.bmm(yr).view(batch_size, h, h, h, h)
+
+    # compute dy(m,n,j,k) = dy(j,k)/dx(m,n) = - y(j,m) * y(n,k)
+    #   = - (y'(m,:))' * y(n,:)
+    #yl = y.transpose(1, 2).repeat(1, 1, h).view(batch_size*h*h, h, 1)
+    #yr = y.repeat(1, h, 1).view(batch_size*h*h, 1, h)
+    #dy = - yl.bmm(yr).view(batch_size, h, h, h, h)
+
+    return dy
+
+
+def batch_pinv_dx(x):
+    """ returns y = (x'*x)^-1 * x' and dy/dx. """
+    # y = (x'*x)^-1 * x'
+    #   = s^-1 * x'
+    #   = b * x'
+    # d{y(j,k)}/d{x(m,n)}
+    #   = d{b(j,i) * x(k,i)}/d{x(m,n)}
+    #   = d{b(j,i)}/d{x(m,n)} * x(k,i) + b(j,i) * d{x(k,i)}/d{x(m,n)}
+    # d{b(j,i)}/d{x(m,n)}
+    #   = d{b(j,i)}/d{s(p,q)} * d{s(p,q)}/d{x(m,n)}
+    #   = -b(j,p)*b(q,i) * d{s(p,q)}/d{x(m,n)}
+    # d{s(p,q)}/d{x(m,n)}
+    #   = d{x(t,p)*x(t,q)}/d{x(m,n)}
+    #   = d{x(t,p)}/d{x(m,n)} * x(t,q) + x(t,p) * d{x(t,q)}/d{x(m,n)}
+    batch_size, h, w = x.size()
+    xt = x.transpose(1, 2)
+    s = xt.bmm(x)
+    b = batch_inverse(s)
+    y = b.bmm(xt)
+
+    # dx/dx
+    ex = torch.eye(h*w).to(x).unsqueeze(0).view(1, h, w, h, w)
+    # ds/dx = dx(t,_)/dx * x(t,_) + x(t,_) * dx(t,_)/dx
+    ex1 = ex.view(1, h, w*h*w) # [t, p*m*n]
+    dx1 = x.transpose(1, 2).matmul(ex1).view(batch_size, w, w, h, w) # [q, p,m,n]
+    ds_dx = dx1.transpose(1, 2) + dx1 # [p, q, m, n]
+    # db/ds
+    db_ds = batch_inverse_dx(b) # [j, i, p, q]
+    # db/dx = db/d{s(p,q)} * d{s(p,q)}/dx
+    db1 = db_ds.view(batch_size, w*w, w*w).bmm(ds_dx.view(batch_size, w*w, h*w))
+    db_dx = db1.view(batch_size, w, w, h, w) # [j, i, m, n]
+    # dy/dx = db(_,i)/dx * x(_,i) + b(_,i) * dx(_,i)/dx
+    dy1 = db_dx.transpose(1, 2).contiguous().view(batch_size, w, w*h*w)
+    dy1 = x.matmul(dy1).view(batch_size, h, w, h, w) # [k, j, m, n]
+    ext = ex.transpose(1, 2).contiguous().view(1, w, h*h*w)
+    dy2 = b.matmul(ext).view(batch_size, w, h, h, w) # [j, k, m, n]
+    dy_dx = dy1.transpose(1, 2) + dy2
+
+    return y, dy_dx
+
+
+class InvMatrix(torch.autograd.Function):
+    """ M(n) -> M(n); x -> x^-1.
+    """
+    @staticmethod
+    def forward(ctx, x):
+        y = batch_inverse(x)
+        ctx.save_for_backward(y)
+        return y
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        y, = ctx.saved_tensors   # v0.4
+        #y, = ctx.saved_variables  # v0.3.1
+        batch_size, h, w = y.size()
+        assert h == w
+
+        # Let y(x) = x^-1 and assume any function f(y(x)).
+        # compute df/dx(m,n)...
+        #   df/dx(m,n) = df/dy(j,k) * dy(j,k)/dx(m,n)
+        # well, df/dy is 'grad_output'
+        # and so we will return 'grad_input = df/dy(j,k) * dy(j,k)/dx(m,n)'
+
+        dy = batch_inverse_dx(y)  # dy(j,k,m,n) = dy(j,k)/dx(m,n)
+        go = grad_output.contiguous().view(batch_size, 1, h*h)  # [1, (j*k)]
+        ym = dy.view(batch_size, h*h, h*h)  # [(j*k), (m*n)]
+        r = go.bmm(ym)  # [1, (m*n)]
+        grad_input = r.view(batch_size, h, h)  # [m, n]
+
+        return grad_input
+
+
+
+if __name__ == '__main__':
+    def test():
+        x = torch.randn(2, 3, 2)
+        x_val = x.requires_grad_()
+
+        s_val = x_val.transpose(1, 2).bmm(x_val)
+        s_inv = InvMatrix.apply(s_val)
+        y_val = s_inv.bmm(x_val.transpose(1, 2))
+        y_val.sum().backward()
+        t1 = x_val.grad
+
+        y, dy_dx = batch_pinv_dx(x)
+        t2 = dy_dx.sum(1).sum(1)
+
+        print(t1)
+        print(t2)
+        print(t1 - t2)
+
+    test()
+
+#EOF