如何在python中加速切片，而不使用for循环

Question

我正在尝试加速以下python代码：

import torch
import numpy as np

A = torch.zeros(11, 16, 64)
B = torch.randn(11, 9, 64)

indices = np.random.randint(0,9,(11,16))

for i in range(len(A)):
    A[i,:,:] = B[i,indices[i],:]

有没有不使用for循环的好方法？这样，它确实很慢，尤其是在处理大数据时。索引是大小为（11,16）的预定义2维矩阵。我需要的是根据索引的顺序将B的元素分配给A。加速之后，A的结果应该与我的结果A完全相同。谢谢！

Answer 1

您可以使用多个多维索引，但是它们的大小必须相同或可以广播。因此，例如

# create a (11, 1) range array that broadcasts with indices which is (11, 16)
indices0 = np.expand_dims(np.arange(indices.shape[0]), 1)
A = B[indices0, indices, :]

或者如果indices是torch.LongTensor

indices0 = torch.arange(indices.shape[0]).unsqueeze(1)
A = B[indices0, indices, :]

Answer 2

即使对于机器学习项目，使用numpy进行切片也足够快。如果要在这种情况下使代码运行更快，则应使用以下方法：

A_length = len(A)
i = 0
while i < A_length:
    A[i,:,:] = B[i,indices[i],:]
    i += 1

range对象使用__iter__和__next__方法来生成迭代索引（在大多数情况下），即使它是用C编写的，它也比仅声明索引计数器要慢，并且每轮都要添加一个步骤。

但是for循环对于您的代码而言更具可读性和简单性，加上使用while循环不会大大提高速度。我认为您不应该使用while循环来稍微提高速度。

但是...

如果您希望您的代码尽可能快地运行：

1. 了解一些性能技巧
2. 考虑使用ssh和GPU远程服务器。它们比您的CPU快得多（如果您使用的是CPU计算机）
3. 学习C或JS，它们是编译语言。 C比未经优化的python快200倍，而使用多核处理和多线程（来自here）的C快大约50000倍]