使用 PyTorch 的具有 D 维的 Rosenbrock 函数

Question

如何使用 PyTorch 实现具有 D 维度的 Rosenbrock 函数？

Create the variables, where D is the number of dimensions and N is the number of elements.

x = (xmax - xmin)*torch.rand(N,D).type(dtype) + xmin

Function :

$\sum_{i=1}^{d-1}\left [ 100{(x_{i+1}-x_{i}^{2})}^{2}+{(x_i-1)}^{2} \right ]$

Using straight Python I'd do something like this:

fit = 0
for i in range(D-1):
    term1 = x[i + 1] - x[i]**2
    term2 = 1 - x[i]
    fit = fit + 100 * term1**2 + term2**2

My attempt using Pytorch:

def Rosenbrock(x):
    return torch.sum(100*(x - x**2)**2 + (x-1)**2)

I just do not know how to do the x[i+1] without using a for loop. How can I deal with it? Thank you!

Answer 1

Numpy 具有 roll 函数，我认为它非常有用。不幸的是，我不知道 pytorch 有任何类似于 numpy.roll 的功能。

在我的尝试中，x 是一个形式为 DxN 的 numpy 数组。首先，我们使用 roll 将第一维 (axis=0) 中的项目向左移动一个位置。像这样，每次我们比较 x_1[i] 就像我们做 x[i+1] 一样。然后，由于该函数只需要 D-1 个元素求和，我们删除了最后一列用 [:-1, :] 对 pytorch 张量进行切片。那么总结真的和你发布的代码很相似，只是在正确的地方把 x 换成了 x_1。

def Rosenbrock(x):
    x_1 = torch.from_numpy(np.roll(x, -1, axis=0)).float()[:-1, :]
    x = torch.from_numpy(x).float()[:-1, :]

    return torch.sum(100 * (x_1 - x ** 2) ** 2 + (x - 1) ** 2, 0)

同样，通过使用 roll，您可以在 numpy 版本中删除 for 循环

使用 PyTorch 的具有 D 维的 Rosenbrock 函数

Rosenbrock function with D dimension using PyTorch

python

pytorch