import numpy as np
import plotly.graph_objects as go
import pandas as pd
import torch
def Rosenbrock(x):
f = ((x[1] - x[0] ** 2) ** 2) * 100 + (1 - x[0]) ** 2
return f
x = np.arange(-2, 2, 0.05)
y = np.arange(-1, 3, 0.05)
X, Y = np.meshgrid(x, y)
Z = Rosenbrock([X, Y])
z_min = Z.min()
x = torch.tensor([-1.2, 2], requires_grad=True)
f = Rosenbrock(x)
optimizer = torch.optim.Adam([x,])
delta_x = torch.tensor([1.])
opt_step = 0
x_trajectory = []
y_trajectory = []
f_trajectory = []
while f > 0.001 + z_min:
x_temp = x.detach().numpy().copy()
if opt_step % 1000 == 0:
x_trajectory.append(x_temp[0])
y_trajectory.append(x_temp[1])
f_trajectory.append(f.detach().numpy().copy())
optimizer.zero_grad()
f.backward()
optimizer.step()
f = Rosenbrock(x)
opt_step = opt_step + 1
Z_df = pd.DataFrame(Z)
fig = go.Figure(data=[go.Surface(z=Z_df, x=X, y=Y)])
fig.update_traces(contours_z=dict(show=True, usecolormap=True,
highlightcolor="limegreen", project_z=True,
start=0, end=2500, size=100))
fig.update_layout(title='Rosenbrock function surface', autosize=True)
fig.add_traces(data=go.Scatter3d(
x=np.array(x_trajectory), y=np.array(y_trajectory), z=np.array(f_trajectory),
marker=dict(
size=4, color=np.array(f_trajectory), colorscale='reds'
),
line=dict(
color='red', width=2
)
))
fig.show()
结果图 |
---|
难点:
- 绘图过程中第三方包只对array和list类型变量操作,所以要把tensor变量转换成array变量,转换方法是np.array(x)
- tensor类型变量进行赋值时,给别人的都是地址,因此使用append函数构造序列,序列中只有一个值,这个值就是最后一位。解决方案可以是,用copy()函数将tensor类型变量中的值赋给一个新的变量,然后再append。 传递地址≠传递值,很有意思的点。