python-梯度下降法求函数极小值

import numpy as np

# 定义func2函数

def func2(x):
    return x[0]**2 + x[1]**2

# 求func2函数在点[3.0, 4.0]处的梯度

def numerical_diff(f, x):
    h = 1e-4
    grap = np.zeros_like(x)
    for idx in range(x.size):
        tmp_val = x[idx]
        x[idx] = tmp_val + h
        fxh1 = f(x)
        
        x[idx] = tmp_val - h
        fxh2 = f(x)
        
        grap[idx] = (fxh1 - fxh2)/(h*2)
        x[idx] = tmp_val
    return grap

# 用梯度法求func2的最小值
def gradient_descent(f, init_x, lr=0.01, step_num=500):
    x = init_x
    for i in range(step_num):
        grap = numerical_diff(f, x)
        x -= lr * grap
    return x

gradient_descent(func2,np.array([3.0, 4.0]))

array([0.00012307, 0.0001641 ])