绘图过程中对tensor类型变量的操纵

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
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()