避免在 autograd 中分配数组

Question

我从autograd 教程中了解到，当数组包含在要区分的目标中时，不支持数组分配。但是，我目前在我的代码中有以下目标函数，我想对 theta 进行区分：

def obj(theta):
    """
    Computes the objective function to be differentiated.

    Args:
        theta: np.array of shape (n, d)

    Return:
        res: np.array of shape (n,)
    """
    theta = np.atleast_2d(theta)
    n = theta.shape[0]

    res = np.zeros(n)  # Scores
    for i in xrange(n):
        res[i] = ... # Do some computations with theta[i, :]

    return res

通常我可以通过对 theta 上的计算进行矢量化来避免 for 循环；但是，在这种情况下，计算已经涉及给定特定行的 theta（作为超参数）的各种线性代数运算（逆运算等），并且我发现很难对所有行的 theta 运算进行矢量化。在这种情况下，我不知道有比用 for 循环逐行填充 res 数组更好的方法。

我尝试了一种简单的方法来避免数组分配，方法是创建一个列表并在每次迭代时将结果附加到该列表中，然后最终在返回 res 时将列表转换为数组，但最后我得到全零梯度...

我想知道此设置中的一般推荐解决方案是什么？

Answer 1

您可以使用numpy.apply_along_axis 为数据中的某个轴应用函数。

def func(row):
    # return the computation result for "row"

def obj(theta):
    """
    Computes the objective function to be differentiated.

    Args:
        theta: np.array of shape (n, d)

    Return:
        res: np.array of shape (n,)
    """
    theta = np.atleast_2d(theta)
    n = theta.shape[0]

    res = np.apply_along_axis(func1d=func, axis=1, arr=a)

    return res