Distribution utilities

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/api/distribution/distribution_utils/" />  <meta property="og:title" content="Keras documentation: Distribution utilities"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Distribution utilities"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Distribution utilities</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-link active" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-sublink" href="/api/models/">Models API</a> <a class="nav-sublink" href="/api/layers/">Layers API</a> <a class="nav-sublink" href="/api/callbacks/">Callbacks API</a> <a class="nav-sublink" href="/api/ops/">Ops API</a> <a class="nav-sublink" href="/api/optimizers/">Optimizers</a> <a class="nav-sublink" href="/api/metrics/">Metrics</a> <a class="nav-sublink" href="/api/losses/">Losses</a> <a class="nav-sublink" href="/api/data_loading/">Data loading</a> <a class="nav-sublink" href="/api/datasets/">Built-in small datasets</a> <a class="nav-sublink" href="/api/applications/">Keras Applications</a> <a class="nav-sublink" href="/api/mixed_precision/">Mixed precision</a> <a class="nav-sublink active" href="/api/distribution/">Multi-device distribution</a> <a class="nav-sublink2" href="/api/distribution/layout_map/">LayoutMap API</a> <a class="nav-sublink2" href="/api/distribution/data_parallel/">DataParallel API</a> <a class="nav-sublink2" href="/api/distribution/model_parallel/">ModelParallel API</a> <a class="nav-sublink2" href="/api/distribution/model_parallel/">ModelParallel API</a> <a class="nav-sublink2 active" href="/api/distribution/distribution_utils/">Distribution utilities</a> <a class="nav-sublink" href="/api/random/">RNG API</a> <a class="nav-sublink" href="/api/utils/">Utilities</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> <span class="k-location-slug-pointer">►</span> <a href='/api/'>Keras 3 API documentation</a> / <a href='/api/distribution/'>Multi-device distribution</a> / Distribution utilities </div> <div class='k-content'> <h1 id="distribution-utilities">Distribution utilities</h1> <p><span style="float:right;"><a href="https://github.com/keras-team/keras/tree/v3.7.0/keras/src/distribution/distribution_lib.py#L829">[source]</a></span></p> <h3 id="setdistribution-function"><code>set_distribution</code> function</h3> <div class="codehilite"><pre><span></span><code><span class="n">keras</span><span class="o">.</span><span class="n">distribution</span><span class="o">.</span><span class="n">set_distribution</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> </code></pre></div> <p>Set the distribution as the global distribution setting.</p> <p><strong>Arguments</strong></p> <ul> <li><strong>value</strong>: a <code>Distribution</code> instance.</li> </ul> <hr /> <p><span style="float:right;"><a href="https://github.com/keras-team/keras/tree/v3.7.0/keras/src/distribution/distribution_lib.py#L823">[source]</a></span></p> <h3 id="distribution-function"><code>distribution</code> function</h3> <div class="codehilite"><pre><span></span><code><span class="n">keras</span><span class="o">.</span><span class="n">distribution</span><span class="o">.</span><span class="n">distribution</span><span class="p">()</span> </code></pre></div> <p>Retrieve the current distribution from global context.</p> <hr /> <p><span style="float:right;"><a href="https://github.com/keras-team/keras/tree/v3.7.0/keras/src/distribution/distribution_lib.py#L26">[source]</a></span></p> <h3 id="listdevices-function"><code>list_devices</code> function</h3> <div class="codehilite"><pre><span></span><code><span class="n">keras</span><span class="o">.</span><span class="n">distribution</span><span class="o">.</span><span class="n">list_devices</span><span class="p">(</span><span class="n">device_type</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span> </code></pre></div> <p>Return all the available devices based on the device type.</p> <p>Note: in a distributed setting, global devices are returned.</p> <p><strong>Arguments</strong></p> <ul> <li><strong>device_type</strong>: string, one of <code>"cpu"</code>, <code>"gpu"</code> or <code>"tpu"</code>. Defaults to <code>"gpu"</code> or <code>"tpu"</code> if available when <code>device_type</code> is not provided. Otherwise will return the <code>"cpu"</code> devices.</li> </ul> <p>Return: List of devices that are available for distribute computation.</p> <hr /> <p><span style="float:right;"><a href="https://github.com/keras-team/keras/tree/v3.7.0/keras/src/distribution/distribution_lib.py#L44">[source]</a></span></p> <h3 id="initialize-function"><code>initialize</code> function</h3> <div class="codehilite"><pre><span></span><code><span class="n">keras</span><span class="o">.</span><span class="n">distribution</span><span class="o">.</span><span class="n">initialize</span><span class="p">(</span> <span class="n">job_addresses</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_processes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">process_id</span><span class="o">=</span><span class="kc">None</span> <span class="p">)</span> </code></pre></div> <p>Initialize the distribution system for multi-host/process setting.</p> <p>Calling <code>initialize</code> will prepare the backend for execution on multi-host GPU or TPUs. It should be called before any computations.</p> <p>Note that the parameters can also be injected via environment variables, which can be better controlled by the launch script at startup time. For certain backend that also rely on the environment variables to configure, Keras will properly forward them.</p> <p><strong>Arguments</strong></p> <ul> <li><strong>job_addresses</strong>: string. Comma separated IP addresses for all the jobs that will form the whole computation cluster. Note that for JAX backend, only the address for job 0 (coodinator) is needed. For certain runtime like cloud TPU, this value can be <code>None</code>, and the backend will figure it out with the TPU environment variables. You can also config this value via environment variable <code>KERAS_DISTRIBUTION_JOB_ADDRESSES</code>.</li> <li><strong>num_processes</strong>: int. The number of worker/processes that will form the whole computation cluster. For certain runtime like cloud TPU, this value can be <code>None</code>, and the backend will figure it out with the TPU environment variables. You can also configure this value via environment variable <code>KERAS_DISTRIBUTION_NUM_PROCESSES</code>.</li> <li><strong>process_id</strong>: int. The ID number of the current worker/process. The value should be ranged from <code>0</code> to <code>num_processes - 1</code>. <code>0</code> will indicate the current worker/process is the master/coordinate job. You can also configure this value via environment variable <code>KERAS_DISTRIBUTION_PROCESS_ID</code>.</li> </ul> <h1 id="example">Example</h1> <p>Suppose there are two GPU processes, and process 0 is running at address <code>10.0.0.1:1234</code>, and process 1 is running at address <code>10.0.0.2:2345</code>. To configure such cluster, you can run - __ On process 0__:</p> <div class="codehilite"><pre><span></span><code><span class="n">keras</span><span class="o">.</span><span class="n">distribute</span><span class="o">.</span><span class="n">initialize</span><span class="p">(</span> <span class="n">job_addresses</span><span class="o">=</span><span class="s2">"10.0.0.1:1234,10.0.0.2:2345"</span><span class="p">,</span> <span class="n">num_processes</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">process_id</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> </code></pre></div> <ul> <li>__ On process 1__:</li> </ul> <div class="codehilite"><pre><span></span><code><span class="n">keras</span><span class="o">.</span><span class="n">distribute</span><span class="o">.</span><span class="n">initialize</span><span class="p">(</span> <span class="n">job_addresses</span><span class="o">=</span><span class="s2">"10.0.0.1:1234,10.0.0.2:2345"</span><span class="p">,</span> <span class="n">num_processes</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">process_id</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span> </code></pre></div> <ul> <li>__ or via the environment variables__:</li> <li><strong>On process 0</strong>:</li> </ul> <div class="codehilite"><pre><span></span><code><span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span> <span class="s2">"KERAS_DISTRIBUTION_JOB_ADDRESSES"</span><span class="p">]</span> <span class="o">=</span> <span class="s2">"10.0.0.1:1234,10.0.0.2:2345"</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s2">"KERAS_DISTRIBUTION_NUM_PROCESSES"</span><span class="p">]</span> <span class="o">=</span> <span class="s2">"2"</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s2">"KERAS_DISTRIBUTION_PROCESS_ID"</span><span class="p">]</span> <span class="o">=</span> <span class="s2">"0"</span> <span class="n">keras</span><span class="o">.</span><span class="n">distribute</span><span class="o">.</span><span class="n">initialize</span><span class="p">()</span> </code></pre></div> <ul> <li>__ On process 1__:</li> </ul> <div class="codehilite"><pre><span></span><code><span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span> <span class="s2">"KERAS_DISTRIBUTION_JOB_ADDRESSES"</span><span class="p">]</span> <span class="o">=</span> <span class="s2">"10.0.0.1:1234,10.0.0.2:2345"</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s2">"KERAS_DISTRIBUTION_NUM_PROCESSES"</span><span class="p">]</span> <span class="o">=</span> <span class="s2">"2"</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s2">"KERAS_DISTRIBUTION_PROCESS_ID"</span><span class="p">]</span> <span class="o">=</span> <span class="s2">"1"</span> <span class="n">keras</span><span class="o">.</span><span class="n">distribute</span><span class="o">.</span><span class="n">initialize</span><span class="p">()</span> </code></pre></div> <p>Also note that for JAX backend, the <code>job_addresses</code> can be further reduced to just the master/coordinator address, which is - __<code>10.0.0.1__:1234</code>.</p> <hr /> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#distribution-utilities'>Distribution utilities</a> </div> <div class='k-outline-depth-3'> <a href='#setdistribution-function'><code>set_distribution</code> function</a> </div> <div class='k-outline-depth-3'> <a href='#distribution-function'><code>distribution</code> function</a> </div> <div class='k-outline-depth-3'> <a href='#listdevices-function'><code>list_devices</code> function</a> </div> <div class='k-outline-depth-3'> <a href='#initialize-function'><code>initialize</code> function</a> </div> <div class='k-outline-depth-1'> <a href='#example'>Example</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>

CINXE.COM

Distribution utilities