广播旋转矩阵乘法

Question

如何以更直接的方式完成用# <----标记的行？

在程序中，x的每一行都是一个点的坐标，rot_mat[0]和rot_mat[1]是两个旋转矩阵。该程序通过每个旋转矩阵对x进行旋转。

如果更改每个旋转矩阵和坐标之间的乘法顺序能使事情更简单，则可以进行更改。我希望每行x或代表一个点的结果。

结果应该与检查结果匹配。

程序如下：

# Rotation of coordinates of 4 points by 
# each of the 2 rotation matrices.
import numpy as np
from scipy.stats import special_ortho_group
rot_mats = special_ortho_group.rvs(dim=3, size=2)  # 2 x 3 x 3
x = np.arange(12).reshape(4, 3)
result = np.dot(rot_mats, x.T).transpose((0, 2, 1))  # <----
print("---- result ----")
print(result)
print("---- check ----")
print(np.dot(x, rot_mats[0].T))
print(np.dot(x, rot_mats[1].T))

结果：

---- result ----
[[[  0.20382264   1.15744672   1.90230739]
  [ -2.68064533   3.71537598   5.38610452]
  [ -5.56511329   6.27330525   8.86990165]
  [ -8.44958126   8.83123451  12.35369878]]

 [[  1.86544623   0.53905202  -1.10884323]
  [  5.59236544  -1.62845022  -4.00918928]
  [  9.31928465  -3.79595246  -6.90953533]
  [ 13.04620386  -5.9634547   -9.80988139]]]
---- check ----
[[  0.20382264   1.15744672   1.90230739]
 [ -2.68064533   3.71537598   5.38610452]
 [ -5.56511329   6.27330525   8.86990165]
 [ -8.44958126   8.83123451  12.35369878]]
[[  1.86544623   0.53905202  -1.10884323]
 [  5.59236544  -1.62845022  -4.00918928]
 [  9.31928465  -3.79595246  -6.90953533]
 [ 13.04620386  -5.9634547   -9.80988139]]

- R zu

---

现在出了什么问题？ - cs95

我希望能够摆脱 transpose((0, 2, 1))。 - R zu

那个转置不是瓶颈，但 np.dot 是。 - Divakar

Answer 1

使用np.tensordot进行涉及这些张量的乘法 -

np.tensordot(rot_mats, x, axes=((2),(1))).swapaxes(1,2)

这里有一些时间上的数据，可以让我们相信为什么tensordot与tensors更搭配 -

In [163]: rot_mats = np.random.rand(20,30,30)
     ...: x = np.random.rand(40,30)

# With numpy.dot
In [164]: %timeit np.dot(rot_mats, x.T).transpose((0, 2, 1))
1000 loops, best of 3: 670 µs per loop

# With numpy.tensordot
In [165]: %timeit np.tensordot(rot_mats, x, axes=((2),(1))).swapaxes(1,2)
10000 loops, best of 3: 75.7 µs per loop

In [166]: rot_mats = np.random.rand(200,300,300)
     ...: x = np.random.rand(400,300)

# With numpy.dot
In [167]: %timeit np.dot(rot_mats, x.T).transpose((0, 2, 1))
1 loop, best of 3: 1.82 s per loop

# With numpy.tensordot
In [168]: %timeit np.tensordot(rot_mats, x, axes=((2),(1))).swapaxes(1,2)
10 loops, best of 3: 185 ms per loop

- Divakar

---

swapaxes会返回一个新的数组而不是视图吗？ - R zu

不要认为它可以是一个视图，因为轴的顺序从0、1、2变成了0、2、1。一个简单的测试是在swapaxes之前和之后对所有内容进行展平，并查看它们是否给出相同的结果。 - R zu

@Rzu 也许你在想其他的观点。我在谈论的是是否为副本。更多信息请参考 - http://scipy-cookbook.readthedocs.io/items/ViewsVsCopies.html - Divakar

@Rzu 我很确定是 ravel 使得产生了一个副本。你可以使用 np.shares_memory 进行检查。 - Paul Panzer

有点惊人的是，tensordot 比 dot 更快，swapaxes 比 transpose 更快，或者两者都更快。 - R zu