PageSourceSearch

https://ml-visualized.com/chapter1/optimizers.html

html ml-visualized.com collected 2026-10-03 19:59:16 UTC 87,423 bytes, 1,046 lines download raw bytes

1
2<!DOCTYPE html>
3
4
5<html lang="en" data-content_root="../" data-theme="light">
6
7  <head>
8    <meta charset="utf-8" />
9    <meta name="viewport" content="width=device-width, initial-scale=1.0" /><meta name="viewport" content="width=device-width, initial-scale=1" />
10
11    <title>Optimizers &#8212; Machine Learning Visualized</title>
12  
13  
14  
15  
15<script data-cfasync="false">
16    document.documentElement.dataset.mode = localStorage.getItem("mode") || "light";
17    document.documentElement.dataset.theme = localStorage.getItem("theme") || "light";
18  </script>
18
19  
20  <!-- Loaded before other Sphinx assets -->
21  <link href="../_static/styles/theme.css?digest=8d27b9dea8ad943066ae" rel="stylesheet" />
22<link href="../_static/styles/bootstrap.css?digest=8d27b9dea8ad943066ae" rel="stylesheet" />
23<link href="../_static/styles/pydata-sphinx-theme.css?digest=8d27b9dea8ad943066ae" rel="stylesheet" />
24
25  
26  <link href="../_static/vendor/fontawesome/6.5.1/css/all.min.css?digest=8d27b9dea8ad943066ae" rel="stylesheet" />
27  <link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.1/webfonts/fa-solid-900.woff2" />
28<link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.1/webfonts/fa-brands-400.woff2" />
29<link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.1/webfonts/fa-regular-400.woff2" />
30
31    <link rel="stylesheet" type="text/css" href="../_static/pygments.css?v=03e43079" />
32    <link rel="stylesheet" type="text/css" href="../_static/styles/sphinx-book-theme.css?v=a3416100" />
33    <link rel="stylesheet" type="text/css" href="../_static/togglebutton.css?v=9c3e77be" />
34    <link rel="stylesheet" type="text/css" href="../_static/copybutton.css?v=76b2166b" />
35    <link rel="stylesheet" type="text/css" href="../_static/mystnb.11b39860a7a0cbfd473a3ad8a317855267ff0bd372690045ca344a6b62be495e.css" />
36    <link rel="stylesheet" type="text/css" href="../_static/sphinx-thebe.css?v=4fa983c6" />
37    <link rel="stylesheet" type="text/css" href="../_static/sphinx-design.min.css?v=95c83b7e" />
38    <link rel="stylesheet" type="text/css" href="../_static/custom.css?v=3a394b7d" />
39  
40  <!-- Pre-loaded scripts that we'll load fully later -->
41  <link rel="preload" as="script" href="../_static/scripts/bootstrap.js?digest=8d27b9dea8ad943066ae" />
42<link rel="preload" as="script" href="../_static/scripts/pydata-sphinx-theme.js?digest=8d27b9dea8ad943066ae" />
43  
43<script src="../_static/vendor/fontawesome/6.5.1/js/all.min.js?digest=8d27b9dea8ad943066ae"></script>
43
44
45    
45<script src="../_static/documentation_options.js?v=9eb32ce0"></script>
45
46    
46<script src="../_static/doctools.js?v=9a2dae69"></script>
46
47    
47<script src="../_static/sphinx_highlight.js?v=dc90522c"></script>
47
48    
48<script src="../_static/clipboard.min.js?v=a7894cd8"></script>
48
49    
49<script src="../_static/copybutton.js?v=f281be69"></script>
49
50    
50<script src="../_static/scripts/sphinx-book-theme.js?v=887ef09a"></script>
50
51    
51<script>let toggleHintShow = 'Click to show';</script>
51
52    
52<script>let toggleHintHide = 'Click to hide';</script>
52
53    
53<script>let toggleOpenOnPrint = 'true';</script>
53
54    
54<script src="../_static/togglebutton.js?v=b8b12719"></script>
54
55    
55<script>var togglebuttonSelector = '.toggle, .admonition.dropdown';</script>
55
56    
56<script src="../_static/design-tabs.js?v=f930bc37"></script>
vendor: 76 bytes, lines 56-57
56
57    <script async="async" src="https://www.googletagmanager.com/gtag/js?id=
57G-9S2CQ3B60T
vendor: 16 bytes, lines 57-58
57"></script>
58    
58<script>
59                
vendor: 177 bytes, lines 59-62
59window.dataLayer = window.dataLayer || [];
60                function gtag(){ dataLayer.push(arguments); }
61                gtag('js', new Date());
62                gtag('config', '
62G-9S2CQ3B60T
vendor: 16 bytes, lines 62-63
62');
63            
63</script>
63
64    
64<script>const THEBE_JS_URL = "https://unpkg.com/[email protected]/lib/index.js"; const thebe_selector = ".thebe,.cell"; const thebe_selector_input = "pre"; const thebe_selector_output = ".output, .cell_output"</script>
64
65    
65<script async="async" src="../_static/sphinx-thebe.js?v=c100c467"></script>
65
66    
66<script>var togglebuttonSelector = '.toggle, .admonition.dropdown';</script>
66
67    
67<script>
68                
vendor: 177 bytes, lines 68-71
68window.dataLayer = window.dataLayer || [];
69                function gtag(){ dataLayer.push(arguments); }
70                gtag('js', new Date());
71                gtag('config', '
71G-9S2CQ3B60T
vendor: 16 bytes, lines 71-72
71');
72            
72</script>
72
73    
73<script>const THEBE_JS_URL = "https://unpkg.com/[email protected]/lib/index.js"; const thebe_selector = ".thebe,.cell"; const thebe_selector_input = "pre"; const thebe_selector_output = ".output, .cell_output"</script>
73
74    
74<script>window.MathJax = {"options": {"processHtmlClass": "tex2jax_process|mathjax_process|math|output_area"}}</script>
74
75    
75<script defer="defer" src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
75
76    
76<script>DOCUMENTATION_OPTIONS.pagename = 'chapter1/optimizers';</script>
76
77    <link rel="icon" href="../_static/logo.png"/>
78    <link rel="index" title="Index" href="../genindex.html" />
79    <link rel="search" title="Search" href="../search.html" />
80    <link rel="next" title="Interactive Linear Regression" href="interactive_linear_regression.html" />
81    <link rel="prev" title="Gradient Descent" href="linear_regression.html" />
82  <meta name="viewport" content="width=device-width, initial-scale=1"/>
83  <meta name="docsearch:language" content="en"/>
84  </head>
85  
86  
87  <body data-bs-spy="scroll" data-bs-target=".bd-toc-nav" data-offset="180" data-bs-root-margin="0px 0px -60%" data-default-mode="light">
88
89  
90  
91  <a id="pst-skip-link" class="skip-link" href="#main-content">Skip to main content</a>
92  
93  <div id="pst-scroll-pixel-helper"></div>
94  
95  <button type="button" class="btn rounded-pill" id="pst-back-to-top">
96    <i class="fa-solid fa-arrow-up"></i>
97    Back to top
98  </button>
99
100  
101  <input type="checkbox"
102          class="sidebar-toggle"
103          name="__primary"
104          id="__primary"/>
105  <label class="overlay overlay-primary" for="__primary"></label>
106  
107  <input type="checkbox"
108          class="sidebar-toggle"
109          name="__secondary"
110          id="__secondary"/>
111  <label class="overlay overlay-secondary" for="__secondary"></label>
112  
113  <div class="search-button__wrapper">
114    <div class="search-button__overlay"></div>
115    <div class="search-button__search-container">
116<form class="bd-search d-flex align-items-center"
117      action="../search.html"
118      method="get">
119  <i class="fa-solid fa-magnifying-glass"></i>
120  <input type="search"
121         class="form-control"
122         name="q"
123         id="search-input"
124         placeholder="Search this book..."
125         aria-label="Search this book..."
126         autocomplete="off"
127         autocorrect="off"
128         autocapitalize="off"
129         spellcheck="false"/>
130  <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd>K</kbd></span>
131</form></div>
132  </div>
133  
134    <header class="bd-header navbar navbar-expand-lg bd-navbar">
135    </header>
136  
137
138  <div class="bd-container">
139    <div class="bd-container__inner bd-page-width">
140      
141      
142      
143      <div class="bd-sidebar-primary bd-sidebar">
144        
145
146  
147  <div class="sidebar-header-items sidebar-primary__section">
148    
149    
150    
151    
152  </div>
153  
154    <div class="sidebar-primary-items__start sidebar-primary__section">
155        <div class="sidebar-primary-item">
156
157  
158
159<a class="navbar-brand logo" href="../index.html">
160  
161  
162  
163  
164  
165  
166    <p class="title logo__title">Machine Learning Visualized</p>
167  
168</a></div>
169        <div class="sidebar-primary-item">
170
171 
171<script>
172 document.write(`
173   <button class="btn navbar-btn search-button-field search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
174    <i class="fa-solid fa-magnifying-glass"></i>
175    <span class="search-button__default-text">Search</span>
176    <span class="search-button__kbd-shortcut"><kbd class="kbd-shortcut__modifier">Ctrl</kbd>+<kbd class="kbd-shortcut__modifier">K</kbd></span>
177   </button>
178 `);
179 </script>
179</div>
180        <div class="sidebar-primary-item"><nav class="bd-links bd-docs-nav" aria-label="Main">
181    <div class="bd-toc-item navbar-nav active">
182        
183        <ul class="nav bd-sidenav bd-sidenav__home-link">
184            <li class="toctree-l1">
185                <a class="reference internal" href="../index.html">
186                    Machine Learning Visualized
187                </a>
188            </li>
189        </ul>
190        <p aria-level="2" class="caption" role="heading"><span class="caption-text">Chapter 1. Optimization</span></p>
191<ul class="current nav bd-sidenav">
192<li class="toctree-l1"><a class="reference internal" href="linear_regression.html">Gradient Descent</a></li>
193<li class="toctree-l1 current active"><a class="current reference internal" href="#">Optimizers</a></li>
194<li class="toctree-l1"><a class="reference internal" href="interactive_linear_regression.html">Interactive Linear Regression</a></li>
195</ul>
196<p aria-level="2" class="caption" role="heading"><span class="caption-text">Chapter 2. Clustering and Reduction</span></p>
197<ul class="nav bd-sidenav">
198<li class="toctree-l1"><a class="reference internal" href="../chapter2/pca.html">Principal Component Analysis</a></li>
199<li class="toctree-l1"><a class="reference internal" href="../chapter2/k_means.html">K-Means Clustering</a></li>
200</ul>
201<p aria-level="2" class="caption" role="heading"><span class="caption-text">Chapter 3. Linear Models</span></p>
202<ul class="nav bd-sidenav">
203<li class="toctree-l1"><a class="reference internal" href="../chapter3/perceptron.html">Perceptron</a></li>
204<li class="toctree-l1"><a class="reference internal" href="../chapter3/logistic_regression.html">Logistic Regression</a></li>
205<li class="toctree-l1"><a class="reference internal" href="../chapter3/interactive_perceptron.html">Interactive Perceptron</a></li>
206<li class="toctree-l1"><a class="reference internal" href="../chapter3/interactive_logistic_regression.html">Interactive Logistic Regression</a></li>
207</ul>
208<p aria-level="2" class="caption" role="heading"><span class="caption-text">Chapter 4. Neural Networks</span></p>
209<ul class="nav bd-sidenav">
210<li class="toctree-l1"><a class="reference internal" href="../chapter4/neural_network.html">Neural Network</a></li>
211<li class="toctree-l1"><a class="reference internal" href="../chapter4/neural_network_weights.html">Neural Network Weights Visualization</a></li>
212<li class="toctree-l1"><a class="reference internal" href="../chapter4/autoencoder.html">Autoencoder</a></li>
213</ul>
214
215    </div>
216</nav></div>
217    </div>
218  
219  
220  <div class="sidebar-primary-items__end sidebar-primary__section">
221  </div>
222  
223  <div id="rtd-footer-container"></div>
224
225
226      </div>
227      
228      <main id="main-content" class="bd-main">
229        
230        
231
232<div class="sbt-scroll-pixel-helper"></div>
233
234          <div class="bd-content">
235            <div class="bd-article-container">
236              
237              <div class="bd-header-article">
238<div class="header-article-items header-article__inner">
239  
240    <div class="header-article-items__start">
241      
242        <div class="header-article-item"><button class="sidebar-toggle primary-toggle btn btn-sm" title="Toggle primary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip">
243  <span class="fa-solid fa-bars"></span>
244</button></div>
245      
246    </div>
247  
248  
249    <div class="header-article-items__end">
250      
251        <div class="header-article-item">
252
253<div class="article-header-buttons">
254
255
256
257
258
259<div class="dropdown dropdown-source-buttons">
260  <button class="btn dropdown-toggle" type="button" data-bs-toggle="dropdown" aria-expanded="false" aria-label="Source repositories">
261    <i class="fab fa-github"></i>
262  </button>
263  <ul class="dropdown-menu">
264      
265      
266      
267      <li><a href="https://github.com/gavinkhung/machine-learning-visualized" target="_blank"
268   class="btn btn-sm btn-source-repository-button dropdown-item"
269   title="Source repository"
270   data-bs-placement="left" data-bs-toggle="tooltip"
271>
272  
273
274<span class="btn__icon-container">
275  <i class="fab fa-github"></i>
276  </span>
277<span class="btn__text-container">Repository</span>
278</a>
279</li>
280      
281      
282      
283      
284      <li>
284<a href="https://github.com/gavinkhung/machine-learning-visualized/issues/new?title=Issue%20on%20page%20%2Fchapter1/optimizers.html&body=Your%20issue%20content%20here." target="_blank"
285   class="btn btn-sm btn-source-issues-button dropdown-item"
286   title="Open an issue"
287   data-bs-placement="left" data-bs-toggle="tooltip"
288>
289  
290
291<span class="btn__icon-container">
292  <i class="fas fa-lightbulb"></i>
293  </span>
294<span class="btn__text-container">Open issue</span>
295</a>
296</li>
297      
298  </ul>
299</div>
300
301
302
303
304
305
306<div class="dropdown dropdown-download-buttons">
307  <button class="btn dropdown-toggle" type="button" data-bs-toggle="dropdown" aria-expanded="false" aria-label="Download this page">
308    <i class="fas fa-download"></i>
309  </button>
310  <ul class="dropdown-menu">
311      
312      
313      
314      <li><a href="../_sources/chapter1/optimizers.ipynb" target="_blank"
315   class="btn btn-sm btn-download-source-button dropdown-item"
316   title="Download source file"
317   data-bs-placement="left" data-bs-toggle="tooltip"
318>
319  
320
321<span class="btn__icon-container">
322  <i class="fas fa-file"></i>
323  </span>
324<span class="btn__text-container">.ipynb</span>
325</a>
326</li>
327      
328      
329      
330      
331      <li>
332<button onclick="window.print()"
333  class="btn btn-sm btn-download-pdf-button dropdown-item"
334  title="Print to PDF"
335  data-bs-placement="left" data-bs-toggle="tooltip"
336>
337  
338
339<span class="btn__icon-container">
340  <i class="fas fa-file-pdf"></i>
341  </span>
342<span class="btn__text-container">.pdf</span>
343</button>
344</li>
345      
346  </ul>
347</div>
348
349
350
351
352<button onclick="toggleFullScreen()"
353  class="btn btn-sm btn-fullscreen-button"
354  title="Fullscreen mode"
355  data-bs-placement="bottom" data-bs-toggle="tooltip"
356>
357  
358
359<span class="btn__icon-container">
360  <i class="fas fa-expand"></i>
361  </span>
362
363</button>
364
365
366
367<script>
368document.write(`
369  <button class="btn btn-sm navbar-btn theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip">
370    <span class="theme-switch nav-link" data-mode="light"><i class="fa-solid fa-sun fa-lg"></i></span>
371    <span class="theme-switch nav-link" data-mode="dark"><i class="fa-solid fa-moon fa-lg"></i></span>
372    <span class="theme-switch nav-link" data-mode="auto"><i class="fa-solid fa-circle-half-stroke fa-lg"></i></span>
373  </button>
374`);
375</script>
375
376
377
378<script>
379document.write(`
380  <button class="btn btn-sm navbar-btn search-button search-button__button" title="Search" aria-label="Search" data-bs-placement="bottom" data-bs-toggle="tooltip">
381    <i class="fa-solid fa-magnifying-glass fa-lg"></i>
382  </button>
383`);
384</script>
384
385<button class="sidebar-toggle secondary-toggle btn btn-sm" title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip">
386    <span class="fa-solid fa-list"></span>
387</button>
388</div></div>
389      
390    </div>
391  
392</div>
393</div>
394              
395              
396
397<div id="jb-print-docs-body" class="onlyprint">
398    <h1>Optimizers</h1>
399    <!-- Table of contents -->
400    <div id="print-main-content">
401        <div id="jb-print-toc">
402            
403            <div>
404                <h2> Contents </h2>
405            </div>
406            <nav aria-label="Page">
407                <ul class="visible nav section-nav flex-column">
408<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#training-dataset">Training Dataset</a></li>
409<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#optimizer-update-rules">Optimizer Update Rules</a><ul class="nav section-nav flex-column">
410<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#stochastic-gradient-descent-sgd">Stochastic Gradient Descent (SGD)</a></li>
411<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#rmsprop">RMSprop</a></li>
412<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#adam">Adam</a></li>
413</ul>
414</li>
415<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#surface-gradient">Surface Gradient</a></li>
416<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#implementing-the-optimizers">Implementing the Optimizers</a></li>
417<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#graphing-functions">Graphing functions</a></li>
418<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#training-the-model">Training the Model</a></li>
419<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#output-gif">Output GIF</a></li>
420<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#pytorch-implementation">PyTorch Implementation</a></li>
421<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#training-the-model-with-pytorch">Training the Model with PyTorch</a></li>
422<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#output-gif-with-pytorch">Output GIF with PyTorch</a></li>
423</ul>
424            </nav>
425        </div>
426    </div>
427</div>
428
429              
430                
431<div id="searchbox"></div>
432                <article class="bd-article">
433                  
434  <section class="tex2jax_ignore mathjax_ignore" id="optimizers">
435<h1>Optimizers<a class="headerlink" href="#optimizers" title="Link to this heading">#</a></h1>
436<p>Optimizers are the algorithms that update a model’s parameters to minimize a function. Building on Gradient Descent, this notebook implements and visualizes more advanced optimizers, like Stochastic Gradient Descent (SGD), RMSprop, and Adam. Rather than fitting a dataset, we will watch each optimizer navigate a 3D surface, stepping from a shared starting point toward lower values of the function. This lets us compare how each optimizer behaves on the exact same problem.</p>
437<p>Importing the libraries</p>
438<div class="cell docutils container">
439<div class="cell_input docutils container">
440<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1"># import numpy as np</span>
441<span class="kn">import</span><span class="w"> </span><span class="nn">autograd.numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
442<span class="kn">from</span><span class="w"> </span><span class="nn">autograd</span><span class="w"> </span><span class="kn">import</span> <span class="n">grad</span>
443<span class="kn">from</span><span class="w"> </span><span class="nn">autograd</span><span class="w"> </span><span class="kn">import</span> <span class="n">elementwise_grad</span> <span class="k">as</span> <span class="n">egrad</span>
444
445<span class="o">%</span><span class="k">matplotlib</span> widget
446<span class="kn">
446import</span><span class="w"> </span><span class="nn">matplotlib.pyplot</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">plt</span>
447<span class="kn">from</span><span class="w"> </span><span class="nn">mpl_toolkits.mplot3d</span><span class="w"> </span><span class="kn">import</span> <span class="n">Axes3D</span>
448
449<span class="kn">from</span><span class="w"> </span><span class="nn">celluloid</span><span class="w"> </span><span class="kn">import</span> <span class="n">Camera</span>
450<span class="kn">import</span><span class="w"> </span><span class="nn">scienceplots</span>
451<span class="kn">from</span><span class="w"> </span><span class="nn">IPython.display</span><span class="w"> </span><span class="kn">import</span> <span class="n">Image</span>
452
453<span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
454<span class="n">plt</span><span class="o">.</span><span class="n">style</span><span class="o">.</span><span class="n">use</span><span class="p">([</span><span class="s2">"science"</span><span class="p">,</span> <span class="s2">"no-latex"</span><span class="p">])</span>
455</pre></div>
456</div>
457</div>
458</div>
459<section id="training-dataset">
460<h2>Training Dataset<a class="headerlink" href="#training-dataset" title="Link to this heading">#</a></h2>
461<p>Instead of a dataset, our optimizers will minimize a mathematical function. We will use the surface:</p>
462<p><span class="math notranslate nohighlight">\(z = x^2 - y^2\)</span></p>
463<p>This is a hyperbolic paraboloid, also known as a saddle. The surface curves upward along the <span class="math notranslate nohighlight">\(x\)</span> axis and downward along the <span class="math notranslate nohighlight">\(y\)</span> axis, creating a saddle point where the gradient is small in some directions. Saddle points are a classic challenge for optimizers: a naive optimizer can stall near the ridge, while adaptive optimizers escape more quickly. This makes the saddle a great way to compare optimizer behavior.</p>
464<div class="cell docutils container">
465<div class="cell_input docutils container">
466<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span><span class="w"> </span><span class="nf">surface_fn</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">
466Y</span><span class="p">):</span>
467    <span class="k">return</span> <span class="n">X</span><span class="o">**</span><span class="mi">2</span> <span class="o">-</span> <span class="n">Y</span><span class="o">**</span><span class="mi">2</span>
468
469
470<span class="k">def</span><span class="w"> </span><span class="nf">generate_function</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="n">surface_fn</span><span class="p">):</span>
471    <span class="n">x</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="o">-</span><span class="mi">5</span><span class="p">,</span> <span class="mi">25</span><span class="p">,</span> <span class="n">dims</span><span class="p">)</span>
472    <span class="n">y</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="o">-</span><span class="mi">5</span><span class="p">,</span> <span class="mi">25</span><span class="p">,</span> <span class="n">dims</span><span class="p">)</span>
473    <span class="n">X</span><span class="p">,</span> <span class="n">
473Y</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
474    <span class="n">Z</span> <span class="o">=</span> <span class="n">surface_fn</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">Y</span><span class="p">)</span>
475
476    <span class="k">return</span> <span class="n">X</span><span class="p">,</span> <span class="n">Y</span><span class="p">,</span> <span class="n">Z</span>
477
478
479<span class="n">dims</span> <span class="o">=</span> <span class="mi">25</span>
480<span class="n">X</span><span class="p">,</span> <span class="n">Y</span><span class="p">,</span> <span class="n">Z</span> <span class="o">=</span> <span class="n">generate_function</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="n">surface_fn</span><span class="p">)</span>
481
482<span class="c1"># Visualize the Hyperbolic Paraboloid</span>
483<span class="n">fig</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">figure</span><span class="p">()</span>
484<span class="n">ax</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">projection</span><span class="o">=</span><span class="s2">"3d"</span><span class="p">)</span>
485<span class="n">ax</span><span class="o">.</span><span class="n">view_init</span><span class="p">(</span><span class="mi">30</span><span class="p">,</span> <span class="mi">130</span><span class="p">)</span>
486
487<span class="n">ax</span><span class="o">.</span><span class="n">plot_surface</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">
487Y</span><span class="p">,</span> <span class="n">Z</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"grey"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.4</span><span class="p">)</span>
488<span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"X"</span><span class="p">)</span>
489<span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s2">"Y"</span><span class="p">)</span>
490<span class="n">ax</span><span class="o">.</span><span class="n">set_zlabel</span><span class="p">(</span><span class="s2">"Z"</span><span class="p">)</span>
491<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
492</pre></div>
493</div>
494</div>
495<div class="cell_output docutils container">
496<script type="application/vnd.jupyter.widget-view+json">{"model_id": "1e6499aef40a467abf8e8d152638e4ce", "version_major": 2, "version_minor": 0}</script>
496</div>
497</div>
498</section>
499<section id="optimizer-update-rules">
500<h2>Optimizer Update Rules<a class="headerlink" href="#optimizer-update-rules" title="Link to this heading">#</a></h2>
501<p>Every optimizer in this notebook follows the same core idea as gradient descent: repeatedly step the parameters <span class="math notranslate nohighlight">\(\theta\)</span> in the direction that decreases the function the most, which is the negative gradient. A single update step is:</p>
502<p><span class="math notranslate nohighlight">\(
503\begin{align*}
504\theta_{t+1} = \theta_t - \eta \nabla_{\theta} f(\theta_t)
505\end{align*}
506\)</span></p>
507<p>where <span class="math notranslate nohighlight">\(\eta\)</span> is the learning rate and <span class="math notranslate nohighlight">\(\nabla_{\theta} f(\theta_t)\)</span> is the gradient of the function with respect to the parameters at step <span class="math notranslate nohighlight">\(t\)</span>. The optimizers below differ in how they use the gradient, such as accumulating past gradients (momentum) or scaling each parameter’s step by a running estimate of its gradient magnitude (adaptive learning rates).</p>
508<section id="stochastic-gradient-descent-sgd">
509<h3>Stochastic Gradient Descent (SGD)<a class="headerlink" href="#stochastic-gradient-descent-sgd" title="Link to this heading">#</a></h3>
510<p>SGD is the simplest optimizer. It takes a step directly proportional to the negative gradient, scaled by the learning rate:</p>
511<p><span class="math notranslate nohighlight">\(
512\begin{align*}
513\theta_{t+1} = \theta_t - \eta \nabla_{\theta} f(\theta_t)
514\end{align*}
515\)</span></p>
516<p>SGD treats every parameter and every direction the same. Because the step depends only on the current gradient and a fixed learning rate, SGD can move very slowly in directions where the gradient is small, such as along the flat ridge of a saddle.</p>
517</section>
518<section id="rmsprop">
519<h3>RMSprop<a class="headerlink" href="#rmsprop" title="Link to this heading">#</a></h3>
520<p>RMSprop gives each parameter its own adaptive learning rate by keeping an exponentially decaying average of recent squared gradients, <span class="math notranslate nohighlight">\(s_t\)</span>:</p>
521<p><span class="math notranslate nohighlight">\(
522\begin{align*}
523s_{t+1} &amp;= \beta s_t + (1 - \beta) (\nabla_{\theta} f(\theta_t))^2 \\
524\theta_{t+1} &amp;= \theta_t - \frac{\eta}{\sqrt{s_{t+1}} + \epsilon} \nabla_{\theta} f(\theta_t)
525\end{align*}
526\)</span></p>
527<p>Dividing by <span class="math notranslate nohighlight">\(\sqrt{s_{t+1}}\)</span> shrinks the step for parameters with large, noisy gradients and enlarges it for parameters with small gradients. The small constant <span class="math notranslate nohighlight">\(\epsilon\)</span> prevents division by zero. This adaptivity helps RMSprop make steady progress even where SGD stalls.</p>
528</section>
529<section id="adam">
530<h3>Adam<a class="headerlink" href="#adam" title="Link to this heading">#</a></h3>
531<p>Adam (Adaptive Moment Estimation) combines momentum with RMSprop’s adaptive learning rates. It tracks an exponentially decaying average of the gradients (the first moment <span class="math notranslate nohighlight">\(m_t\)</span>) and of the squared gradients (the second moment <span class="math notranslate nohighlight">\(v_t\)</span>):</p>
532<p><span class="math notranslate nohighlight">\(
533\begin{align*}
534m_{t+1} &amp;= \beta_1 m_t + (1 - \beta_1) \nabla_{\theta} f(\theta_t) \\
535v_{t+1} &amp;= \beta_2 v_t + (1 - \beta_2) (\nabla_{\theta} f(\theta_t))^2
536\end{align*}
537\)</span></p>
538<p>Because <span class="math notranslate nohighlight">\(m_t\)</span> and <span class="math notranslate nohighlight">\(v_t\)</span> start at zero, they are biased toward zero early in training, so Adam applies a bias correction:</p>
539<p><span class="math notranslate nohighlight">\(
540\begin{align*}
541\hat{m}_{t+1} = \frac{m_{t+1}}{1 - \beta_1^{t+1}}, \quad \hat{v}_{t+1} = \frac{v_{t+1}}{1 - \beta_2^{t+1}}
542\end{align*}
543\)</span></p>
544<p>The parameters are then updated using the corrected moments:</p>
545<p><span class="math notranslate nohighlight">\(
546\begin{align*}
547\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{\hat{v}_{t+1}} + \epsilon} \hat{m}_{t+1}
548\end{align*}
549\)</span></p>
550<p>Combining momentum and adaptive scaling makes Adam a robust default optimizer that often converges quickly across many problems.</p>
551</section>
552</section>
553<section id="surface-gradient">
554<h2>Surface Gradient<a class="headerlink" href="#surface-gradient" title="Link to this heading">#</a></h2>
555<p>Each optimizer needs the gradient of the surface with respect to the current point <span class="math notranslate nohighlight">\(\theta = (x, y)\)</span>. For our surface <span class="math notranslate nohighlight">\(z = x^2 - y^2\)</span>, the partial derivatives are:</p>
556<p><span class="math notranslate nohighlight">\(
557\begin{align*}
558\frac{\partial z}{\partial x} = 2x, \quad \frac{\partial z}{\partial y} = -2y
559\end{align*}
560\)</span></p>
561<p>so the gradient is <span class="math notranslate nohighlight">\(\nabla z = (2x, -2y)\)</span>. The gradient grows with <span class="math notranslate nohighlight">\(|x|\)</span> along the upward-curving axis and with <span class="math notranslate nohighlight">\(|y|\)</span> along the downward-curving axis. This is the signal every optimizer below uses to decide its next step.</p>
562</section>
563<section id="implementing-the-optimizers">
564<h2>
564Implementing the Optimizers<a class="headerlink" href="#implementing-the-optimizers" title="Link to this heading">#</a></h2>
565<p>Let’s implement each optimizer from scratch as a class with a <code class="docutils literal notranslate"><span class="pre">step</span></code> method that takes the current parameters and gradient and returns the updated parameters. The stateful optimizers (RMSprop and Adam) keep their running averages between steps.</p>
566<div class="cell docutils container">
567<div class="cell_input docutils container">
568<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span><span class="w"> </span><span class="nf">surface_gradient</span><span class="p">(</span><span class="n">pt</span><span class="p">):</span>
569    <span class="n">x</span><span class="p">,</span> <span class="n">y</span> <span class="o">=</span> <span class="n">pt</span>
570    <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="mi">2</span> <span class="o">*</span> <span class="n">x</span><span class="p">,</span> <span class="o">-</span><span class="mi">2</span> <span class="o">*</span> <span class="n">y</span><span class="p">])</span>
571
572
573<span class="k">class</span><span class="w"> </span><span class="nc">SGD</span><span class="p">:</span>
574    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">learning_rate</span><span class="p">):</span>
575        <span class="bp">self</span><span class="o">.</span><span class="n">learning_rate</span> <span class="o">=</span> <span class="n">learning_rate</span>
576
577    <span class="k">def</span><span class="w"> </span><span class="nf">step</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">gradient</span><span class="p">):</span>
578        <span class="k">return</span> <span class="n">params</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">learning_rate</span> <span class="o">*</span> <span class="n">gradient</span>
579
580
581<span class="k">class</span><span class="w"> </span><span class="nc">RMSprop</span><span class="p">:</span>
582    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">learning_rate</span><span class="p">,</span> <span class="n">beta</span><span class="o">=</span><span class="mf">0.99</span><span class="p">,</span> <span class="n">epsilon</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">):</span>
583        <span class="bp">self</span><span class="o">.</span><span class="n">learning_rate</span> <span class="o">=</span> <span class="n">learning_rate</span>
584        <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">=</span> <span class="n">beta</span>
585        <span class="bp">self</span><span class="o">.</span><span class="n">epsilon</span> <span class="o">=</span> <span class="n">epsilon</span>
586        <span class="bp">self</span><span class="o">.</span><span class="n">s</span> <span class="o">=</span> <span class="mi">0</span>
587
588    <span class="k">def</span><span class="w"> </span><span class="nf">step</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">gradient</span><span class="p">):</span>
589        <span class="bp">self</span><span class="o">.</span><span class="n">s</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">s</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span><span class="p">)</span> <span class="o">*</span> <span class="n">gradient</span><span class="o">
589**</span><span class="mi">2</span>
590        <span class="k">return</span> <span class="n">params</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">learning_rate</span> <span class="o">*</span> <span class="n">gradient</span> <span class="o">/</span> <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">s</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">epsilon</span><span class="p">)</span>
591
592
593<span class="k">class</span><span class="w"> </span><span class="nc">Adam</span><span class="p">:</span>
594    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">learning_rate</span><span class="p">,</span> <span class="n">beta1</span><span class="o">=</span><span class="mf">0.9</span><span class="p">,</span> <span class="n">beta2</span><span class="o">=</span><span class="mf">0.999</span><span class="p">,</span> <span class="n">epsilon</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">):</span>
595        <span class="bp">self</span><span class="o">.</span><span class="n">learning_rate</span> <span class="o">=</span> <span class="n">learning_rate</span>
596        <span class="bp">self</span><span class="o">.</span><span class="n">beta1</span> <span class="o">=</span> <span class="n">beta1</span>
597        <span class="bp">self</span><span class="o">.</span><span class="n">beta2</span> <span class="o">=</span> <span class="n">beta2</span>
598        <span class="bp">self</span><span class="o">.</span><span class="n">epsilon</span> <span class="o">=</span> <span class="n">epsilon</span>
599        <span class="bp">self</span><span class="o">.</span><span class="n">m</span> <span class="o">=</span> <span class="mi">0</span>
600        <span class="bp">self</span><span class="o">.</span><span class="n">v</span> <span class="o">=</span> <span class="mi">0</span>
601        <span class="bp">self</span><span class="o">.</span><span class="n">t</span> <span class="o">=</span> <span class="mi">0</span>
602
603    <span class="k">def</span><span class="w"> </span><span class="nf">step</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">gradient</span><span class="p">):</span>
604        <span class="bp">self</span><span class="o">.</span><span class="n">t</span> <span class="o">+=</span> <span class="mi">1</span>
605        <span class="bp">self</span><span class="o">.</span><span class="n">m</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta1</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">m</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta1</span><span class="p">)</span> <span class="o">*</span> <span class="n">gradient</span>
606        <span class="bp">self</span><span class="o">.</span><span class="n">v</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta2</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">v</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta2</span><span class="p">)</span> <span class="o">*</span> <span class="n">gradient</span><span class="o">**</span><span class="mi">2</span>
607        <span class="n">m_hat</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">m</span> <span class="o">/</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="bp">
607self</span><span class="o">.</span><span class="n">beta1</span><span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">t</span><span class="p">)</span>
608        <span class="n">v_hat</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">v</span> <span class="o">/</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta2</span><span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">t</span><span class="p">)</span>
609        <span class="k">return</span> <span class="n">params</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">learning_rate</span> <span class="o">*</span> <span class="n">m_hat</span> <span class="o">/</span> <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">v_hat</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">epsilon</span><span class="p">)</span>
610</pre></div>
611</div>
612</div>
613</div>
614</section>
615<section id="graphing-functions">
616<h2>Graphing functions<a class="headerlink" href="#graphing-functions" title="Link to this heading">#</a></h2>
617<p>Let’s define helper functions to plot the graphs. <code class="docutils literal notranslate"><span class="pre">create_scatter_and_3d_plot</span></code> sets up two views: a 2D plot tracking each optimizer’s <span class="math notranslate nohighlight">\(z\)</span> value over the epochs, and a 3D plot showing each optimizer’s current position on the surface. <code class="docutils literal notranslate"><span class="pre">show_epoch</span></code> controls how often a frame is captured for the animation, sampling frequently early in training and less often later.</p>
618<div class="cell docutils container">
619<div class="cell_input docutils container">
620<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span><span class="w"> </span><span class="nf">create_scatter_and_3d_plot</span><span class="p">():</span>
621    <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">16</span> <span class="o">/</span> <span class="mf">9.0</span> <span class="o">*</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">4</span> <span class="o">*</span> <span class="mi">1</span><span class="p">))</span>
622    <span class="n">fig</span><span class="o">.</span><span class="n">suptitle</span><span class="p">(</span><span class="s2">"Optimizers"</span><span class="p">)</span>
623
624    <span class="n">ylim</span> <span class="o">=</span> <span class="p">(</span><span class="o">-</span><span class="mi">5</span><span class="p">,</span> <span class="mi">25</span><span class="p">)</span>
625    <span class="n">zlim</span> <span class="o">=</span> <span class="p">(</span><span class="o">-</span><span class="mi">650</span><span class="p">,</span> <span class="mi">550</span><span class="p">)</span>
626
627    <span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"Epoch"</span><span class="p">,</span> <span class="n">fontweight</span><span class="o">=</span><span class="s2">"normal"</span><span class="p">)</span>
628    <span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s2">"Z"</span><span class="p">,</span> <span class="n">fontweight</span><span class="o">=</span><span class="s2">"normal"</span><span class="p">)</span>
629    <span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">set_title</span><span class="p">(</span><span class="s2">"Z Value"</span><span class="p">)</span>
630    <span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">set_ylim</span><span class="p">((</span><span class="o">-</span><span class="mi">650</span><span class="p">,</span> <span class="mi">450</span><span class="p">))</span>
631
632    <span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">axis</span><span class="p">(</span><span class="s2">"off"</span><span class="p">)</span>
633    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">axis</span><span class="p">(</span><span class="s2">"off"</span><span class="p">)</span>
634
635    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">projection</span><span class="o">=</span><span class="s2">"3d"</span><span class="p">)</span>
636    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"X"</span><span class="p">)</span>
637    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s2">
637"Y"</span><span class="p">)</span>
638    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">set_zlabel</span><span class="p">(</span><span class="s2">"Z"</span><span class="p">)</span>
639    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">set_title</span><span class="p">(</span><span class="s2">"Current Position"</span><span class="p">)</span>
640    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">view_init</span><span class="p">(</span><span class="mi">30</span><span class="p">,</span> <span class="mi">130</span><span class="p">)</span>
641    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">set_xlim</span><span class="p">(</span><span class="n">ylim</span><span class="p">)</span>
642    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">set_ylim</span><span class="p">(</span><span class="n">ylim</span><span class="p">)</span>
643    <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">set_zlim</span><span class="p">(</span><span class="n">zlim</span><span class="p">)</span>
644
645    <span class="n">camera</span> <span class="o">=</span> <span class="n">Camera</span><span class="p">(</span><span class="n">fig</span><span class="p">)</span>
646    <span class="k">return</span> <span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">ax</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span> <span class="n">camera</span>
647
648
649<span class="k">def</span><span class="w"> </span><span class="nf">show_epoch</span><span class="p">(</span><span class="n">epoch</span><span class="p">):</span>
650    <span class="k">return</span> <span class="p">(</span>
651        <span class="p">(</span><span class="n">epoch</span> <span class="o">&lt;</span> <span class="mi">25</span> <span class="ow">and</span> <span class="n">epoch</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">==</span> <span class="mi">0</span><span class="p">)</span>
652        <span class="ow">or</span> <span class="p">(</span><span class="n">epoch</span> <span class="o">&lt;</span> <span class="mi">50</span> <span class="ow">and</span> <span class="n">epoch</span> <span class="o">%</span> <span class="mi">5</span> <span class="o">==</span> <span class="mi">0</span><span class="p">)</span>
653        <span class="ow">or</span> <span class="p">(</span><span class="n">epoch</span> <span class="o">&lt;=</span> <span class="mi">100</span> <span class="ow">and</span> <span class="n">epoch</span> <span class="o">%</span> <span class="mi">15</span> <span class="o">==</span> <span class="mi">0</span><span class="p">)</span>
654        <span class="ow">or</span> <span class="p">(</span><span class="n">epoch</span> <span class="o">&lt;=</span> <span class="mi">500</span> <span class="ow">and</span> <span class="n">epoch</span> <span class="o">%</span> <span class="mi">45</span> <span class="o">==</span> <span class="mi">0</span><span class="p">)</span>
655        <span class="ow">or</span> <span class="p">(</span><span class="n">epoch</span> <span class="o">&lt;=</span> <span class="mi">1000</span> <span class="ow">and</span> <span class="n">epoch</span> <span class="o">%</span> <span class="mi">50</span> <span class="o">==</span> <span class="mi">0</span><span class="p">)</span>
656        <span class="ow">or</span> <span class="n">epoch</span> <span class="o">%</span> <span class="mi">250</span> <span class="o">==</span> <span class="mi">0</span>
657    <span class="p">)</span>
658</pre></div>
659</div>
660</div>
661</div>
662</section>
663<section id="training-the-model">
664<h2>Training the Model<a class="headerlink" href="#training-the-model" title="Link to this heading">#</a></h2>
665<p>The <code class="docutils literal notranslate"><span class="pre">fit</span></code> function runs every optimizer in parallel from the same starting point. In each epoch it evaluates the surface, computes the gradient, and asks each optimizer for its next position, recording the trajectory so the optimizers can be animated side by side.</p>
666<div class="cell docutils container">
667<div class="cell_input docutils container">
668<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span><span class="w"> </span><span class="nf">fit</span><span class="p">(</span><span class="n">optimizers</span><span class="p">,</span> <span class="n">surface_fn</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">Y</span><span class="p">,</span> <span class="n">Z</span><span class="p">,</span> <span class="n">initial_pt</span><span class="p">,</span> <span class="n">epochs</span><span class="p">,</span> <span class="n">output_filename</span><span class="p">,</span> <span class="n">colors</span><span class="p">):</span>
669    <span class="n">optimizer_names</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">optimizers</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
670    <span class="n">points</span> <span class="o">=</span> <span class="p">[</span><span class="n">initial_pt</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">optimizers</span><span class="p">]</span>
671
672    <span class="n">z_idx</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">epochs</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
673    <span class="n">z_vals</span> <span class="o">=</span> <span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">full</span><span class="p">(</span><span class="n">epochs</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">optimizer_names</span><span class="p">))]</span>
674
675    <span class="n">z_ax</span><span class="p">,</span> <span class="n">predictions_ax</span><span class="p">,</span> <span class="n">camera</span> <span class="o">=</span> <span class="n">create_scatter_and_3d_plot</span><span class="p">()</span>
676
677    <span class="k">for</span> <span class="n">idx</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">epochs</span><span class="p">):</span>
678        <span class="n">prediction_legends</span> <span class="o">=</span> <span class="p">[]</span>
679        <span class="n">z_legends</span> <span class="o">=</span> <span class="p">[]</span>
680
681        <span class="k">for</span> <span class="n">opt_idx</span><span class="p">,</span> <span class="n">optimizer</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">optimizers</span><span class="o">.</span><span class="n">values</span><span class="p">()):</span>
682            <span class="n">pt</span> <span class="o">=</span> <span class="n">points</span><span class="p">[</span><span class="n">opt_idx</span><span class="p">]</span>
683            <span class="n">z</span> <span class="o">=</span> <span class="n">surface_fn</span><span class="p">(</span><span class="n">pt</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">pt</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
684            <span class="n">gradient</span> <span class="o">=</span> <span class="n">surface_gradient</span><span class="p">(</span><span class="n">pt</span><span class="p">)</span>
685            <span class="n">points</span><span class="p">[</span><span class="n">opt_idx</span><span class="p">]</span> <span class="o">=</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">pt</span><span class="p">,</span> <span class="n">gradient</span><span class="p">)</span>
686
687            <span class="k">if</span> <span class="n">show_epoch</span><span class="p">(</span><span class="n">idx</span><span class="p">):</span>
688                <span class="n">z_clamped</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">maximum</span><span class="p">(</span><span class="o">-</span><span class="mi">1000</span><span class="p">,</span> <span class="n">z</span><span class="p">)</span>
689
690                <span class="n">prediction_legend</span> <span class="o">=</span> <span class="n">predictions_ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
691                    <span class="n">pt</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">pt</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">z_clamped</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="n">colors</span><span class="p">[</span><span class="n">opt_idx</span><span class="p">]</span>
692                <span class="p">)</span>
693                <span class="n">prediction_legends</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">prediction_legend</span><span class="p">)</span>
694
695                <span class="n">z_vals</span><span class="p">[</span><span class="n">opt_idx</span><span class="p">][</span><span class="n">idx</span><span class="p">]</span> <span class="o">=</span> <span class="n">z_clamped</span>
696                <span class="n">visible</span> <span class="o">=</span> <span class="n">z_vals</span><span class="p">[</span><span class="n">opt_idx</span><span class="p">]</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span>
697                <span class="n">z_legend</span> <span class="o">=</span> <span class="n">z_ax</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
698                    <span class="n">z_idx</span><span class="p">[</span><span class="n">visible</span><span class="p">][:</span> <span class="n">idx</span> <span class="o">+</span> <span class="mi">1</span><span class="p">],</span>
699                    <span class="n">z_vals</span><span class="p">[</span><span class="n">opt_idx</span><span class="p">][</span><span class="n">visible</span><span class="p">][:</span> <span class="n">idx</span> <span class="o">+</span> <span class="mi">1</span><span class="p">],</span>
700                    <span class="n">color</span><span class="o">=</span><span class="n">colors</span><span class="p">[</span><span class="n">opt_idx</span><span class="p">],</span>
701                    <span class="n">alpha</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
702                <span class="p">)</span>
703                <span class="n">z_legends</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">z_legend</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
704
705        <span class="k">if</span> <span class="n">show_epoch</span><span class="p">(</span><span class="n">idx</span><span class="p">):</span>
706            <span class="n">predictions_ax</span><span class="o">.</span><span class="n">plot_surface</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">
706Y</span><span class="p">,</span> <span class="n">Z</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"grey"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>
707            <span class="n">predictions_ax</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span><span class="n">prediction_legends</span><span class="p">,</span> <span class="n">optimizer_names</span><span class="p">,</span> <span class="n">loc</span><span class="o">=</span><span class="s2">"upper right"</span><span class="p">)</span>
708            <span class="n">z_ax</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span><span class="n">z_legends</span><span class="p">,</span> <span class="n">optimizer_names</span><span class="p">,</span> <span class="n">loc</span><span class="o">=</span><span class="s2">"upper right"</span><span class="p">)</span>
709            <span class="n">camera</span><span class="o">.</span><span class="n">snap</span><span class="p">()</span>
710
711    <span class="n">animation</span> <span class="o">=</span> <span class="n">camera</span><span class="o">.</span><span class="n">animate</span><span class="p">()</span>
712    <span class="n">animation</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">output_filename</span><span class="p">,</span> <span class="n">writer</span><span class="o">=</span><span class="s2">"pillow"</span><span class="p">)</span>
713    <span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
714</pre></div>
715</div>
716</div>
717</div>
718<p>Let’s train all three optimizers from the same starting point and compare them. We start near the saddle’s ridge at <span class="math notranslate nohighlight">\((20.75, 0.5)\)</span>, where the gradient along the <span class="math notranslate nohighlight">\(y\)</span> axis is small. SGD, Adam, and RMSprop are each given the same learning rate, so the only difference is the update rule.</p>
719<div class="cell docutils container">
720<div class="cell_input docutils container">
721<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">initial_pt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="mf">20.75</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">])</span>
722
723<span class="n">learning_rate</span> <span class="o">=</span> <span class="mf">0.005</span>
724<span class="n">epochs</span> <span class="o">=</span> <span class="mi">4000</span>
725
726<span class="n">optimizers</span> <span class="o">=</span> <span class="p">{</span>
727    <span class="s2">"SGD"</span><span class="p">:</span> <span class="n">SGD</span><span class="p">(</span><span class="n">learning_rate</span><span class="p">),</span>
728    <span class="s2">"Adam"</span><span class="p">:</span> <span class="n">Adam</span><span class="p">(</span><span class="n">learning_rate</span><span class="p">),</span>
729    <span class="s2">"RMSprop"</span><span class="p">:</span> <span class="n">RMSprop</span><span class="p">(</span><span class="n">learning_rate</span><span class="p">),</span>
730<span class="p">}</span>
731<span class="n">colors</span> <span class="o">=</span> <span class="p">[</span><span class="s2">"blue"</span><span class="p">,</span> <span class="s2">"red"</span><span class="p">,</span> <span class="s2">"green"</span><span class="p">]</span>
732
733<span class="n">output_filename</span> <span class="o">=</span> <span class="s2">"optimizers.gif"</span>
734<span class="n">fit</span><span class="p">(</span><span class="n">optimizers</span><span class="p">,</span> <span class="n">surface_fn</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">
734Y</span><span class="p">,</span> <span class="n">Z</span><span class="p">,</span> <span class="n">initial_pt</span><span class="p">,</span> <span class="n">epochs</span><span class="p">,</span> <span class="n">output_filename</span><span class="p">,</span> <span class="n">colors</span><span class="p">)</span>
735</pre></div>
736</div>
737</div>
738<div class="cell_output docutils container">
739<script type="application/vnd.jupyter.widget-view+json">{"model_id": "845de8218c6a4a479641c33fb925ca12", "version_major": 2, "version_minor": 0}</script>
739</div>
740</div>
741</section>
742<section id="output-gif">
743<h2>Output GIF<a class="headerlink" href="#output-gif" title="Link to this heading">#</a></h2>
744<p>The animation below shows our from-scratch SGD, Adam, and RMSprop descending the saddle surface from the same starting point. The left plot tracks how quickly each optimizer drives the <span class="math notranslate nohighlight">\(z\)</span> value down, while the right plot shows each optimizer’s current position on the surface.</p>
745<div class="cell docutils container">
746<div class="cell_input docutils container">
747<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">Image</span><span class="p">(</span><span class="n">filename</span><span class="o">=</span><span class="n">output_filename</span><span class="p">)</span>
748</pre></div>
749</div>
750</div>
751<div class="cell_output docutils container">
752<img alt="../_images/9b79b273ed072157c89430e60a543aafbda3c7a40cee0b19e6dfebabd727b7b5.gif" src="../_images/9b79b273ed072157c89430e60a543aafbda3c7a40cee0b19e6dfebabd727b7b5.gif" />
753</div>
754</div>
755</section>
756<section id="pytorch-implementation">
757<h2>PyTorch Implementation<a class="headerlink" href="#pytorch-implementation" title="Link to this heading">#</a></h2>
758<p>Machine learning libraries, such as PyTorch, provide highly optimized implementations of these optimizers. We will use PyTorch to run the same three optimizers on the same surface and confirm our from-scratch results.</p>
759<div class="cell docutils container">
760<div class="cell_input docutils container">
761<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
762<span class="kn">import</span><span class="w"> </span><span class="nn">torch.nn</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">nn</span>
763
764<span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
765</pre></div>
766</div>
767</div>
768<div class="cell_output docutils container">
769<div class="output text_plain highlight-myst-ansi notranslate"><div class="highlight"><pre><span></span>&lt;torch._C.Generator at 0xffff2d619e70&gt;
770</pre></div>
771</div>
772</div>
773</div>
774<p>We represent the optimization problem as a small model whose only parameters are the coordinates of a point on the surface. The <code class="docutils literal notranslate"><span class="pre">forward</span></code> method evaluates the surface function at the current point, which is the value we want to minimize.</p>
775<div class="cell docutils container">
776<div class="cell_input docutils container">
777<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">class</span><span class="w"> </span><span class="nc">TorchNet</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
778    <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">initial_pt</span><span class="p">):</span>
779        <span class="nb">super</span><span class="p">(</span><span class="n">TorchNet</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
780
781        <span class="bp">self</span><span class="o">.</span><span class="n">pts</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="n">initial_pt</span><span class="p">)</span>
782
783    <span class="k">def</span><span class="w"> </span><span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">surface_fn</span><span class="p">):</span>
784        <span class="k">return</span> <span class="n">surface_fn</span><span class="p">(</span><span class="bp">
784self</span><span class="o">.</span><span class="n">pts</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">pts</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
785
786    <span class="k">def</span><span class="w"> </span><span class="nf">get_pts</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
787        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">pts</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>
788</pre></div>
789</div>
790</div>
791</div>
792<p>The training function creates one model and one optimizer per optimizer type, all starting from the same initial point. In each epoch, it evaluates the surface at every optimizer’s current point, backpropagates the gradient, and steps each optimizer.</p>
793<div class="cell docutils container">
794<div class="cell_input docutils container">
795<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span><span class="w"> </span><span class="nf">torch_fit</span><span class="p">(</span>
796    <span class="n">optimizer_types</span><span class="p">,</span> <span class="n">surface_fn</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">
796Y</span><span class="p">,</span> <span class="n">Z</span><span class="p">,</span> <span class="n">epochs</span><span class="p">,</span> <span class="n">learning_rate</span><span class="p">,</span> <span class="n">output_filename</span><span class="p">,</span> <span class="n">colors</span>
797<span class="p">):</span>
798    <span class="n">torch_models</span> <span class="o">=</span> <span class="p">[]</span>
799    <span class="n">torch_optimizers</span> <span class="o">=</span> <span class="p">[]</span>
800    <span class="n">optimizer_names</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">optimizer_types</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
801
802    <span class="n">z_idx</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">epochs</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
803    <span class="n">z_vals</span> <span class="o">=</span> <span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">full</span><span class="p">(</span><span class="n">epochs</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">optimizer_names</span><span class="p">))]</span>
804
805    <span class="k">for</span> <span class="n">optimizer</span> <span class="ow">in</span> <span class="n">optimizer_types</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
806        <span class="n">torch_model</span> <span class="o">=</span> <span class="n">TorchNet</span><span class="p">(</span><span class="n">initial_pt_tensor</span><span class="o">.</span><span class="n">clone</span><span class="p">())</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
807        <span class="n">torch_optimizer</span> <span class="o">=</span> <span class="n">optimizer</span><span class="p">(</span><span class="n">torch_model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="n">lr</span><span class="o">=</span><span class="n">learning_rate</span><span class="p">)</span>
808
809        <span class="n">torch_models</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">torch_model</span><span class="p">)</span>
810        <span class="n">torch_optimizers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">torch_optimizer</span><span class="p">)</span>
811
812    <span class="n">z_ax</span><span class="p">,</span> <span class="n">predictions_ax1</span><span class="p">,</span> <span class="n">camera1</span> <span class="o">=</span> <span class="n">create_scatter_and_3d_plot</span><span class="p">()</span>
813
814    <span class="k">for</span> <span class="n">idx</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">epochs</span><span class="p">):</span>
815        <span class="n">predictions_legends</span> <span class="o">=</span> <span class="p">[]</span>
816        <span class="n">z_legends</span> <span class="o">=</span> <span class="p">[]</span>
817
818        <span class="k">for</span> <span class="n">torch_idx</span><span class="p">,</span> <span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span>
819            <span class="nb">zip</span><span class="p">(</span><span class="n">torch_models</span><span class="p">,</span> <span class="n">torch_optimizers</span><span class="p">)</span>
820        <span class="p">):</span>
821            <span class="n">loss</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="n">surface_fn</span><span class="p">)</span>
822
823            <span class="c1"># optimizer</span>
824            <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>
825            <span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
826            <span class="n">optimizer</span><span class="o">.</span><span class="n">step</span><span class="p">()</span>
827
828            <span class="k">if</span> <span class="n">show_epoch</span><span class="p">(</span><span class="n">idx</span><span class="p">):</span>
829                <span class="n">loss_np</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">maximum</span><span class="p">(</span><span class="o">-</span><span class="mi">1000</span><span class="p">,</span> <span class="n">loss</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
830
831                <span class="n">model_pts</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">get_pts</span><span class="p">()</span>
832                <span class="n">predictions_legends_idx</span> <span class="o">=</span> <span class="n">predictions_ax1</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
833                    <span class="n">model_pts</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">model_pts</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">loss_np</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="n">colors</span><span class="p">[</span><span class="n">torch_idx</span><span class="p">]</span>
834                <span class="p">)</span>
835                <span class="n">predictions_legends</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">predictions_legends_idx</span><span class="p">)</span>
836
837                <span class="n">z_vals</span><span class="p">[</span><span class="n">torch_idx</span><span class="p">][</span><span class="n">idx</span><span class="p">]</span> <span class="o">=</span> <span class="n">loss_np</span>
838                <span class="n">visible_mse</span> <span class="o">=</span> <span class="n">z_vals</span><span class="p">[</span><span class="n">torch_idx</span><span class="p">]</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span>
839
840                <span class="n">z_legends_idx</span> <span class="o">=</span> <span class="n">z_ax</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
841                    <span class="n">z_idx</span><span class="p">[</span><span class="n">visible_mse</span><span class="p">][:</span> <span class="n">idx</span> <span class="o">+</span> <span class="mi">1</span><span class="p">],</span>
842                    <span class="n">z_vals</span><span class="p">[</span><span class="n">torch_idx</span><span class="p">][</span><span class="n">visible_mse</span><span class="p">][:</span> <span class="n">idx</span> <span class="o">+</span> <span class="mi">1</span><span class="p">],</span>
843                    <span class="n">color</span><span class="o">=</span><span class="n">colors</span><span class="p">[</span><span class="n">torch_idx</span><span class="p">],</span>
844                    <span class="n">alpha</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
845                <span class="p">)</span>
846                <span class="n">z_legends</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">z_legends_idx</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
847
848        <span class="k">if</span> <span class="n">show_epoch</span><span class="p">(</span><span class="n">idx</span><span class="p">):</span>
849            <span class="n">predictions_ax1</span><span class="o">.</span><span class="n">plot_surface</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">Y</span><span class="p">,</span> <span class="n">Z</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"grey"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>
850            <span class="n">predictions_ax1</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span>
851                <span class="n">predictions_legends</span><span class="p">,</span> <span class="n">optimizer_names</span><span class="p">,</span> <span class="n">loc</span><span class="o">=</span><span class="s2">"upper right"</span>
852            <span class="p">)</span>
853            <span class="n">z_ax</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span><span class="n">z_legends</span><span class="p">,</span> <span class="n">optimizer_names</span><span class="p">,</span> <span class="n">loc</span><span class="o">=</span><span class="s2">"upper right"</span><span class="p">)</span>
854
855            <span class="n">camera1</span><span class="o">.</span><span class="n">snap</span><span class="p">()</span>
856
857    <span class="n">animation1</span> <span class="o">=</span> <span class="n">camera1</span><span class="o">.</span><span class="n">animate</span><span class="p">()</span>
858    <span class="n">animation1</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">output_filename</span><span class="p">,</span> <span class="n">writer</span><span class="o">=</span><span class="s2">"pillow"</span><span class="p">)</span>
859    <span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
860</pre></div>
861</div>
862</div>
863</div>
864</section>
865<section id="training-the-model-with-pytorch">
866<h2>Training the Model with PyTorch<a class="headerlink" href="#training-the-model-with-pytorch" title="Link to this heading">#</a></h2>
867<p>Let’s run PyTorch’s SGD, Adam, and RMSprop from the same starting point and learning rate we used for our from-scratch implementation.</p>
868<div class="cell docutils container">
869<div class="cell_input docutils container">
870<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">device</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">(</span><span class="s2">"cuda:0"</span> <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">is_available</span><span class="p">()</span> <span class="k">else</span> <span class="s2">"cpu"</span><span class="p">)</span>
871
872<span class="n">initial_pt</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="mf">20.75</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">])</span>
873<span class="n">initial_pt_tensor</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">initial_pt</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">
873float32</span><span class="p">)</span>
874
875<span class="n">optimizer_types</span> <span class="o">=</span> <span class="p">{</span>
876    <span class="s2">"SGD"</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">optim</span><span class="o">.</span><span class="n">SGD</span><span class="p">,</span>
877    <span class="s2">"Adam"</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">optim</span><span class="o">.</span><span class="n">Adam</span><span class="p">,</span>
878    <span class="s2">"RMSprop"</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">optim</span><span class="o">.</span><span class="n">RMSprop</span><span class="p">,</span>
879<span class="p">}</span>
880<span class="n">colors</span> <span class="o">=</span> <span class="p">[</span><span class="s2">"blue"</span><span class="p">,</span> <span class="s2">"red"</span><span class="p">,</span> <span class="s2">"green"</span><span class="p">]</span>
881
882<span class="n">learning_rate</span> <span class="o">=</span> <span class="mf">0.005</span>
883<span class="n">epochs</span> <span class="o">=</span> <span class="mi">4000</span>
884
885<span class="n">output_filename_pytorch</span> <span class="o">=</span> <span class="s2">"optimizers_pytorch.gif"</span>
886<span class="n">torch_fit</span><span class="p">(</span>
887    <span class="n">optimizer_types</span><span class="p">,</span>
888    <span class="n">surface_fn</span><span class="p">,</span>
889    <span class="n">X</span><span class="p">,</span>
890    <span class="n">
890Y</span><span class="p">,</span>
891    <span class="n">Z</span><span class="p">,</span>
892    <span class="n">epochs</span><span class="p">,</span>
893    <span class="n">learning_rate</span><span class="p">,</span>
894    <span class="n">output_filename_pytorch</span><span class="p">,</span>
895    <span class="n">colors</span><span class="p">,</span>
896<span class="p">)</span>
897</pre></div>
898</div>
899</div>
900<div class="cell_output docutils container">
901<script type="application/vnd.jupyter.widget-view+json">{"model_id": "205d8dafe73c456b94ea4608541a8b83", "version_major": 2, "version_minor": 0}</script>
901</div>
902</div>
903</section>
904<section id="output-gif-with-pytorch">
905<h2>Output GIF with PyTorch<a class="headerlink" href="#output-gif-with-pytorch" title="Link to this heading">#</a></h2>
906<p>The PyTorch optimizers produce trajectories that match our from-scratch implementation, confirming our derivations of SGD, RMSprop, and Adam.</p>
907<div class="cell docutils container">
908<div class="cell_input docutils container">
909<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">Image</span><span class="p">(</span><span class="n">filename</span><span class="o">=</span><span class="n">output_filename_pytorch</span><span class="p">)</span>
910</pre></div>
911</div>
912</div>
913<div class="cell_output docutils container">
914<img alt="../_images/bbbc788851805c529bbdaa488ebb7cb8e820be8c61d00adcb8dbf6fd44163fec.gif" src="../_images/bbbc788851805c529bbdaa488ebb7cb8e820be8c61d00adcb8dbf6fd44163fec.gif" />
915</div>
916</div>
917</section>
918</section>
919
920    
920<script type="text/x-thebe-config">
921    {
922        requestKernel: true,
923        binderOptions: {
924            repo: "binder-examples/jupyter-stacks-datascience",
925            ref: "master",
926        },
927        codeMirrorConfig: {
928            theme: "abcdef",
929            mode: "python"
930        },
931        kernelOptions: {
932            name: "python3",
933            path: "./chapter1"
934        },
935        predefinedOutput: true
936    }
937    </script>
937
938    
938<script>kernelName = 'python3'</script>
938
939
940                </article>
941              
942
943              
944              
945              
946              
947                <footer class="prev-next-footer">
948                  
949<div class="prev-next-area">
950    <a class="left-prev"
951       href="linear_regression.html"
952       title="previous page">
953      <i class="fa-solid fa-angle-left"></i>
954      <div class="prev-next-info">
955        <p class="prev-next-subtitle">previous</p>
956        <p class="prev-next-title">Gradient Descent</p>
957      </div>
958    </a>
959    <a class="right-next"
960       href="interactive_linear_regression.html"
961       title="next page">
962      <div class="prev-next-info">
963        <p class="prev-next-subtitle">next</p>
964        <p class="prev-next-title">Interactive Linear Regression</p>
965      </div>
966      <i class="fa-solid fa-angle-right"></i>
967    </a>
968</div>
969                </footer>
970              
971            </div>
972            
973            
974              
975                <div class="bd-sidebar-secondary bd-toc"><div class="sidebar-secondary-items sidebar-secondary__inner">
976
977
978  <div class="sidebar-secondary-item">
979  <div class="page-toc tocsection onthispage">
980    <i class="fa-solid fa-list"></i> Contents
981  </div>
982  <nav class="bd-toc-nav page-toc">
983    <ul class="visible nav section-nav flex-column">
984<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#training-dataset">Training Dataset</a></li>
985<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#optimizer-update-rules">Optimizer Update Rules</a><ul class="nav section-nav flex-column">
986<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#stochastic-gradient-descent-sgd">Stochastic Gradient Descent (SGD)</a></li>
987<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#rmsprop">RMSprop</a></li>
988<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#adam">Adam</a></li>
989</ul>
990</li>
991<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#surface-gradient">Surface Gradient</a></li>
992<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#implementing-the-optimizers">Implementing the Optimizers</a></li>
993<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#graphing-functions">Graphing functions</a></li>
994<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#training-the-model">Training the Model</a></li>
995<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#output-gif">Output GIF</a></li>
996<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#pytorch-implementation">PyTorch Implementation</a></li>
997<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#training-the-model-with-pytorch">Training the Model with PyTorch</a></li>
998<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#output-gif-with-pytorch">Output GIF with PyTorch</a></li>
999</ul>
1000  </nav></div>
1001
1002</div></div>
1003              
1004            
1005          </div>
1006          <footer class="bd-footer-content">
1007            
1008<div class="bd-footer-content__inner container">
1009  
1010  <div class="footer-item">
1011    
1012<p class="component-author">
1013By Gavin Hung
1014</p>
1015
1016  </div>
1017  
1018  <div class="footer-item">
1019    
1020
1021  </div>
1022  
1023  <div class="footer-item">
1024    
1025  </div>
1026  
1027  <div class="footer-item">
1028    
1029  </div>
1030  
1031</div>
1032          </footer>
1033        
1034
1035      </main>
1036    </div>
1037  </div>
1038  
1039  <!-- Scripts loaded after <body> so the DOM is not blocked -->
1040  
1040<script src="../_static/scripts/bootstrap.js?digest=8d27b9dea8ad943066ae"></script>
vendor: 1 bytes, line 1040
1040
1041<script src="../_static/scripts/pydata-sphinx-theme.js?digest=8d27b9dea8ad943066ae"></script>
1041
1042
1043  <footer class="bd-footer">
1044  </footer>
1045  </body>
1046</html>

Line numbers count LF bytes from the start of the resource, as the search results do. Vendor segments are library code the classifier recognised; they are stored but not indexed. Bytes are shown as Latin1 characters, one per byte.