python中循环的矢量化

我有一个下面的循环，其中我正在计算不同大小批次的

softmax

转换，如下所示

import numpy as np 
def softmax(Z,arr):
    """
    :param Z:  numpy array of any shape (output from hidden layer)
    :param arr: numpy array of any shape (start, end)
    :return A: output of multinum_logit(Z,arr), same shape as Z
    :return cache: returns Z as well, useful during back propagation
    """
    A = np.zeros(Z.shape)
    for i in prange(len(arr)):
        shiftx = Z[:,arr[i,1]:arr[i,2]+1] - np.max(Z[:,int(arr[i,1]):int(arr[i,2])+1])
        A[:,arr[i,1]:arr[i,2]+1] = np.exp(shiftx)/np.exp(shiftx).sum()
    cache = Z
    return A,cache

因为这个for循环没有矢量化，所以它是我代码中的瓶颈。什么是一个可能的解决方案，使它更快。我试过使用

numba

的

@jit

，这使它稍微快一点，但还不够。我想知道是否有其他方法使它更快或矢量化/并行化

函数的示例输入数据

Z = np.random.random([1,10000])
arr = np.zeros([100,3])
arr[:,0] = 1
temp = int(Z.shape[1]/arr.shape[0])
for i in range(arr.shape[0]):
    arr[i,1] = i*temp
    arr[i,2] = (i+1)*temp-1
arr = arr.astype(int)

编辑：

---

我忘了在这里强调我的班级数量在变化。例如，批次1有10个类，批次2可能有15个类。所以我传递一个数组arr，它跟踪哪些行属于batch1，以此类推。这些批次不同于传统神经网络框架中的批次

---

在上面的示例中，

arr

跟踪行的起始索引和结束索引。因此，softmax函数中的分母将仅为索引位于起始索引和结束索引之间的观测值的总和。

这是一个矢量化的softmax函数。这是斯坦福大学cs231n课程conv网络作业的实施

该函数接受可优化的参数、输入数据、目标和正则化器。（您可以忽略正则化器，因为它引用了某些cs231n指定所独占的另一个类）

它返回参数的损失和梯度

def softmax_loss_vectorized(W, X, y, reg):
  """
  Softmax loss function, vectorized version.
  Inputs and outputs are the same as softmax_loss_naive.
  """
  # Initialize the loss and gradient to zero.

  loss = 0.0
  dW = np.zeros_like(W)

  num_train = X.shape[0]

  scores = X.dot(W)

  shift_scores = scores - np.amax(scores,axis=1).reshape(-1,1)

  softmax = np.exp(shift_scores)/np.sum(np.exp(shift_scores), axis=1).reshape(-1,1)

  loss = -np.sum(np.log(softmax[range(num_train), list(y)]))

  loss /= num_train

  loss += 0.5* reg * np.sum(W * W)

  dSoftmax = softmax.copy()

  dSoftmax[range(num_train), list(y)] += -1

  dW = (X.T).dot(dSoftmax)
  dW = dW/num_train + reg * W

  return loss, dW

为了便于比较，这里是同一方法的一个简单（非矢量化）实现

def softmax_loss_naive(W, X, y, reg):
  """
  Softmax loss function, naive implementation (with loops)
  Inputs have dimension D, there are C classes, and we operate on minibatches
  of N examples.
  Inputs:
  - W: A numpy array of shape (D, C) containing weights.
  - X: A numpy array of shape (N, D) containing a minibatch of data.
  - y: A numpy array of shape (N,) containing training labels; y[i] = c means
    that X[i] has label c, where 0 <= c < C.
  - reg: (float) regularization strength
  Returns a tuple of:
  - loss as single float
  - gradient with respect to weights W; an array of same shape as W
  """

  loss = 0.0
  dW = np.zeros_like(W)

  num_train = X.shape[0]
  num_classes = W.shape[1]

  for i in xrange(num_train):
      scores = X[i].dot(W)

      shift_scores = scores - max(scores)

      loss_i = -shift_scores[y[i]] + np.log(sum(np.exp(shift_scores)))
      loss += loss_i
      for j in xrange(num_classes):
          softmax = np.exp(shift_scores[j])/sum(np.exp(shift_scores))
          if j==y[i]:

              dW[:,j] += (-1 + softmax) * X[i]
          else:
              dW[:,j] += softmax *X[i]

  loss /= num_train

  loss += 0.5 * reg * np.sum(W * W)

  dW /= num_train + reg * W

  return loss, dW

def softmax_loss_naive（W、X、y、reg）：
"""
Softmax损耗函数，朴素实现（带循环）
输入有维度D，有C类，我们对小批量进行操作
有N个例子。
投入：
-W：包含权重的形状（D，C）的numpy数组。
-X：一个形状（N，D）的numpy数组，包含一小批数据。
-y：包含训练标签的形状（N，）的numpy数组；y[i]=c表示
X[i]有标签c，其中0我忘了在这里强调我的类的数量是变化的。例如，批1有10个类，批2可能有15个类。所以我传递一个数组arr，它跟踪哪些行属于批1等等。这些批与传统神经网络框架中的批不同，如果运行tmax每批单独计算，这个选项可以在上面的代码中工作我想我正在这样做。我正在为特定批次的行计算矢量化softmax。是的，你在上面的代码中这样做。我不确定你为什么添加额外的注释-我想你建议它改变了我答案的有效性。我正在寻找一些东西这个for循环更快。我知道它不能按照您建议的方式进行矢量化。根据您的回答，我觉得我的问题不清楚，因此我附加了我的问题1）数组顺序对性能有重大影响。在C顺序数组中，最后一个dim的变化最快。因此我建议使用zo exchange（z阵列中的dim0和dim1）。2）如果您的arr阵列确实是常规阵列，您可以将其集成到softmax函数中，只使用循环。这将提供另一个加速。