Skip to content

Instantly share code, notes, and snippets.

@vankesteren
Last active September 24, 2022 13:22
Show Gist options
  • Save vankesteren/4f1060943ddc1369e50b8b9ec3285968 to your computer and use it in GitHub Desktop.
Save vankesteren/4f1060943ddc1369e50b8b9ec3285968 to your computer and use it in GitHub Desktop.
Gradient descent for linear regression using Zygote AutoDiff
# Gradient descent with autodiff for linear regression
using Zygote
# Data
X = randn(1000, 10)
b = (1:10)
y = X * b + randn(1000)
# MSE for linear model
function mse(bhat)
res = y - X * bhat
res'res / length(res)
end
# Gradient w.r.t. parameters
grad(bhat) = gradient(mse, bhat)[1]
# Initialize params & control
bhat = zeros(10)
step = 1.0e-3
maxit = 1_000_000
tol = eps(1.0)
# gradient descent
for i = 1:maxit
global bhat
Δ = step .* grad(bhat)
if'Δ < tol)
break
end
bhat -= Δ
print("loss ", i, " ", mse(bhat), "\n")
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment