目标函数

所需模块numpy,matplotlib,torch

import numpyas np

from mpl_toolkits.mplot3dimport  Axes3D

from matplotlibimport pyplotas plt

import torch

定义函数

def himmelblau(x):

return (x[0]**2 +3*x[1] -18)**2 +(x[0] +x[1]**2 -11)**2

画出初略图

x = np.arange(-6,6,0.1)

y = np.arange(-6,6,0.1)

print('x,y range:',x.shape,y.shape)

X,Y = np.meshgrid(x,y)

print('X,Y maps:',X.shape,Y.shape)

Z = himmelblau([X,Y])

fig = plt.figure('himmelblau')

ax = fig.gca(projection='3d')

ax.plot_surface(X, Y, Z)

ax.view_init(60, -30)

ax.set_xlabel('x')

ax.set_ylabel('y')

plt.show()

效果图

利用梯度下降求最小值

x = torch.tensor([5.,2.],requires_grad=True)

optimizer = torch.optim.Adam([x],lr=1e-3)

for stepin range(40000):

pred = himmelblau(x)

optimizer.zero_grad()

pred.backward()

optimizer.step()

if step %2000 ==0:

print ('step {}: x = {}, f(x) = {}'

              .format(step, x.tolist(), pred.item()))

运行结果

step 0: x = [4.999000072479248, 1.9989999532699585], f(x) = 173.0

step 2000: x = [3.8468220233917236, 1.9102824926376343], f(x) = 18.69461441040039

step 4000: x = [3.1986043453216553, 2.73587703704834], f(x) = 0.29346320033073425

step 6000: x = [3.093132495880127, 2.8116304874420166], f(x) = 8.19714114186354e-06

step 8000: x = [3.0925745964050293, 2.8120086193084717], f(x) = 2.812157617881894e-09

step 10000: x = [3.0925676822662354, 2.812013626098633], f(x) = 3.637978807091713e-10

step 12000: x = [3.0925650596618652, 2.8120155334472656], f(x) = 4.729372449219227e-11

step 14000: x = [3.092564105987549, 2.812016248703003], f(x) = 0.0

step 16000: x = [3.092564105987549, 2.812016248703003], f(x) = 0.0

step 18000: x = [3.092564105987549, 2.812016248703003], f(x) = 0.0