from scipy.spatial.transform import Rotation as R
class MNCriterion(nn.Module):
def __init__(self, t_loss_fn=nn.L1Loss(), q_loss_fn=nn.L1Loss(), sax=0.0, saq=0.0, srx=0, srq=0.0, learn_beta=False, learn_gamma=False):
super(MapNetCriterion, self).__init__()
self.t_loss_fn = t_loss_fn
self.q_loss_fn = q_loss_fn
self.sax = nn.Parameter(torch.Tensor([sax]), requires_grad=learn_beta)
self.saq = nn.Parameter(torch.Tensor([saq]), requires_grad=learn_beta)
self.srx = nn.Parameter(torch.Tensor([srx]), requires_grad=learn_gamma)
self.srq = nn.Parameter(torch.Tensor([srq]), requires_grad=learn_gamma)
self.register_buffer("conj", torch.tensor([1,-1,-1,-1]))
def pose_loss(self, pose_a, pose_b):
quat_a = pose_a[3:]
quat_b = pose_b[3:]
r = R.from_quat(quat_a)
q_mat = r.as_matrix().T
q_conj = torch.mul(quat_a, self.conj)
sub = pose_b[:3] - pose_a[:3]
delta_position = torch.matmul(q_mat, sub.T)
delta_quat = self.hamilton_product(q_conj, quat_b)
return np.concatenate((delta_position, delta_quat)).unsqueeze(0)
def forward(self, pred, targ):
p_loss = self.pose_loss(pred, targ)
......
I get an error
r = R.from_quat(quat_a))
TypeError: can’t convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
Since Scipy expects the numpy array, how do I deal with this?