02-DP-linear-fit.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only if using JuliaBox\n",
    "pop!(DEPOT_PATH)\n",
    "push!(DEPOT_PATH, \"/home/jrun/.julia2/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear fit by differentiable programming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Zygote, LinearAlgebra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a data type representing our model (optional)\n",
    "struct LinearRegression\n",
    "    weights::Matrix\n",
    "    bias::Array{Float64,0}\n",
    "end\n",
    "LinearRegression(nparams) = LinearRegression(randn(1, nparams), fill(0.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define how to \"run\" the model forwards to generate a prediction\n",
    "function predict(model::LinearRegression, X)\n",
    "    return model.weights * X .+ model.bias[]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the \"loss\" to be minimized\n",
    "function loss(model::LinearRegression, X, Y)\n",
    "    return norm(predict(model, X) .- Y, 2)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up a sample problem.\n",
    "# \"Ground truth\" values we will try to learn/recover\n",
    "weights_gt = [1.0, 2.7, 0.3, 1.2]'\n",
    "bias_gt = 0.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a fake dataset of many observations\n",
    "X = randn(length(weights_gt), 10000)\n",
    "Y = weights_gt * X .+ bias_gt\n",
    "# Add a little bit of noise to `X` so that we do not have an exact solution,\n",
    "# but must instead do a least-squares fit\n",
    "X .+= 0.01 .* randn(size(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start with a random model to train\n",
    "model = LinearRegression(size(X, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate gradient on `model` for the first example in our training set\n",
    "grads = Zygote.gradient(m -> loss(m, X[:, 1], Y[1]), model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The `grads` object is a tuple containing one element per argument to\n",
    "# `gradient`, so we take the first one to get the gradient on `model`\n",
    "grads = grads[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define an update rule that will allow us to modify the weights\n",
    "# of our model a bit according to the gradients\n",
    "function sgd_update!(model::LinearRegression, grads, η = 0.001)\n",
    "    model.weights .-= η .* grads.weights\n",
    "    model.bias .-= η .* grads.bias\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx in 1:size(X, 2)\n",
    "    grads = Zygote.gradient(m -> loss(m, X[:, idx], Y[idx]), model)[1]\n",
    "    sgd_update!(model, grads)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_gt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.0.3",
   "language": "julia",
   "name": "julia-1.0"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.0.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}