PyTorch-SGD-求函数最优值

面试遇到的一个基础问题:使用PyTorch求一个函数的最优值

代码如下:

import torch.optim as optim
import torch

# 参数为x,(1,1)
x = torch.randn([1],requires_grad=True)

def function(x):
    return x**2 - 2*x + 1

optimizer = optim.SGD([x], lr=0.01, momentum=0)

for step in range(11):
    y = torch.cos(x)
    optimizer.zero_grad()
    y.backward()
    optimizer.step()   #求一次SDG
    print('step {}: x={},y={}'.format(step, x[0], y[0]))

这里将x作为需要优化的参数,每次梯度下降更新x的值,直到目标函数的值达到最小

这里y即是目标函数

分享到