Pytorch-如何有效地形成此变换张量?

时间:2019-02-14 18:32:51

标签: pytorch

我有一组旋转矩阵Rs:

Rs.shape = [62x3x3]

和一组翻译组件Js:

Js.shape = [62x3]

我一直在尝试找到一种有效的方法来将它们组合到[62x4x4]矩阵中,该矩阵是62个均匀的变换矩阵。目前,我正在使用愚蠢的for循环:

def make_A(R, t):
    R_homo = torch.cat([R, torch.zeros(1, 3).cuda()], dim = 0)
    t_homo = torch.cat([t.view(3,1), torch.ones(1, 1).cuda()], dim = 0)
    return torch.cat([R_homo, t_homo], dim=1)

transforms = self.NUM_JOINTS*[None]
for idj in range(0, self.NUM_JOINTS):
    transforms[idj] = make_A(Rs[idj, :], Js[idj,:])

FinalMatrix = torch.stack(transforms, dim=0)

这是非常低效的,大约需要10毫秒的时间才能形成。我该如何张紧呢?

1 个答案:

答案 0 :(得分:0)

不确定是否有助于提高效率,但这应该使代码矢量化:

def make_A(Rs, Js):
    R_homo = torch.cat((Rs, torch.zeros(Rs.shape[0], 1, 3)), dim=1)
    t_homo = torch.cat((Js, torch.ones(Js.shape[0], 1)), dim=1)
    return torch.cat((R_homo, t_homo.unsqueeze(2)), dim=2)