@@ -1,4 +1,4 @@
 
                             # Sphinx build info version 1
                
 
                             # This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done.
                
 
                            -config: 1ceb2b6b1de4e767037aa12749cd16be
                
 
                            +config: 338ace02139fbf63879758a90ee35818
                
 
                             tags: 645f666f9bcd5a90fca523b33c5a78b7
                
 
            
          
@@ -1,13 +1,13 @@
 
                             <!DOCTYPE html>
                
 
                             <html class="writer-html5" lang="en" >
                
 
                             <head>
                
 
                            -  <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
                
 
                            -
                
 
                            +  <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>Contribution Guidelines &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>Contribution Guidelines &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -15,7 +15,9 @@
 
                                     <script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
                
 
                                     <script src="_static/jquery.js"></script>
                
 
                                     <script src="_static/underscore.js"></script>
                
 
                            +        <script src="_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="_static/doctools.js"></script>
                
 
                            +        <script src="_static/sphinx_highlight.js"></script>
                
 
                                 <script src="_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="genindex.html" />
                
 
                                 <link rel="search" title="Search" href="search.html" /> 
                
@@ -36,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -86,8 +87,8 @@
 
                                       <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
                
 
                                        <div itemprop="articleBody">
                
 
                            -  <section id="contribution-guidelines">
                
 
                            -<h1>Contribution Guidelines<a class="headerlink" href="#contribution-guidelines" title="Permalink to this headline"></a></h1>
                
 
                            +  <div class="section" id="contribution-guidelines">
                
 
                            +<h1>Contribution Guidelines<a class="headerlink" href="#contribution-guidelines" title="Permalink to this heading"></a></h1>
                
 
                             <p>We appreciate all contributions. If you are planning to contribute back bug-fixes, please do so without any further discussion. If you plan to contribute new features, utility functions or extensions, please first open an issue and discuss the feature with us.</p>
                
 
                             <p>Here are a few more things to know:</p>
                
 
                             <ul class="simple">
                
@@ -95,8 +96,8 @@
 
                             <li><p><span class="xref myst">Jupyter Notebooks Contribution</span></p></li>
                
 
                             <li><p><span class="xref myst">Code Style Guidelines</span></p></li>
                
 
                             </ul>
                
 
                            -<section id="how-to-contribute">
                
 
                            -<h2>How to Contribute<a class="headerlink" href="#how-to-contribute" title="Permalink to this headline"></a></h2>
                
 
                            +<div class="section" id="how-to-contribute">
                
 
                            +<h2>How to Contribute<a class="headerlink" href="#how-to-contribute" title="Permalink to this heading"></a></h2>
                
 
                             <p>Here is a simple guideline to get you started with your first contribution</p>
                
 
                             <ol class="arabic simple">
                
 
                             <li><p>Use <a class="reference external" href="https://github.com/Deci-AI/super-gradients/issues">issues</a> to discuss the suggested changes. Create an issue describing changes if necessary and add labels to ease orientation.</p></li>
                
@@ -130,9 +131,9 @@
 
                             <ol class="arabic simple" start="7">
                
 
                             <li><p>Create a pull request against <b>master</b> branch.</p></li>
                
 
                             </ol>
                
 
                            -</section>
                
 
                            -<section id="jupyter-notebooks-contribution">
                
 
                            -<h2>Jupyter Notebooks Contribution<a class="headerlink" href="#jupyter-notebooks-contribution" title="Permalink to this headline"></a></h2>
                
 
                            +</div>
                
 
                            +<div class="section" id="jupyter-notebooks-contribution">
                
 
                            +<h2>Jupyter Notebooks Contribution<a class="headerlink" href="#jupyter-notebooks-contribution" title="Permalink to this heading"></a></h2>
                
 
                             <p>Pulling updates from remote might cause merge conflicts with jupyter notebooks. The tool <a class="reference external" href="https://nbdime.readthedocs.io/en/latest/">nbdime</a> might solve this.</p>
                
 
                             <ul class="simple">
                
 
                             <li><p>Installing nbdime</p></li>
                
@@ -146,9 +147,9 @@
 
                             <div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">nbdiff</span> <span class="n">notebook_1</span><span class="o">.</span><span class="n">ipynb</span> <span class="n">notebook_2</span><span class="o">.</span><span class="n">ipynb</span>
                
 
                             </pre></div>
                
 
                             </div>
                
 
                            -</section>
                
 
                            -<section id="code-style-guidelines">
                
 
                            -<h2>Code Style Guidelines<a class="headerlink" href="#code-style-guidelines" title="Permalink to this headline"></a></h2>
                
 
                            +</div>
                
 
                            +<div class="section" id="code-style-guidelines">
                
 
                            +<h2>Code Style Guidelines<a class="headerlink" href="#code-style-guidelines" title="Permalink to this heading"></a></h2>
                
 
                             <p>We are working hard to make sure all the code in this repository is readable, maintainable and testable.
                
 
                             We follow the Google docstring guidelines outlined on this <a class="reference external" href="https://github.com/google/styleguide/blob/gh-pages/pyguide.md#38-comments-and-docstrings">styleguide</a> page. For example:</p>
                
 
                             <div class="highlight-python notranslate"><div class="highlight"><pre><span></span>  <span class="k">def</span> <span class="nf">python_function</span><span class="p">(</span><span class="n">first_argument</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">second_argument</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
                
@@ -167,9 +168,9 @@ We follow the Google docstring guidelines outlined on this <a class="reference e
 
                             <span class="sd">      &quot;&quot;&quot;</span>
                
 
                             </pre></div>
                
 
                             </div>
                
 
                            -</section>
                
 
                            -<section id="documentation">
                
 
                            -<h2>Documentation<a class="headerlink" href="#documentation" title="Permalink to this headline"></a></h2>
                
 
                            +</div>
                
 
                            +<div class="section" id="documentation">
                
 
                            +<h2>Documentation<a class="headerlink" href="#documentation" title="Permalink to this heading"></a></h2>
                
 
                             <p>We use  GitHub Pages for technical documentation hosting on https://deci-ai.github.io/super-gradients/welcome.html <br>
                
 
                             To generate the docs based on the current work tree, run: <br>
                
 
                             <code>./scripts/generate_docs.sh</code> <br><br>
                
@@ -177,8 +178,8 @@ And the documentation will automatically update, based on <code>documentation/</
 
                             The new documentation HTML will be generated to <code>docs/</code>. <br>
                
 
                             Once <code>docs/</code> is committed and pushed, GitHub Pages will use it.<br>
                
 
                             The step of documentation update is currently manual.</p>
                
 
                            -</section>
                
 
                            -</section>
                
 
                            +</div>
                
 
                            +</div>
                
 
                                        </div>
                
@@ -208,4 +209,4 @@ The step of documentation update is currently manual.</p>
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -1,13 +1,13 @@
 
                             <!DOCTYPE html>
                
 
                             <html class="writer-html5" lang="en" >
                
 
                             <head>
                
 
                            -  <meta charset="utf-8" /><meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
                
 
                            -
                
 
                            +  <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>&lt;no title&gt; &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>&lt;no title&gt; &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -15,7 +15,9 @@
 
                                     <script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
                
 
                                     <script src="_static/jquery.js"></script>
                
 
                                     <script src="_static/underscore.js"></script>
                
 
                            +        <script src="_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="_static/doctools.js"></script>
                
 
                            +        <script src="_static/sphinx_highlight.js"></script>
                
 
                                 <script src="_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="genindex.html" />
                
 
                                 <link rel="search" title="Search" href="search.html" /> 
                
@@ -36,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -294,4 +295,4 @@ limitations under the License.</p>
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>Overview: module code &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>Overview: module code &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../" id="documentation_options" src="../_static/documentation_options.js"></script>
                
 
                                     <script src="../_static/jquery.js"></script>
                
 
                                     <script src="../_static/underscore.js"></script>
                
 
                            +        <script src="../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../_static/doctools.js"></script>
                
 
                            +        <script src="../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -85,10 +87,10 @@
 
                                        <div itemprop="articleBody">
                
 
                               <h1>All modules for which code is available</h1>
                
 
                            -<ul><li><a href="super_gradients/common/abstractions/abstract_logger.html">super_gradients.common.abstractions.abstract_logger</a></li>
                
 
                            -<li><a href="super_gradients/common/auto_logging/auto_logger.html">super_gradients.common.auto_logging.auto_logger</a></li>
                
 
                            +<ul><li><a href="super_gradients/common/auto_logging/auto_logger.html">super_gradients.common.auto_logging.auto_logger</a></li>
                
 
                            +<li><a href="super_gradients/common/auto_logging/console_logging.html">super_gradients.common.auto_logging.console_logging</a></li>
                
 
                             <li><a href="super_gradients/common/aws_connection/aws_connector.html">super_gradients.common.aws_connection.aws_connector</a></li>
                
 
                            -<li><a href="super_gradients/common/aws_connection/aws_secrets_manager_connector.html">super_gradients.common.aws_connection.aws_secrets_manager_connector</a></li>
                
 
                            +<li><a href="super_gradients/common/crash_handler/crash_handler.html">super_gradients.common.crash_handler.crash_handler</a></li>
                
 
                             <li><a href="super_gradients/common/data_connection/s3_connector.html">super_gradients.common.data_connection.s3_connector</a></li>
                
 
                             <li><a href="super_gradients/common/data_interface/adnn_model_repository_data_interface.html">super_gradients.common.data_interface.adnn_model_repository_data_interface</a></li>
                
 
                             <li><a href="super_gradients/common/data_interface/dataset_data_interface.html">super_gradients.common.data_interface.dataset_data_interface</a></li>
                
@@ -96,37 +98,29 @@
 
                             <li><a href="super_gradients/common/data_types/enum/evaluation_type.html">super_gradients.common.data_types.enum.evaluation_type</a></li>
                
 
                             <li><a href="super_gradients/common/data_types/enum/multi_gpu_mode.html">super_gradients.common.data_types.enum.multi_gpu_mode</a></li>
                
 
                             <li><a href="super_gradients/common/data_types/enum/strict_load.html">super_gradients.common.data_types.enum.strict_load</a></li>
                
 
                            -<li><a href="super_gradients/common/decorators/deci_logger.html">super_gradients.common.decorators.deci_logger</a></li>
                
 
                            +<li><a href="super_gradients/common/data_types/enum/upsample_mode.html">super_gradients.common.data_types.enum.upsample_mode</a></li>
                
 
                             <li><a href="super_gradients/common/decorators/explicit_params_validator.html">super_gradients.common.decorators.explicit_params_validator</a></li>
                
 
                             <li><a href="super_gradients/common/decorators/singleton.html">super_gradients.common.decorators.singleton</a></li>
                
 
                             <li><a href="super_gradients/common/environment/env_helpers.html">super_gradients.common.environment.env_helpers</a></li>
                
 
                            -<li><a href="super_gradients/training/datasets/all_datasets.html">super_gradients.training.datasets.all_datasets</a></li>
                
 
                            -<li><a href="super_gradients/training/datasets/auto_augment.html">super_gradients.training.datasets.auto_augment</a></li>
                
 
                            +<li><a href="super_gradients/common/object_names.html">super_gradients.common.object_names</a></li>
                
 
                            +<li><a href="super_gradients/training/dataloaders/dataloaders.html">super_gradients.training.dataloaders.dataloaders</a></li>
                
 
                            +<li><a href="super_gradients/training/datasets/classification_datasets/cifar.html">super_gradients.training.datasets.classification_datasets.cifar</a></li>
                
 
                            +<li><a href="super_gradients/training/datasets/classification_datasets/imagenet_dataset.html">super_gradients.training.datasets.classification_datasets.imagenet_dataset</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/data_augmentation.html">super_gradients.training.datasets.data_augmentation</a></li>
                
 
                            -<li><a href="super_gradients/training/datasets/dataset_interfaces/dataset_interface.html">super_gradients.training.datasets.dataset_interfaces.dataset_interface</a></li>
                
 
                            -<li><a href="super_gradients/training/datasets/datasets_utils.html">super_gradients.training.datasets.datasets_utils</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/detection_datasets/coco_detection.html">super_gradients.training.datasets.detection_datasets.coco_detection</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/detection_datasets/detection_dataset.html">super_gradients.training.datasets.detection_datasets.detection_dataset</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/detection_datasets/pascal_voc_detection.html">super_gradients.training.datasets.detection_datasets.pascal_voc_detection</a></li>
                
 
                            -<li><a href="super_gradients/training/datasets/mixup.html">super_gradients.training.datasets.mixup</a></li>
                
 
                            -<li><a href="super_gradients/training/datasets/segmentation_datasets/cityscape_segmentation.html">super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/segmentation_datasets/coco_segmentation.html">super_gradients.training.datasets.segmentation_datasets.coco_segmentation</a></li>
                
 
                            -<li><a href="super_gradients/training/datasets/segmentation_datasets/pascal_aug_segmentation.html">super_gradients.training.datasets.segmentation_datasets.pascal_aug_segmentation</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/segmentation_datasets/pascal_voc_segmentation.html">super_gradients.training.datasets.segmentation_datasets.pascal_voc_segmentation</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/segmentation_datasets/segmentation_dataset.html">super_gradients.training.datasets.segmentation_datasets.segmentation_dataset</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/segmentation_datasets/supervisely_persons_segmentation.html">super_gradients.training.datasets.segmentation_datasets.supervisely_persons_segmentation</a></li>
                
 
                             <li><a href="super_gradients/training/datasets/sg_dataset.html">super_gradients.training.datasets.sg_dataset</a></li>
                
 
                            -<li><a href="super_gradients/training/exceptions/dataset_exceptions.html">super_gradients.training.exceptions.dataset_exceptions</a></li>
                
 
                            -<li><a href="super_gradients/training/exceptions/sg_model_exceptions.html">super_gradients.training.exceptions.sg_model_exceptions</a></li>
                
 
                            -<li><a href="super_gradients/training/kd_model/kd_model.html">super_gradients.training.kd_model.kd_model</a></li>
                
 
                            -<li><a href="super_gradients/training/legacy/utils.html">super_gradients.training.legacy.utils</a></li>
                
 
                            +<li><a href="super_gradients/training/kd_trainer/kd_trainer.html">super_gradients.training.kd_trainer.kd_trainer</a></li>
                
 
                             <li><a href="super_gradients/training/losses/bce_dice_loss.html">super_gradients.training.losses.bce_dice_loss</a></li>
                
 
                            -<li><a href="super_gradients/training/losses/ddrnet_loss.html">super_gradients.training.losses.ddrnet_loss</a></li>
                
 
                             <li><a href="super_gradients/training/losses/dice_ce_edge_loss.html">super_gradients.training.losses.dice_ce_edge_loss</a></li>
                
 
                             <li><a href="super_gradients/training/losses/focal_loss.html">super_gradients.training.losses.focal_loss</a></li>
                
 
                             <li><a href="super_gradients/training/losses/kd_losses.html">super_gradients.training.losses.kd_losses</a></li>
                
 
                             <li><a href="super_gradients/training/losses/label_smoothing_cross_entropy_loss.html">super_gradients.training.losses.label_smoothing_cross_entropy_loss</a></li>
                
 
                            -<li><a href="super_gradients/training/losses/ohem_ce_loss.html">super_gradients.training.losses.ohem_ce_loss</a></li>
                
 
                             <li><a href="super_gradients/training/losses/r_squared_loss.html">super_gradients.training.losses.r_squared_loss</a></li>
                
 
                             <li><a href="super_gradients/training/losses/shelfnet_ohem_loss.html">super_gradients.training.losses.shelfnet_ohem_loss</a></li>
                
 
                             <li><a href="super_gradients/training/losses/shelfnet_semantic_encoding_loss.html">super_gradients.training.losses.shelfnet_semantic_encoding_loss</a></li>
                
@@ -134,26 +128,13 @@
 
                             <li><a href="super_gradients/training/losses/yolox_loss.html">super_gradients.training.losses.yolox_loss</a></li>
                
 
                             <li><a href="super_gradients/training/metrics/classification_metrics.html">super_gradients.training.metrics.classification_metrics</a></li>
                
 
                             <li><a href="super_gradients/training/metrics/detection_metrics.html">super_gradients.training.metrics.detection_metrics</a></li>
                
 
                            -<li><a href="super_gradients/training/metrics/metric_utils.html">super_gradients.training.metrics.metric_utils</a></li>
                
 
                             <li><a href="super_gradients/training/metrics/segmentation_metrics.html">super_gradients.training.metrics.segmentation_metrics</a></li>
                
 
                            -<li><a href="super_gradients/training/models/sg_module.html">super_gradients.training.models.sg_module</a></li>
                
 
                            -<li><a href="super_gradients/training/sg_model/sg_model.html">super_gradients.training.sg_model.sg_model</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/callbacks.html">super_gradients.training.utils.callbacks</a></li>
                
 
                            +<li><a href="super_gradients/training/sg_trainer/sg_trainer.html">super_gradients.training.sg_trainer.sg_trainer</a></li>
                
 
                            +<li><a href="super_gradients/training/training_hyperparams/training_hyperparams.html">super_gradients.training.training_hyperparams.training_hyperparams</a></li>
                
 
                            +<li><a href="super_gradients/training/transforms/transforms.html">super_gradients.training.transforms.transforms</a></li>
                
 
                             <li><a href="super_gradients/training/utils/checkpoint_utils.html">super_gradients.training.utils.checkpoint_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/detection_utils.html">super_gradients.training.utils.detection_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/distributed_training_utils.html">super_gradients.training.utils.distributed_training_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/early_stopping.html">super_gradients.training.utils.early_stopping</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/ema.html">super_gradients.training.utils.ema</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/export_utils.html">super_gradients.training.utils.export_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/module_utils.html">super_gradients.training.utils.module_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/optimizer_utils.html">super_gradients.training.utils.optimizer_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/optimizers/rmsprop_tf.html">super_gradients.training.utils.optimizers.rmsprop_tf</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/regularization_utils.html">super_gradients.training.utils.regularization_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/segmentation_utils.html">super_gradients.training.utils.segmentation_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/sg_model_utils.html">super_gradients.training.utils.sg_model_utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/ssd_utils.html">super_gradients.training.utils.ssd_utils</a></li>
                
 
                             <li><a href="super_gradients/training/utils/utils.html">super_gradients.training.utils.utils</a></li>
                
 
                            -<li><a href="super_gradients/training/utils/weight_averaging_utils.html">super_gradients.training.utils.weight_averaging_utils</a></li>
                
 
                            +<li><a href="super_gradients/training/utils/version_utils.html">super_gradients.training.utils.version_utils</a></li>
                
 
                             </ul>
                
 
                                        </div>
                
@@ -183,4 +164,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>logging &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../" id="documentation_options" src="../_static/documentation_options.js"></script>
        <script src="../_static/jquery.js"></script>
        <script src="../_static/underscore.js"></script>
        <script src="../_static/doctools.js"></script>
    <script src="../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../genindex.html" />
    <link rel="search" title="Search" href="../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../welcome.html">SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../super_gradients.common.html">Common</a></li>
<li class="toctree-l1"><a class="reference internal" href="../super_gradients.training.html">Training</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="index.html">Module code</a> &raquo;</li>
      <li>logging</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for logging</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2001-2019 by Vinay Sajip. All Rights Reserved.</span>
<span class="c1">#</span>
<span class="c1"># Permission to use, copy, modify, and distribute this software and its</span>
<span class="c1"># documentation for any purpose and without fee is hereby granted,</span>
<span class="c1"># provided that the above copyright notice appear in all copies and that</span>
<span class="c1"># both that copyright notice and this permission notice appear in</span>
<span class="c1"># supporting documentation, and that the name of Vinay Sajip</span>
<span class="c1"># not be used in advertising or publicity pertaining to distribution</span>
<span class="c1"># of the software without specific, written prior permission.</span>
<span class="c1"># VINAY SAJIP DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING</span>
<span class="c1"># ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL</span>
<span class="c1"># VINAY SAJIP BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR</span>
<span class="c1"># ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER</span>
<span class="c1"># IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT</span>
<span class="c1"># OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.</span>

<span class="sd">&quot;&quot;&quot;</span>
<span class="sd">Logging package for Python. Based on PEP 282 and comments thereto in</span>
<span class="sd">comp.lang.python.</span>

<span class="sd">Copyright (C) 2001-2019 Vinay Sajip. All Rights Reserved.</span>

<span class="sd">To use, simply &#39;import logging&#39; and log away!</span>
<span class="sd">&quot;&quot;&quot;</span>

<span class="kn">import</span> <span class="nn">sys</span><span class="o">,</span> <span class="nn">os</span><span class="o">,</span> <span class="nn">time</span><span class="o">,</span> <span class="nn">io</span><span class="o">,</span> <span class="nn">re</span><span class="o">,</span> <span class="nn">traceback</span><span class="o">,</span> <span class="nn">warnings</span><span class="o">,</span> <span class="nn">weakref</span><span class="o">,</span> <span class="nn">collections.abc</span>

<span class="kn">from</span> <span class="nn">string</span> <span class="kn">import</span> <span class="n">Template</span>
<span class="kn">from</span> <span class="nn">string</span> <span class="kn">import</span> <span class="n">Formatter</span> <span class="k">as</span> <span class="n">StrFormatter</span>


<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;BASIC_FORMAT&#39;</span><span class="p">,</span> <span class="s1">&#39;BufferingFormatter&#39;</span><span class="p">,</span> <span class="s1">&#39;CRITICAL&#39;</span><span class="p">,</span> <span class="s1">&#39;DEBUG&#39;</span><span class="p">,</span> <span class="s1">&#39;ERROR&#39;</span><span class="p">,</span>
           <span class="s1">&#39;FATAL&#39;</span><span class="p">,</span> <span class="s1">&#39;FileHandler&#39;</span><span class="p">,</span> <span class="s1">&#39;Filter&#39;</span><span class="p">,</span> <span class="s1">&#39;Formatter&#39;</span><span class="p">,</span> <span class="s1">&#39;Handler&#39;</span><span class="p">,</span> <span class="s1">&#39;INFO&#39;</span><span class="p">,</span>
           <span class="s1">&#39;LogRecord&#39;</span><span class="p">,</span> <span class="s1">&#39;Logger&#39;</span><span class="p">,</span> <span class="s1">&#39;LoggerAdapter&#39;</span><span class="p">,</span> <span class="s1">&#39;NOTSET&#39;</span><span class="p">,</span> <span class="s1">&#39;NullHandler&#39;</span><span class="p">,</span>
           <span class="s1">&#39;StreamHandler&#39;</span><span class="p">,</span> <span class="s1">&#39;WARN&#39;</span><span class="p">,</span> <span class="s1">&#39;WARNING&#39;</span><span class="p">,</span> <span class="s1">&#39;addLevelName&#39;</span><span class="p">,</span> <span class="s1">&#39;basicConfig&#39;</span><span class="p">,</span>
           <span class="s1">&#39;captureWarnings&#39;</span><span class="p">,</span> <span class="s1">&#39;critical&#39;</span><span class="p">,</span> <span class="s1">&#39;debug&#39;</span><span class="p">,</span> <span class="s1">&#39;disable&#39;</span><span class="p">,</span> <span class="s1">&#39;error&#39;</span><span class="p">,</span>
           <span class="s1">&#39;exception&#39;</span><span class="p">,</span> <span class="s1">&#39;fatal&#39;</span><span class="p">,</span> <span class="s1">&#39;getLevelName&#39;</span><span class="p">,</span> <span class="s1">&#39;getLogger&#39;</span><span class="p">,</span> <span class="s1">&#39;getLoggerClass&#39;</span><span class="p">,</span>
           <span class="s1">&#39;info&#39;</span><span class="p">,</span> <span class="s1">&#39;log&#39;</span><span class="p">,</span> <span class="s1">&#39;makeLogRecord&#39;</span><span class="p">,</span> <span class="s1">&#39;setLoggerClass&#39;</span><span class="p">,</span> <span class="s1">&#39;shutdown&#39;</span><span class="p">,</span>
           <span class="s1">&#39;warn&#39;</span><span class="p">,</span> <span class="s1">&#39;warning&#39;</span><span class="p">,</span> <span class="s1">&#39;getLogRecordFactory&#39;</span><span class="p">,</span> <span class="s1">&#39;setLogRecordFactory&#39;</span><span class="p">,</span>
           <span class="s1">&#39;lastResort&#39;</span><span class="p">,</span> <span class="s1">&#39;raiseExceptions&#39;</span><span class="p">]</span>

<span class="kn">import</span> <span class="nn">threading</span>

<span class="n">__author__</span>  <span class="o">=</span> <span class="s2">&quot;Vinay Sajip &lt;vinay_sajip@red-dove.com&gt;&quot;</span>
<span class="n">__status__</span>  <span class="o">=</span> <span class="s2">&quot;production&quot;</span>
<span class="c1"># The following module attributes are no longer updated.</span>
<span class="n">__version__</span> <span class="o">=</span> <span class="s2">&quot;0.5.1.2&quot;</span>
<span class="n">__date__</span>    <span class="o">=</span> <span class="s2">&quot;07 February 2010&quot;</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Miscellaneous module data</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="c1">#</span>
<span class="c1">#_startTime is used as the base when calculating the relative time of events</span>
<span class="c1">#</span>
<span class="n">_startTime</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span>

<span class="c1">#</span>
<span class="c1">#raiseExceptions is used to see if exceptions during handling should be</span>
<span class="c1">#propagated</span>
<span class="c1">#</span>
<span class="n">raiseExceptions</span> <span class="o">=</span> <span class="kc">True</span>

<span class="c1">#</span>
<span class="c1"># If you don&#39;t want threading information in the log, set this to zero</span>
<span class="c1">#</span>
<span class="n">logThreads</span> <span class="o">=</span> <span class="kc">True</span>

<span class="c1">#</span>
<span class="c1"># If you don&#39;t want multiprocessing information in the log, set this to zero</span>
<span class="c1">#</span>
<span class="n">logMultiprocessing</span> <span class="o">=</span> <span class="kc">True</span>

<span class="c1">#</span>
<span class="c1"># If you don&#39;t want process information in the log, set this to zero</span>
<span class="c1">#</span>
<span class="n">logProcesses</span> <span class="o">=</span> <span class="kc">True</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Level related stuff</span>
<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#</span>
<span class="c1"># Default levels and level names, these can be replaced with any positive set</span>
<span class="c1"># of values having corresponding names. There is a pseudo-level, NOTSET, which</span>
<span class="c1"># is only really there as a lower limit for user-defined levels. Handlers and</span>
<span class="c1"># loggers are initialized with NOTSET so that they will log all messages, even</span>
<span class="c1"># at user-defined levels.</span>
<span class="c1">#</span>

<span class="n">CRITICAL</span> <span class="o">=</span> <span class="mi">50</span>
<span class="n">FATAL</span> <span class="o">=</span> <span class="n">CRITICAL</span>
<span class="n">ERROR</span> <span class="o">=</span> <span class="mi">40</span>
<span class="n">WARNING</span> <span class="o">=</span> <span class="mi">30</span>
<span class="n">WARN</span> <span class="o">=</span> <span class="n">WARNING</span>
<span class="n">INFO</span> <span class="o">=</span> <span class="mi">20</span>
<span class="n">DEBUG</span> <span class="o">=</span> <span class="mi">10</span>
<span class="n">NOTSET</span> <span class="o">=</span> <span class="mi">0</span>

<span class="n">_levelToName</span> <span class="o">=</span> <span class="p">{</span>
    <span class="n">CRITICAL</span><span class="p">:</span> <span class="s1">&#39;CRITICAL&#39;</span><span class="p">,</span>
    <span class="n">ERROR</span><span class="p">:</span> <span class="s1">&#39;ERROR&#39;</span><span class="p">,</span>
    <span class="n">WARNING</span><span class="p">:</span> <span class="s1">&#39;WARNING&#39;</span><span class="p">,</span>
    <span class="n">INFO</span><span class="p">:</span> <span class="s1">&#39;INFO&#39;</span><span class="p">,</span>
    <span class="n">DEBUG</span><span class="p">:</span> <span class="s1">&#39;DEBUG&#39;</span><span class="p">,</span>
    <span class="n">NOTSET</span><span class="p">:</span> <span class="s1">&#39;NOTSET&#39;</span><span class="p">,</span>
<span class="p">}</span>
<span class="n">_nameToLevel</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s1">&#39;CRITICAL&#39;</span><span class="p">:</span> <span class="n">CRITICAL</span><span class="p">,</span>
    <span class="s1">&#39;FATAL&#39;</span><span class="p">:</span> <span class="n">FATAL</span><span class="p">,</span>
    <span class="s1">&#39;ERROR&#39;</span><span class="p">:</span> <span class="n">ERROR</span><span class="p">,</span>
    <span class="s1">&#39;WARN&#39;</span><span class="p">:</span> <span class="n">WARNING</span><span class="p">,</span>
    <span class="s1">&#39;WARNING&#39;</span><span class="p">:</span> <span class="n">WARNING</span><span class="p">,</span>
    <span class="s1">&#39;INFO&#39;</span><span class="p">:</span> <span class="n">INFO</span><span class="p">,</span>
    <span class="s1">&#39;DEBUG&#39;</span><span class="p">:</span> <span class="n">DEBUG</span><span class="p">,</span>
    <span class="s1">&#39;NOTSET&#39;</span><span class="p">:</span> <span class="n">NOTSET</span><span class="p">,</span>
<span class="p">}</span>

<span class="k">def</span> <span class="nf">getLevelName</span><span class="p">(</span><span class="n">level</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Return the textual or numeric representation of logging level &#39;level&#39;.</span>

<span class="sd">    If the level is one of the predefined levels (CRITICAL, ERROR, WARNING,</span>
<span class="sd">    INFO, DEBUG) then you get the corresponding string. If you have</span>
<span class="sd">    associated levels with names using addLevelName then the name you have</span>
<span class="sd">    associated with &#39;level&#39; is returned.</span>

<span class="sd">    If a numeric value corresponding to one of the defined levels is passed</span>
<span class="sd">    in, the corresponding string representation is returned.</span>

<span class="sd">    If a string representation of the level is passed in, the corresponding</span>
<span class="sd">    numeric value is returned.</span>

<span class="sd">    If no matching numeric or string value is passed in, the string</span>
<span class="sd">    &#39;Level %s&#39; % level is returned.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># See Issues #22386, #27937 and #29220 for why it&#39;s this way</span>
    <span class="n">result</span> <span class="o">=</span> <span class="n">_levelToName</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">result</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">result</span>
    <span class="n">result</span> <span class="o">=</span> <span class="n">_nameToLevel</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">result</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">result</span>
    <span class="k">return</span> <span class="s2">&quot;Level </span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">level</span>

<span class="k">def</span> <span class="nf">addLevelName</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">levelName</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Associate &#39;levelName&#39; with &#39;level&#39;.</span>

<span class="sd">    This is used when converting levels to text during message formatting.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_acquireLock</span><span class="p">()</span>
    <span class="k">try</span><span class="p">:</span>    <span class="c1">#unlikely to cause an exception, but you never know...</span>
        <span class="n">_levelToName</span><span class="p">[</span><span class="n">level</span><span class="p">]</span> <span class="o">=</span> <span class="n">levelName</span>
        <span class="n">_nameToLevel</span><span class="p">[</span><span class="n">levelName</span><span class="p">]</span> <span class="o">=</span> <span class="n">level</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="n">_releaseLock</span><span class="p">()</span>

<span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">sys</span><span class="p">,</span> <span class="s1">&#39;_getframe&#39;</span><span class="p">):</span>
    <span class="n">currentframe</span> <span class="o">=</span> <span class="k">lambda</span><span class="p">:</span> <span class="n">sys</span><span class="o">.</span><span class="n">_getframe</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span> <span class="c1">#pragma: no cover</span>
    <span class="k">def</span> <span class="nf">currentframe</span><span class="p">():</span>
        <span class="sd">&quot;&quot;&quot;Return the frame object for the caller&#39;s stack frame.&quot;&quot;&quot;</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">Exception</span>
        <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">sys</span><span class="o">.</span><span class="n">exc_info</span><span class="p">()[</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">tb_frame</span><span class="o">.</span><span class="n">f_back</span>

<span class="c1">#</span>
<span class="c1"># _srcfile is used when walking the stack to check when we&#39;ve got the first</span>
<span class="c1"># caller stack frame, by skipping frames whose filename is that of this</span>
<span class="c1"># module&#39;s source. It therefore should contain the filename of this module&#39;s</span>
<span class="c1"># source file.</span>
<span class="c1">#</span>
<span class="c1"># Ordinarily we would use __file__ for this, but frozen modules don&#39;t always</span>
<span class="c1"># have __file__ set, for some reason (see Issue #21736). Thus, we get the</span>
<span class="c1"># filename from a handy code object from a function defined in this module.</span>
<span class="c1"># (There&#39;s no particular reason for picking addLevelName.)</span>
<span class="c1">#</span>

<span class="n">_srcfile</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">normcase</span><span class="p">(</span><span class="n">addLevelName</span><span class="o">.</span><span class="vm">__code__</span><span class="o">.</span><span class="n">co_filename</span><span class="p">)</span>

<span class="c1"># _srcfile is only used in conjunction with sys._getframe().</span>
<span class="c1"># To provide compatibility with older versions of Python, set _srcfile</span>
<span class="c1"># to None if _getframe() is not available; this value will prevent</span>
<span class="c1"># findCaller() from being called. You can also do this if you want to avoid</span>
<span class="c1"># the overhead of fetching caller information, even when _getframe() is</span>
<span class="c1"># available.</span>
<span class="c1">#if not hasattr(sys, &#39;_getframe&#39;):</span>
<span class="c1">#    _srcfile = None</span>


<span class="k">def</span> <span class="nf">_checkLevel</span><span class="p">(</span><span class="n">level</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="n">level</span>
    <span class="k">elif</span> <span class="nb">str</span><span class="p">(</span><span class="n">level</span><span class="p">)</span> <span class="o">==</span> <span class="n">level</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">level</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">_nameToLevel</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Unknown level: </span><span class="si">%r</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">level</span><span class="p">)</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="n">_nameToLevel</span><span class="p">[</span><span class="n">level</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;Level not an integer or a valid string: </span><span class="si">%r</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">level</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">rv</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Thread-related stuff</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="c1">#</span>
<span class="c1">#_lock is used to serialize access to shared data structures in this module.</span>
<span class="c1">#This needs to be an RLock because fileConfig() creates and configures</span>
<span class="c1">#Handlers, and so might arbitrary user threads. Since Handler code updates the</span>
<span class="c1">#shared dictionary _handlers, it needs to acquire the lock. But if configuring,</span>
<span class="c1">#the lock would already have been acquired - so we need an RLock.</span>
<span class="c1">#The same argument applies to Loggers and Manager.loggerDict.</span>
<span class="c1">#</span>
<span class="n">_lock</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">RLock</span><span class="p">()</span>

<span class="k">def</span> <span class="nf">_acquireLock</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Acquire the module-level lock for serializing access to shared data.</span>

<span class="sd">    This should be released with _releaseLock().</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_lock</span><span class="p">:</span>
        <span class="n">_lock</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>

<span class="k">def</span> <span class="nf">_releaseLock</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Release the module-level lock acquired by calling _acquireLock().</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_lock</span><span class="p">:</span>
        <span class="n">_lock</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>


<span class="c1"># Prevent a held logging lock from blocking a child from logging.</span>

<span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">os</span><span class="p">,</span> <span class="s1">&#39;register_at_fork&#39;</span><span class="p">):</span>  <span class="c1"># Windows and friends.</span>
    <span class="k">def</span> <span class="nf">_register_at_fork_reinit_lock</span><span class="p">(</span><span class="n">instance</span><span class="p">):</span>
        <span class="k">pass</span>  <span class="c1"># no-op when os.register_at_fork does not exist.</span>
<span class="k">else</span><span class="p">:</span>
    <span class="c1"># A collection of instances with a _at_fork_reinit method (logging.Handler)</span>
    <span class="c1"># to be called in the child after forking.  The weakref avoids us keeping</span>
    <span class="c1"># discarded Handler instances alive.</span>
    <span class="n">_at_fork_reinit_lock_weakset</span> <span class="o">=</span> <span class="n">weakref</span><span class="o">.</span><span class="n">WeakSet</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_register_at_fork_reinit_lock</span><span class="p">(</span><span class="n">instance</span><span class="p">):</span>
        <span class="n">_acquireLock</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">_at_fork_reinit_lock_weakset</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">instance</span><span class="p">)</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">_releaseLock</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_after_at_fork_child_reinit_locks</span><span class="p">():</span>
        <span class="k">for</span> <span class="n">handler</span> <span class="ow">in</span> <span class="n">_at_fork_reinit_lock_weakset</span><span class="p">:</span>
            <span class="n">handler</span><span class="o">.</span><span class="n">_at_fork_reinit</span><span class="p">()</span>

        <span class="c1"># _acquireLock() was called in the parent before forking.</span>
        <span class="c1"># The lock is reinitialized to unlocked state.</span>
        <span class="n">_lock</span><span class="o">.</span><span class="n">_at_fork_reinit</span><span class="p">()</span>

    <span class="n">os</span><span class="o">.</span><span class="n">register_at_fork</span><span class="p">(</span><span class="n">before</span><span class="o">=</span><span class="n">_acquireLock</span><span class="p">,</span>
                        <span class="n">after_in_child</span><span class="o">=</span><span class="n">_after_at_fork_child_reinit_locks</span><span class="p">,</span>
                        <span class="n">after_in_parent</span><span class="o">=</span><span class="n">_releaseLock</span><span class="p">)</span>


<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   The logging record</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="k">class</span> <span class="nc">LogRecord</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A LogRecord instance represents an event being logged.</span>

<span class="sd">    LogRecord instances are created every time something is logged. They</span>
<span class="sd">    contain all the information pertinent to the event being logged. The</span>
<span class="sd">    main information passed in is in msg and args, which are combined</span>
<span class="sd">    using str(msg) % args to create the message field of the record. The</span>
<span class="sd">    record also includes information such as when the record was created,</span>
<span class="sd">    the source line where the logging call was made, and any exception</span>
<span class="sd">    information to be logged.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">pathname</span><span class="p">,</span> <span class="n">lineno</span><span class="p">,</span>
                 <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="p">,</span> <span class="n">func</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">sinfo</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize a logging record with interesting information.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ct</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">=</span> <span class="n">name</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">msg</span> <span class="o">=</span> <span class="n">msg</span>
        <span class="c1">#</span>
        <span class="c1"># The following statement allows passing of a dictionary as a sole</span>
        <span class="c1"># argument, so that you can do something like</span>
        <span class="c1">#  logging.debug(&quot;a %(a)d b %(b)s&quot;, {&#39;a&#39;:1, &#39;b&#39;:2})</span>
        <span class="c1"># Suggested by Stefan Behnel.</span>
        <span class="c1"># Note that without the test for args[0], we get a problem because</span>
        <span class="c1"># during formatting, we test to see if the arg is present using</span>
        <span class="c1"># &#39;if self.args:&#39;. If the event being logged is e.g. &#39;Value is %d&#39;</span>
        <span class="c1"># and if the passed arg fails &#39;if self.args:&#39; then no formatting</span>
        <span class="c1"># is done. For example, logger.warning(&#39;Value is %d&#39;, 0) would log</span>
        <span class="c1"># &#39;Value is %d&#39; instead of &#39;Value is 0&#39;.</span>
        <span class="c1"># For the use case of passing a dictionary, this should not be a</span>
        <span class="c1"># problem.</span>
        <span class="c1"># Issue #21172: a request was made to relax the isinstance check</span>
        <span class="c1"># to hasattr(args[0], &#39;__getitem__&#39;). However, the docs on string</span>
        <span class="c1"># formatting still seem to suggest a mapping object is required.</span>
        <span class="c1"># Thus, while not removing the isinstance check, it does now look</span>
        <span class="c1"># for collections.abc.Mapping rather than, as before, dict.</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">args</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">args</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">collections</span><span class="o">.</span><span class="n">abc</span><span class="o">.</span><span class="n">Mapping</span><span class="p">)</span>
            <span class="ow">and</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]):</span>
            <span class="n">args</span> <span class="o">=</span> <span class="n">args</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">args</span> <span class="o">=</span> <span class="n">args</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">levelname</span> <span class="o">=</span> <span class="n">getLevelName</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">levelno</span> <span class="o">=</span> <span class="n">level</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pathname</span> <span class="o">=</span> <span class="n">pathname</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">basename</span><span class="p">(</span><span class="n">pathname</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">module</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">splitext</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
        <span class="k">except</span> <span class="p">(</span><span class="ne">TypeError</span><span class="p">,</span> <span class="ne">ValueError</span><span class="p">,</span> <span class="ne">AttributeError</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="n">pathname</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">module</span> <span class="o">=</span> <span class="s2">&quot;Unknown module&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">exc_info</span> <span class="o">=</span> <span class="n">exc_info</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">exc_text</span> <span class="o">=</span> <span class="kc">None</span>      <span class="c1"># used to cache the traceback text</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stack_info</span> <span class="o">=</span> <span class="n">sinfo</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lineno</span> <span class="o">=</span> <span class="n">lineno</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">funcName</span> <span class="o">=</span> <span class="n">func</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">created</span> <span class="o">=</span> <span class="n">ct</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">msecs</span> <span class="o">=</span> <span class="p">(</span><span class="n">ct</span> <span class="o">-</span> <span class="nb">int</span><span class="p">(</span><span class="n">ct</span><span class="p">))</span> <span class="o">*</span> <span class="mi">1000</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">relativeCreated</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">created</span> <span class="o">-</span> <span class="n">_startTime</span><span class="p">)</span> <span class="o">*</span> <span class="mi">1000</span>
        <span class="k">if</span> <span class="n">logThreads</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">thread</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">get_ident</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">threadName</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">current_thread</span><span class="p">()</span><span class="o">.</span><span class="n">name</span>
        <span class="k">else</span><span class="p">:</span> <span class="c1"># pragma: no cover</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">thread</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">threadName</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">logMultiprocessing</span><span class="p">:</span> <span class="c1"># pragma: no cover</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">processName</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">processName</span> <span class="o">=</span> <span class="s1">&#39;MainProcess&#39;</span>
            <span class="n">mp</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">modules</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;multiprocessing&#39;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">mp</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="c1"># Errors may occur if multiprocessing has not finished loading</span>
                <span class="c1"># yet - e.g. if a custom import hook causes third-party code</span>
                <span class="c1"># to run when multiprocessing calls import. See issue 8200</span>
                <span class="c1"># for an example</span>
                <span class="k">try</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">processName</span> <span class="o">=</span> <span class="n">mp</span><span class="o">.</span><span class="n">current_process</span><span class="p">()</span><span class="o">.</span><span class="n">name</span>
                <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span> <span class="c1">#pragma: no cover</span>
                    <span class="k">pass</span>
        <span class="k">if</span> <span class="n">logProcesses</span> <span class="ow">and</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">os</span><span class="p">,</span> <span class="s1">&#39;getpid&#39;</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">process</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">getpid</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">process</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="s1">&#39;&lt;LogRecord: </span><span class="si">%s</span><span class="s1">, </span><span class="si">%s</span><span class="s1">, </span><span class="si">%s</span><span class="s1">, </span><span class="si">%s</span><span class="s1">, &quot;</span><span class="si">%s</span><span class="s1">&quot;&gt;&#39;</span><span class="o">%</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">levelno</span><span class="p">,</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">pathname</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">lineno</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">msg</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">getMessage</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the message for this LogRecord.</span>

<span class="sd">        Return the message for this LogRecord after merging any user-supplied</span>
<span class="sd">        arguments with the message.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">msg</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">msg</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">args</span><span class="p">:</span>
            <span class="n">msg</span> <span class="o">=</span> <span class="n">msg</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">args</span>
        <span class="k">return</span> <span class="n">msg</span>

<span class="c1">#</span>
<span class="c1">#   Determine which class to use when instantiating log records.</span>
<span class="c1">#</span>
<span class="n">_logRecordFactory</span> <span class="o">=</span> <span class="n">LogRecord</span>

<span class="k">def</span> <span class="nf">setLogRecordFactory</span><span class="p">(</span><span class="n">factory</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Set the factory to be used when instantiating a log record.</span>

<span class="sd">    :param factory: A callable which will be called to instantiate</span>
<span class="sd">    a log record.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">global</span> <span class="n">_logRecordFactory</span>
    <span class="n">_logRecordFactory</span> <span class="o">=</span> <span class="n">factory</span>

<span class="k">def</span> <span class="nf">getLogRecordFactory</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Return the factory to be used when instantiating a log record.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">return</span> <span class="n">_logRecordFactory</span>

<span class="k">def</span> <span class="nf">makeLogRecord</span><span class="p">(</span><span class="nb">dict</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Make a LogRecord whose attributes are defined by the specified dictionary,</span>
<span class="sd">    This function is useful for converting a logging event received over</span>
<span class="sd">    a socket connection (which is sent as a dictionary) into a LogRecord</span>
<span class="sd">    instance.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">rv</span> <span class="o">=</span> <span class="n">_logRecordFactory</span><span class="p">(</span><span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">,</span> <span class="p">(),</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
    <span class="n">rv</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="nb">dict</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">rv</span>


<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Formatter classes and functions</span>
<span class="c1">#---------------------------------------------------------------------------</span>
<span class="n">_str_formatter</span> <span class="o">=</span> <span class="n">StrFormatter</span><span class="p">()</span>
<span class="k">del</span> <span class="n">StrFormatter</span>


<span class="k">class</span> <span class="nc">PercentStyle</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>

    <span class="n">default_format</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">%(message)s</span><span class="s1">&#39;</span>
    <span class="n">asctime_format</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">%(asctime)s</span><span class="s1">&#39;</span>
    <span class="n">asctime_search</span> <span class="o">=</span> <span class="s1">&#39;%(asctime)&#39;</span>
    <span class="n">validation_pattern</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="sa">r</span><span class="s1">&#39;%\(\w+\)[#0+ -]*(\*|\d+)?(\.(\*|\d+))?[diouxefgcrsa%]&#39;</span><span class="p">,</span> <span class="n">re</span><span class="o">.</span><span class="n">I</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fmt</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span> <span class="o">=</span> <span class="n">fmt</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">default_format</span>

    <span class="k">def</span> <span class="nf">usesTime</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">asctime_search</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="mi">0</span>

    <span class="k">def</span> <span class="nf">validate</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Validate the input format, ensure it matches the correct style&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">validation_pattern</span><span class="o">.</span><span class="n">search</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid format &#39;</span><span class="si">%s</span><span class="s2">&#39; for &#39;</span><span class="si">%s</span><span class="s2">&#39; style&quot;</span> <span class="o">%</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">default_format</span><span class="p">[</span><span class="mi">0</span><span class="p">]))</span>

    <span class="k">def</span> <span class="nf">_format</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span> <span class="o">%</span> <span class="n">record</span><span class="o">.</span><span class="vm">__dict__</span>

    <span class="k">def</span> <span class="nf">format</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_format</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
        <span class="k">except</span> <span class="ne">KeyError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Formatting field not found in record: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">e</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">StrFormatStyle</span><span class="p">(</span><span class="n">PercentStyle</span><span class="p">):</span>
    <span class="n">default_format</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">{message}</span><span class="s1">&#39;</span>
    <span class="n">asctime_format</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">{asctime}</span><span class="s1">&#39;</span>
    <span class="n">asctime_search</span> <span class="o">=</span> <span class="s1">&#39;{asctime&#39;</span>

    <span class="n">fmt_spec</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="sa">r</span><span class="s1">&#39;^(.?[&lt;&gt;=^])?[+ -]?#?0?(\d+|{\w+})?[,_]?(\.(\d+|{\w+}))?[bcdefgnosx%]?$&#39;</span><span class="p">,</span> <span class="n">re</span><span class="o">.</span><span class="n">I</span><span class="p">)</span>
    <span class="n">field_spec</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="sa">r</span><span class="s1">&#39;^(\d+|\w+)(\.\w+|\[[^]]+\])*$&#39;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_format</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">record</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">validate</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Validate the input format, ensure it is the correct string formatting style&quot;&quot;&quot;</span>
        <span class="n">fields</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">_</span><span class="p">,</span> <span class="n">fieldname</span><span class="p">,</span> <span class="n">spec</span><span class="p">,</span> <span class="n">conversion</span> <span class="ow">in</span> <span class="n">_str_formatter</span><span class="o">.</span><span class="n">parse</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span><span class="p">):</span>
                <span class="k">if</span> <span class="n">fieldname</span><span class="p">:</span>
                    <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">field_spec</span><span class="o">.</span><span class="n">match</span><span class="p">(</span><span class="n">fieldname</span><span class="p">):</span>
                        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;invalid field name/expression: </span><span class="si">%r</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">fieldname</span><span class="p">)</span>
                    <span class="n">fields</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">fieldname</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">conversion</span> <span class="ow">and</span> <span class="n">conversion</span> <span class="ow">not</span> <span class="ow">in</span> <span class="s1">&#39;rsa&#39;</span><span class="p">:</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;invalid conversion: </span><span class="si">%r</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">conversion</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">spec</span> <span class="ow">and</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">fmt_spec</span><span class="o">.</span><span class="n">match</span><span class="p">(</span><span class="n">spec</span><span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;bad specifier: </span><span class="si">%r</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">spec</span><span class="p">)</span>
        <span class="k">except</span> <span class="ne">ValueError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;invalid format: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">e</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">fields</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;invalid format: no fields&#39;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">StringTemplateStyle</span><span class="p">(</span><span class="n">PercentStyle</span><span class="p">):</span>
    <span class="n">default_format</span> <span class="o">=</span> <span class="s1">&#39;$</span><span class="si">{message}</span><span class="s1">&#39;</span>
    <span class="n">asctime_format</span> <span class="o">=</span> <span class="s1">&#39;$</span><span class="si">{asctime}</span><span class="s1">&#39;</span>
    <span class="n">asctime_search</span> <span class="o">=</span> <span class="s1">&#39;$</span><span class="si">{asctime}</span><span class="s1">&#39;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fmt</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span> <span class="o">=</span> <span class="n">fmt</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">default_format</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_tpl</span> <span class="o">=</span> <span class="n">Template</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">usesTime</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">fmt</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span>
        <span class="k">return</span> <span class="n">fmt</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s1">&#39;$asctime&#39;</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">fmt</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">asctime_format</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="mi">0</span>

    <span class="k">def</span> <span class="nf">validate</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">pattern</span> <span class="o">=</span> <span class="n">Template</span><span class="o">.</span><span class="n">pattern</span>
        <span class="n">fields</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">pattern</span><span class="o">.</span><span class="n">finditer</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span><span class="p">):</span>
            <span class="n">d</span> <span class="o">=</span> <span class="n">m</span><span class="o">.</span><span class="n">groupdict</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s1">&#39;named&#39;</span><span class="p">]:</span>
                <span class="n">fields</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">d</span><span class="p">[</span><span class="s1">&#39;named&#39;</span><span class="p">])</span>
            <span class="k">elif</span> <span class="n">d</span><span class="p">[</span><span class="s1">&#39;braced&#39;</span><span class="p">]:</span>
                <span class="n">fields</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">d</span><span class="p">[</span><span class="s1">&#39;braced&#39;</span><span class="p">])</span>
            <span class="k">elif</span> <span class="n">m</span><span class="o">.</span><span class="n">group</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">==</span> <span class="s1">&#39;$&#39;</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;invalid format: bare </span><span class="se">\&#39;</span><span class="s1">$</span><span class="se">\&#39;</span><span class="s1"> not allowed&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">fields</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;invalid format: no fields&#39;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_format</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_tpl</span><span class="o">.</span><span class="n">substitute</span><span class="p">(</span><span class="o">**</span><span class="n">record</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">)</span>


<span class="n">BASIC_FORMAT</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="si">%(levelname)s</span><span class="s2">:</span><span class="si">%(name)s</span><span class="s2">:</span><span class="si">%(message)s</span><span class="s2">&quot;</span>

<span class="n">_STYLES</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s1">&#39;%&#39;</span><span class="p">:</span> <span class="p">(</span><span class="n">PercentStyle</span><span class="p">,</span> <span class="n">BASIC_FORMAT</span><span class="p">),</span>
    <span class="s1">&#39;{&#39;</span><span class="p">:</span> <span class="p">(</span><span class="n">StrFormatStyle</span><span class="p">,</span> <span class="s1">&#39;</span><span class="si">{levelname}</span><span class="s1">:</span><span class="si">{name}</span><span class="s1">:</span><span class="si">{message}</span><span class="s1">&#39;</span><span class="p">),</span>
    <span class="s1">&#39;$&#39;</span><span class="p">:</span> <span class="p">(</span><span class="n">StringTemplateStyle</span><span class="p">,</span> <span class="s1">&#39;$</span><span class="si">{levelname}</span><span class="s1">:$</span><span class="si">{name}</span><span class="s1">:$</span><span class="si">{message}</span><span class="s1">&#39;</span><span class="p">),</span>
<span class="p">}</span>

<span class="k">class</span> <span class="nc">Formatter</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Formatter instances are used to convert a LogRecord to text.</span>

<span class="sd">    Formatters need to know how a LogRecord is constructed. They are</span>
<span class="sd">    responsible for converting a LogRecord to (usually) a string which can</span>
<span class="sd">    be interpreted by either a human or an external system. The base Formatter</span>
<span class="sd">    allows a formatting string to be specified. If none is supplied, the</span>
<span class="sd">    style-dependent default value, &quot;%(message)s&quot;, &quot;{message}&quot;, or</span>
<span class="sd">    &quot;${message}&quot;, is used.</span>

<span class="sd">    The Formatter can be initialized with a format string which makes use of</span>
<span class="sd">    knowledge of the LogRecord attributes - e.g. the default value mentioned</span>
<span class="sd">    above makes use of the fact that the user&#39;s message and arguments are pre-</span>
<span class="sd">    formatted into a LogRecord&#39;s message attribute. Currently, the useful</span>
<span class="sd">    attributes in a LogRecord are described by:</span>

<span class="sd">    %(name)s            Name of the logger (logging channel)</span>
<span class="sd">    %(levelno)s         Numeric logging level for the message (DEBUG, INFO,</span>
<span class="sd">                        WARNING, ERROR, CRITICAL)</span>
<span class="sd">    %(levelname)s       Text logging level for the message (&quot;DEBUG&quot;, &quot;INFO&quot;,</span>
<span class="sd">                        &quot;WARNING&quot;, &quot;ERROR&quot;, &quot;CRITICAL&quot;)</span>
<span class="sd">    %(pathname)s        Full pathname of the source file where the logging</span>
<span class="sd">                        call was issued (if available)</span>
<span class="sd">    %(filename)s        Filename portion of pathname</span>
<span class="sd">    %(module)s          Module (name portion of filename)</span>
<span class="sd">    %(lineno)d          Source line number where the logging call was issued</span>
<span class="sd">                        (if available)</span>
<span class="sd">    %(funcName)s        Function name</span>
<span class="sd">    %(created)f         Time when the LogRecord was created (time.time()</span>
<span class="sd">                        return value)</span>
<span class="sd">    %(asctime)s         Textual time when the LogRecord was created</span>
<span class="sd">    %(msecs)d           Millisecond portion of the creation time</span>
<span class="sd">    %(relativeCreated)d Time in milliseconds when the LogRecord was created,</span>
<span class="sd">                        relative to the time the logging module was loaded</span>
<span class="sd">                        (typically at application startup time)</span>
<span class="sd">    %(thread)d          Thread ID (if available)</span>
<span class="sd">    %(threadName)s      Thread name (if available)</span>
<span class="sd">    %(process)d         Process ID (if available)</span>
<span class="sd">    %(message)s         The result of record.getMessage(), computed just as</span>
<span class="sd">                        the record is emitted</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">converter</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">localtime</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fmt</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">datefmt</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">style</span><span class="o">=</span><span class="s1">&#39;%&#39;</span><span class="p">,</span> <span class="n">validate</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the formatter with specified format strings.</span>

<span class="sd">        Initialize the formatter either with the specified format string, or a</span>
<span class="sd">        default as described above. Allow for specialized date formatting with</span>
<span class="sd">        the optional datefmt argument. If datefmt is omitted, you get an</span>
<span class="sd">        ISO8601-like (or RFC 3339-like) format.</span>

<span class="sd">        Use a style parameter of &#39;%&#39;, &#39;{&#39; or &#39;$&#39; to specify that you want to</span>
<span class="sd">        use one of %-formatting, :meth:`str.format` (``{}``) formatting or</span>
<span class="sd">        :class:`string.Template` formatting in your format string.</span>

<span class="sd">        .. versionchanged:: 3.2</span>
<span class="sd">           Added the ``style`` parameter.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">style</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">_STYLES</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Style must be one of: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="s1">&#39;,&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span>
                             <span class="n">_STYLES</span><span class="o">.</span><span class="n">keys</span><span class="p">()))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_style</span> <span class="o">=</span> <span class="n">_STYLES</span><span class="p">[</span><span class="n">style</span><span class="p">][</span><span class="mi">0</span><span class="p">](</span><span class="n">fmt</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">validate</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_style</span><span class="o">.</span><span class="n">validate</span><span class="p">()</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_fmt</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_style</span><span class="o">.</span><span class="n">_fmt</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">datefmt</span> <span class="o">=</span> <span class="n">datefmt</span>

    <span class="n">default_time_format</span> <span class="o">=</span> <span class="s1">&#39;%Y-%m-</span><span class="si">%d</span><span class="s1"> %H:%M:%S&#39;</span>
    <span class="n">default_msec_format</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">%s</span><span class="s1">,</span><span class="si">%03d</span><span class="s1">&#39;</span>

    <span class="k">def</span> <span class="nf">formatTime</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">,</span> <span class="n">datefmt</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the creation time of the specified LogRecord as formatted text.</span>

<span class="sd">        This method should be called from format() by a formatter which</span>
<span class="sd">        wants to make use of a formatted time. This method can be overridden</span>
<span class="sd">        in formatters to provide for any specific requirement, but the</span>
<span class="sd">        basic behaviour is as follows: if datefmt (a string) is specified,</span>
<span class="sd">        it is used with time.strftime() to format the creation time of the</span>
<span class="sd">        record. Otherwise, an ISO8601-like (or RFC 3339-like) format is used.</span>
<span class="sd">        The resulting string is returned. This function uses a user-configurable</span>
<span class="sd">        function to convert the creation time to a tuple. By default,</span>
<span class="sd">        time.localtime() is used; to change this for a particular formatter</span>
<span class="sd">        instance, set the &#39;converter&#39; attribute to a function with the same</span>
<span class="sd">        signature as time.localtime() or time.gmtime(). To change it for all</span>
<span class="sd">        formatters, for example if you want all logging times to be shown in GMT,</span>
<span class="sd">        set the &#39;converter&#39; attribute in the Formatter class.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ct</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">converter</span><span class="p">(</span><span class="n">record</span><span class="o">.</span><span class="n">created</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">datefmt</span><span class="p">:</span>
            <span class="n">s</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">strftime</span><span class="p">(</span><span class="n">datefmt</span><span class="p">,</span> <span class="n">ct</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">s</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">strftime</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">default_time_format</span><span class="p">,</span> <span class="n">ct</span><span class="p">)</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">default_msec_format</span><span class="p">:</span>
                <span class="n">s</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">default_msec_format</span> <span class="o">%</span> <span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="n">record</span><span class="o">.</span><span class="n">msecs</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">s</span>

    <span class="k">def</span> <span class="nf">formatException</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ei</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Format and return the specified exception information as a string.</span>

<span class="sd">        This default implementation just uses</span>
<span class="sd">        traceback.print_exception()</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">sio</span> <span class="o">=</span> <span class="n">io</span><span class="o">.</span><span class="n">StringIO</span><span class="p">()</span>
        <span class="n">tb</span> <span class="o">=</span> <span class="n">ei</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
        <span class="c1"># See issues #9427, #1553375. Commented out for now.</span>
        <span class="c1">#if getattr(self, &#39;fullstack&#39;, False):</span>
        <span class="c1">#    traceback.print_stack(tb.tb_frame.f_back, file=sio)</span>
        <span class="n">traceback</span><span class="o">.</span><span class="n">print_exception</span><span class="p">(</span><span class="n">ei</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">ei</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">tb</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="n">sio</span><span class="p">)</span>
        <span class="n">s</span> <span class="o">=</span> <span class="n">sio</span><span class="o">.</span><span class="n">getvalue</span><span class="p">()</span>
        <span class="n">sio</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">s</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">:]</span> <span class="o">==</span> <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span><span class="p">:</span>
            <span class="n">s</span> <span class="o">=</span> <span class="n">s</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">s</span>

    <span class="k">def</span> <span class="nf">usesTime</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Check if the format uses the creation time of the record.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_style</span><span class="o">.</span><span class="n">usesTime</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">formatMessage</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_style</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">formatStack</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">stack_info</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This method is provided as an extension point for specialized</span>
<span class="sd">        formatting of stack information.</span>

<span class="sd">        The input data is a string as returned from a call to</span>
<span class="sd">        :func:`traceback.print_stack`, but with the last trailing newline</span>
<span class="sd">        removed.</span>

<span class="sd">        The base implementation just returns the value passed in.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">stack_info</span>

    <span class="k">def</span> <span class="nf">format</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Format the specified record as text.</span>

<span class="sd">        The record&#39;s attribute dictionary is used as the operand to a</span>
<span class="sd">        string formatting operation which yields the returned string.</span>
<span class="sd">        Before formatting the dictionary, a couple of preparatory steps</span>
<span class="sd">        are carried out. The message attribute of the record is computed</span>
<span class="sd">        using LogRecord.getMessage(). If the formatting string uses the</span>
<span class="sd">        time (as determined by a call to usesTime(), formatTime() is</span>
<span class="sd">        called to format the event time. If there is exception information,</span>
<span class="sd">        it is formatted using formatException() and appended to the message.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">record</span><span class="o">.</span><span class="n">message</span> <span class="o">=</span> <span class="n">record</span><span class="o">.</span><span class="n">getMessage</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">usesTime</span><span class="p">():</span>
            <span class="n">record</span><span class="o">.</span><span class="n">asctime</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">formatTime</span><span class="p">(</span><span class="n">record</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">datefmt</span><span class="p">)</span>
        <span class="n">s</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">formatMessage</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">record</span><span class="o">.</span><span class="n">exc_info</span><span class="p">:</span>
            <span class="c1"># Cache the traceback text to avoid converting it multiple times</span>
            <span class="c1"># (it&#39;s constant anyway)</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">record</span><span class="o">.</span><span class="n">exc_text</span><span class="p">:</span>
                <span class="n">record</span><span class="o">.</span><span class="n">exc_text</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">formatException</span><span class="p">(</span><span class="n">record</span><span class="o">.</span><span class="n">exc_info</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">record</span><span class="o">.</span><span class="n">exc_text</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">s</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">:]</span> <span class="o">!=</span> <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span><span class="p">:</span>
                <span class="n">s</span> <span class="o">=</span> <span class="n">s</span> <span class="o">+</span> <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span>
            <span class="n">s</span> <span class="o">=</span> <span class="n">s</span> <span class="o">+</span> <span class="n">record</span><span class="o">.</span><span class="n">exc_text</span>
        <span class="k">if</span> <span class="n">record</span><span class="o">.</span><span class="n">stack_info</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">s</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">:]</span> <span class="o">!=</span> <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span><span class="p">:</span>
                <span class="n">s</span> <span class="o">=</span> <span class="n">s</span> <span class="o">+</span> <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span>
            <span class="n">s</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">formatStack</span><span class="p">(</span><span class="n">record</span><span class="o">.</span><span class="n">stack_info</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">s</span>

<span class="c1">#</span>
<span class="c1">#   The default formatter to use when no other is specified</span>
<span class="c1">#</span>
<span class="n">_defaultFormatter</span> <span class="o">=</span> <span class="n">Formatter</span><span class="p">()</span>

<span class="k">class</span> <span class="nc">BufferingFormatter</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A formatter suitable for formatting a number of records.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">linefmt</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Optionally specify a formatter which will be used to format each</span>
<span class="sd">        individual record.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">linefmt</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">linefmt</span> <span class="o">=</span> <span class="n">linefmt</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">linefmt</span> <span class="o">=</span> <span class="n">_defaultFormatter</span>

    <span class="k">def</span> <span class="nf">formatHeader</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">records</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the header string for the specified records.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="s2">&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">formatFooter</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">records</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Return the footer string for the specified records.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="s2">&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">format</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">records</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Format the specified records and return the result as a string.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">records</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">rv</span> <span class="o">=</span> <span class="n">rv</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">formatHeader</span><span class="p">(</span><span class="n">records</span><span class="p">)</span>
            <span class="k">for</span> <span class="n">record</span> <span class="ow">in</span> <span class="n">records</span><span class="p">:</span>
                <span class="n">rv</span> <span class="o">=</span> <span class="n">rv</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">linefmt</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
            <span class="n">rv</span> <span class="o">=</span> <span class="n">rv</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">formatFooter</span><span class="p">(</span><span class="n">records</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">rv</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Filter classes and functions</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="k">class</span> <span class="nc">Filter</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Filter instances are used to perform arbitrary filtering of LogRecords.</span>

<span class="sd">    Loggers and Handlers can optionally use Filter instances to filter</span>
<span class="sd">    records as desired. The base filter class only allows events which are</span>
<span class="sd">    below a certain point in the logger hierarchy. For example, a filter</span>
<span class="sd">    initialized with &quot;A.B&quot; will allow events logged by loggers &quot;A.B&quot;,</span>
<span class="sd">    &quot;A.B.C&quot;, &quot;A.B.C.D&quot;, &quot;A.B.D&quot; etc. but not &quot;A.BB&quot;, &quot;B.A.B&quot; etc. If</span>
<span class="sd">    initialized with the empty string, all events are passed.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s1">&#39;&#39;</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize a filter.</span>

<span class="sd">        Initialize with the name of the logger which, together with its</span>
<span class="sd">        children, will have its events allowed through the filter. If no</span>
<span class="sd">        name is specified, allow every event.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">=</span> <span class="n">name</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">nlen</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">filter</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Determine if the specified record is to be logged.</span>

<span class="sd">        Returns True if the record should be logged, or False otherwise.</span>
<span class="sd">        If deemed appropriate, the record may be modified in-place.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">nlen</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">==</span> <span class="n">record</span><span class="o">.</span><span class="n">name</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="k">elif</span> <span class="n">record</span><span class="o">.</span><span class="n">name</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">name</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">nlen</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">False</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">record</span><span class="o">.</span><span class="n">name</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">nlen</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;.&quot;</span><span class="p">)</span>

<span class="k">class</span> <span class="nc">Filterer</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A base class for loggers and handlers which allows them to share</span>
<span class="sd">    common code.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the list of filters to be an empty list.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">filters</span> <span class="o">=</span> <span class="p">[]</span>

    <span class="k">def</span> <span class="nf">addFilter</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="nb">filter</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Add the specified filter to this handler.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="nb">filter</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">filters</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">filters</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">filter</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">removeFilter</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="nb">filter</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Remove the specified filter from this handler.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">filter</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">filters</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">filters</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="nb">filter</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">filter</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Determine if a record is loggable by consulting all the filters.</span>

<span class="sd">        The default is to allow the record to be logged; any filter can veto</span>
<span class="sd">        this and the record is then dropped. Returns a zero value if a record</span>
<span class="sd">        is to be dropped, else non-zero.</span>

<span class="sd">        .. versionchanged:: 3.2</span>

<span class="sd">           Allow filters to be just callables.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="k">for</span> <span class="n">f</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">filters</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="s1">&#39;filter&#39;</span><span class="p">):</span>
                <span class="n">result</span> <span class="o">=</span> <span class="n">f</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">result</span> <span class="o">=</span> <span class="n">f</span><span class="p">(</span><span class="n">record</span><span class="p">)</span> <span class="c1"># assume callable - will raise if not</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">result</span><span class="p">:</span>
                <span class="n">rv</span> <span class="o">=</span> <span class="kc">False</span>
                <span class="k">break</span>
        <span class="k">return</span> <span class="n">rv</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Handler classes and functions</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="n">_handlers</span> <span class="o">=</span> <span class="n">weakref</span><span class="o">.</span><span class="n">WeakValueDictionary</span><span class="p">()</span>  <span class="c1">#map of handler names to handlers</span>
<span class="n">_handlerList</span> <span class="o">=</span> <span class="p">[]</span> <span class="c1"># added to allow handlers to be removed in reverse of order initialized</span>

<span class="k">def</span> <span class="nf">_removeHandlerRef</span><span class="p">(</span><span class="n">wr</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Remove a handler reference from the internal cleanup list.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This function can be called during module teardown, when globals are</span>
    <span class="c1"># set to None. It can also be called from another thread. So we need to</span>
    <span class="c1"># pre-emptively grab the necessary globals and check if they&#39;re None,</span>
    <span class="c1"># to prevent race conditions and failures during interpreter shutdown.</span>
    <span class="n">acquire</span><span class="p">,</span> <span class="n">release</span><span class="p">,</span> <span class="n">handlers</span> <span class="o">=</span> <span class="n">_acquireLock</span><span class="p">,</span> <span class="n">_releaseLock</span><span class="p">,</span> <span class="n">_handlerList</span>
    <span class="k">if</span> <span class="n">acquire</span> <span class="ow">and</span> <span class="n">release</span> <span class="ow">and</span> <span class="n">handlers</span><span class="p">:</span>
        <span class="n">acquire</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">wr</span> <span class="ow">in</span> <span class="n">handlers</span><span class="p">:</span>
                <span class="n">handlers</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">wr</span><span class="p">)</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">release</span><span class="p">()</span>

<span class="k">def</span> <span class="nf">_addHandlerRef</span><span class="p">(</span><span class="n">handler</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Add a handler to the internal cleanup list using a weak reference.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_acquireLock</span><span class="p">()</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">_handlerList</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">weakref</span><span class="o">.</span><span class="n">ref</span><span class="p">(</span><span class="n">handler</span><span class="p">,</span> <span class="n">_removeHandlerRef</span><span class="p">))</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="n">_releaseLock</span><span class="p">()</span>

<span class="k">class</span> <span class="nc">Handler</span><span class="p">(</span><span class="n">Filterer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Handler instances dispatch logging events to specific destinations.</span>

<span class="sd">    The base handler class. Acts as a placeholder which defines the Handler</span>
<span class="sd">    interface. Handlers can optionally use Formatter instances to format</span>
<span class="sd">    records as desired. By default, no formatter is specified; in this case,</span>
<span class="sd">    the &#39;raw&#39; message as determined by record.message is logged.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="o">=</span><span class="n">NOTSET</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initializes the instance - basically setting the formatter to None</span>
<span class="sd">        and the filter list to empty.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">Filterer</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_name</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">level</span> <span class="o">=</span> <span class="n">_checkLevel</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">formatter</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="c1"># Add the handler to the global _handlerList (for cleanup on shutdown)</span>
        <span class="n">_addHandlerRef</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">createLock</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">get_name</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_name</span>

    <span class="k">def</span> <span class="nf">set_name</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
        <span class="n">_acquireLock</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_name</span> <span class="ow">in</span> <span class="n">_handlers</span><span class="p">:</span>
                <span class="k">del</span> <span class="n">_handlers</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">_name</span><span class="p">]</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_name</span> <span class="o">=</span> <span class="n">name</span>
            <span class="k">if</span> <span class="n">name</span><span class="p">:</span>
                <span class="n">_handlers</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">_releaseLock</span><span class="p">()</span>

    <span class="n">name</span> <span class="o">=</span> <span class="nb">property</span><span class="p">(</span><span class="n">get_name</span><span class="p">,</span> <span class="n">set_name</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">createLock</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Acquire a thread lock for serializing access to the underlying I/O.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lock</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">RLock</span><span class="p">()</span>
        <span class="n">_register_at_fork_reinit_lock</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_at_fork_reinit</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lock</span><span class="o">.</span><span class="n">_at_fork_reinit</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">acquire</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Acquire the I/O thread lock.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">lock</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">lock</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">release</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Release the I/O thread lock.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">lock</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">lock</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">setLevel</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set the logging level of this handler.  level must be an int or a str.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">level</span> <span class="o">=</span> <span class="n">_checkLevel</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">format</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Format the specified record.</span>

<span class="sd">        If a formatter is set, use it. Otherwise, use the default formatter</span>
<span class="sd">        for the module.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">formatter</span><span class="p">:</span>
            <span class="n">fmt</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">formatter</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">fmt</span> <span class="o">=</span> <span class="n">_defaultFormatter</span>
        <span class="k">return</span> <span class="n">fmt</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">emit</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Do whatever it takes to actually log the specified logging record.</span>

<span class="sd">        This version is intended to be implemented by subclasses and so</span>
<span class="sd">        raises a NotImplementedError.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s1">&#39;emit must be implemented &#39;</span>
                                  <span class="s1">&#39;by Handler subclasses&#39;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">handle</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Conditionally emit the specified logging record.</span>

<span class="sd">        Emission depends on filters which may have been added to the handler.</span>
<span class="sd">        Wrap the actual emission of the record with acquisition/release of</span>
<span class="sd">        the I/O thread lock. Returns whether the filter passed the record for</span>
<span class="sd">        emission.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">rv</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">rv</span>

    <span class="k">def</span> <span class="nf">setFormatter</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fmt</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set the formatter for this handler.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">formatter</span> <span class="o">=</span> <span class="n">fmt</span>

    <span class="k">def</span> <span class="nf">flush</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Ensure all logging output has been flushed.</span>

<span class="sd">        This version does nothing and is intended to be implemented by</span>
<span class="sd">        subclasses.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">pass</span>

    <span class="k">def</span> <span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Tidy up any resources used by the handler.</span>

<span class="sd">        This version removes the handler from an internal map of handlers,</span>
<span class="sd">        _handlers, which is used for handler lookup by name. Subclasses</span>
<span class="sd">        should ensure that this gets called from overridden close()</span>
<span class="sd">        methods.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1">#get the module data lock, as we&#39;re updating a shared structure.</span>
        <span class="n">_acquireLock</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>    <span class="c1">#unlikely to raise an exception, but you never know...</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_name</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">_name</span> <span class="ow">in</span> <span class="n">_handlers</span><span class="p">:</span>
                <span class="k">del</span> <span class="n">_handlers</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">_name</span><span class="p">]</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">_releaseLock</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">handleError</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Handle errors which occur during an emit() call.</span>

<span class="sd">        This method should be called from handlers when an exception is</span>
<span class="sd">        encountered during an emit() call. If raiseExceptions is false,</span>
<span class="sd">        exceptions get silently ignored. This is what is mostly wanted</span>
<span class="sd">        for a logging system - most users will not care about errors in</span>
<span class="sd">        the logging system, they are more interested in application errors.</span>
<span class="sd">        You could, however, replace this with a custom handler if you wish.</span>
<span class="sd">        The record which was being processed is passed in to this method.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">raiseExceptions</span> <span class="ow">and</span> <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="p">:</span>  <span class="c1"># see issue 13807</span>
            <span class="n">t</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">tb</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">exc_info</span><span class="p">()</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;--- Logging error ---</span><span class="se">\n</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="n">traceback</span><span class="o">.</span><span class="n">print_exception</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">tb</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="p">)</span>
                <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;Call stack:</span><span class="se">\n</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="c1"># Walk the stack frame up until we&#39;re out of logging,</span>
                <span class="c1"># so as to print the calling context.</span>
                <span class="n">frame</span> <span class="o">=</span> <span class="n">tb</span><span class="o">.</span><span class="n">tb_frame</span>
                <span class="k">while</span> <span class="p">(</span><span class="n">frame</span> <span class="ow">and</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">frame</span><span class="o">.</span><span class="n">f_code</span><span class="o">.</span><span class="n">co_filename</span><span class="p">)</span> <span class="o">==</span>
                       <span class="n">__path__</span><span class="p">[</span><span class="mi">0</span><span class="p">]):</span>
                    <span class="n">frame</span> <span class="o">=</span> <span class="n">frame</span><span class="o">.</span><span class="n">f_back</span>
                <span class="k">if</span> <span class="n">frame</span><span class="p">:</span>
                    <span class="n">traceback</span><span class="o">.</span><span class="n">print_stack</span><span class="p">(</span><span class="n">frame</span><span class="p">,</span> <span class="n">file</span><span class="o">=</span><span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="c1"># couldn&#39;t find the right stack frame, for some reason</span>
                    <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;Logged from file </span><span class="si">%s</span><span class="s1">, line </span><span class="si">%s</span><span class="se">\n</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="p">(</span>
                                     <span class="n">record</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="n">record</span><span class="o">.</span><span class="n">lineno</span><span class="p">))</span>
                <span class="c1"># Issue 18671: output logging message and arguments</span>
                <span class="k">try</span><span class="p">:</span>
                    <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;Message: </span><span class="si">%r</span><span class="se">\n</span><span class="s1">&#39;</span>
                                     <span class="s1">&#39;Arguments: </span><span class="si">%s</span><span class="se">\n</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="n">record</span><span class="o">.</span><span class="n">msg</span><span class="p">,</span>
                                                          <span class="n">record</span><span class="o">.</span><span class="n">args</span><span class="p">))</span>
                <span class="k">except</span> <span class="ne">RecursionError</span><span class="p">:</span>  <span class="c1"># See issue 36272</span>
                    <span class="k">raise</span>
                <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
                    <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;Unable to print the message and arguments&#39;</span>
                                     <span class="s1">&#39; - possible formatting error.</span><span class="se">\n</span><span class="s1">Use the&#39;</span>
                                     <span class="s1">&#39; traceback above to help find the error.</span><span class="se">\n</span><span class="s1">&#39;</span>
                                    <span class="p">)</span>
            <span class="k">except</span> <span class="ne">OSError</span><span class="p">:</span> <span class="c1">#pragma: no cover</span>
                <span class="k">pass</span>    <span class="c1"># see issue 5971</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="k">del</span> <span class="n">t</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">tb</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">level</span> <span class="o">=</span> <span class="n">getLevelName</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">level</span><span class="p">)</span>
        <span class="k">return</span> <span class="s1">&#39;&lt;</span><span class="si">%s</span><span class="s1"> (</span><span class="si">%s</span><span class="s1">)&gt;&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> <span class="n">level</span><span class="p">)</span>

<span class="k">class</span> <span class="nc">StreamHandler</span><span class="p">(</span><span class="n">Handler</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A handler class which writes logging records, appropriately formatted,</span>
<span class="sd">    to a stream. Note that this class does not close the stream, as</span>
<span class="sd">    sys.stdout or sys.stderr may be used.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">terminator</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\n</span><span class="s1">&#39;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">stream</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the handler.</span>

<span class="sd">        If stream is not specified, sys.stderr is used.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">Handler</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">stream</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">stream</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="o">=</span> <span class="n">stream</span>

    <span class="k">def</span> <span class="nf">flush</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Flushes the stream.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="ow">and</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">stream</span><span class="p">,</span> <span class="s2">&quot;flush&quot;</span><span class="p">):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">stream</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">emit</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Emit a record.</span>

<span class="sd">        If a formatter is specified, it is used to format the record.</span>
<span class="sd">        The record is then written to the stream with a trailing newline.  If</span>
<span class="sd">        exception information is present, it is formatted using</span>
<span class="sd">        traceback.print_exception and appended to the stream.  If the stream</span>
<span class="sd">        has an &#39;encoding&#39; attribute, it is used to determine how to do the</span>
<span class="sd">        output to the stream.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">msg</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
            <span class="n">stream</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">stream</span>
            <span class="c1"># issue 35046: merged two stream.writes into one.</span>
            <span class="n">stream</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">msg</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">terminator</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
        <span class="k">except</span> <span class="ne">RecursionError</span><span class="p">:</span>  <span class="c1"># See issue 36272</span>
            <span class="k">raise</span>
        <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">handleError</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">setStream</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">stream</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Sets the StreamHandler&#39;s stream to the specified value,</span>
<span class="sd">        if it is different.</span>

<span class="sd">        Returns the old stream, if the stream was changed, or None</span>
<span class="sd">        if it wasn&#39;t.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">stream</span> <span class="ow">is</span> <span class="bp">self</span><span class="o">.</span><span class="n">stream</span><span class="p">:</span>
            <span class="n">result</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">result</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">stream</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="o">=</span> <span class="n">stream</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">result</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">level</span> <span class="o">=</span> <span class="n">getLevelName</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">level</span><span class="p">)</span>
        <span class="n">name</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">stream</span><span class="p">,</span> <span class="s1">&#39;name&#39;</span><span class="p">,</span> <span class="s1">&#39;&#39;</span><span class="p">)</span>
        <span class="c1">#  bpo-36015: name can be an int</span>
        <span class="n">name</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">name</span><span class="p">:</span>
            <span class="n">name</span> <span class="o">+=</span> <span class="s1">&#39; &#39;</span>
        <span class="k">return</span> <span class="s1">&#39;&lt;</span><span class="si">%s</span><span class="s1"> </span><span class="si">%s</span><span class="s1">(</span><span class="si">%s</span><span class="s1">)&gt;&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">FileHandler</span><span class="p">(</span><span class="n">StreamHandler</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A handler class which writes formatted logging records to disk files.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filename</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s1">&#39;a&#39;</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">delay</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">errors</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Open the specified file and use it as the stream for logging.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># Issue #27493: add support for Path objects to be passed in</span>
        <span class="n">filename</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">fspath</span><span class="p">(</span><span class="n">filename</span><span class="p">)</span>
        <span class="c1">#keep the absolute path, otherwise derived classes which use this</span>
        <span class="c1">#may come a cropper when the current directory changes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">baseFilename</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">filename</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">=</span> <span class="n">mode</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">encoding</span> <span class="o">=</span> <span class="n">encoding</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">errors</span> <span class="o">=</span> <span class="n">errors</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">delay</span> <span class="o">=</span> <span class="n">delay</span>
        <span class="k">if</span> <span class="n">delay</span><span class="p">:</span>
            <span class="c1">#We don&#39;t open the stream, but we still need to call the</span>
            <span class="c1">#Handler constructor to set level, formatter, lock etc.</span>
            <span class="n">Handler</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">StreamHandler</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_open</span><span class="p">())</span>

    <span class="k">def</span> <span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Closes the stream.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">stream</span><span class="p">:</span>
                    <span class="k">try</span><span class="p">:</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
                    <span class="k">finally</span><span class="p">:</span>
                        <span class="n">stream</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">stream</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="o">=</span> <span class="kc">None</span>
                        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">stream</span><span class="p">,</span> <span class="s2">&quot;close&quot;</span><span class="p">):</span>
                            <span class="n">stream</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="c1"># Issue #19523: call unconditionally to</span>
                <span class="c1"># prevent a handler leak when delay is set</span>
                <span class="n">StreamHandler</span><span class="o">.</span><span class="n">close</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_open</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Open the current base file with the (original) mode and encoding.</span>
<span class="sd">        Return the resulting stream.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="nb">open</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">baseFilename</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">encoding</span><span class="p">,</span>
                    <span class="n">errors</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">errors</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">emit</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Emit a record.</span>

<span class="sd">        If the stream was not opened because &#39;delay&#39; was specified in the</span>
<span class="sd">        constructor, open it before calling the superclass&#39;s emit.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_open</span><span class="p">()</span>
        <span class="n">StreamHandler</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">level</span> <span class="o">=</span> <span class="n">getLevelName</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">level</span><span class="p">)</span>
        <span class="k">return</span> <span class="s1">&#39;&lt;</span><span class="si">%s</span><span class="s1"> </span><span class="si">%s</span><span class="s1"> (</span><span class="si">%s</span><span class="s1">)&gt;&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">baseFilename</span><span class="p">,</span> <span class="n">level</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_StderrHandler</span><span class="p">(</span><span class="n">StreamHandler</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This class is like a StreamHandler using sys.stderr, but always uses</span>
<span class="sd">    whatever sys.stderr is currently set to rather than the value of</span>
<span class="sd">    sys.stderr at handler construction time.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="o">=</span><span class="n">NOTSET</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the handler.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">Handler</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">)</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">stream</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span>


<span class="n">_defaultLastResort</span> <span class="o">=</span> <span class="n">_StderrHandler</span><span class="p">(</span><span class="n">WARNING</span><span class="p">)</span>
<span class="n">lastResort</span> <span class="o">=</span> <span class="n">_defaultLastResort</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Manager classes and functions</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="k">class</span> <span class="nc">PlaceHolder</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    PlaceHolder instances are used in the Manager logger hierarchy to take</span>
<span class="sd">    the place of nodes for which no loggers have been defined. This class is</span>
<span class="sd">    intended for internal use only and not as part of the public API.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">alogger</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize with the specified logger being a child of this placeholder.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">loggerMap</span> <span class="o">=</span> <span class="p">{</span> <span class="n">alogger</span> <span class="p">:</span> <span class="kc">None</span> <span class="p">}</span>

    <span class="k">def</span> <span class="nf">append</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">alogger</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Add the specified logger as a child of this placeholder.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">alogger</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">loggerMap</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">loggerMap</span><span class="p">[</span><span class="n">alogger</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>

<span class="c1">#</span>
<span class="c1">#   Determine which class to use when instantiating loggers.</span>
<span class="c1">#</span>

<span class="k">def</span> <span class="nf">setLoggerClass</span><span class="p">(</span><span class="n">klass</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Set the class to be used when instantiating a logger. The class should</span>
<span class="sd">    define __init__() such that only a name argument is required, and the</span>
<span class="sd">    __init__() should call Logger.__init__()</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">klass</span> <span class="o">!=</span> <span class="n">Logger</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">klass</span><span class="p">,</span> <span class="n">Logger</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;logger not derived from logging.Logger: &quot;</span>
                            <span class="o">+</span> <span class="n">klass</span><span class="o">.</span><span class="vm">__name__</span><span class="p">)</span>
    <span class="k">global</span> <span class="n">_loggerClass</span>
    <span class="n">_loggerClass</span> <span class="o">=</span> <span class="n">klass</span>

<span class="k">def</span> <span class="nf">getLoggerClass</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Return the class to be used when instantiating a logger.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">_loggerClass</span>

<span class="k">class</span> <span class="nc">Manager</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    There is [under normal circumstances] just one Manager instance, which</span>
<span class="sd">    holds the hierarchy of loggers.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">rootnode</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the manager with the root node of the logger hierarchy.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">root</span> <span class="o">=</span> <span class="n">rootnode</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">disable</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">emittedNoHandlerWarning</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">loggerClass</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logRecordFactory</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">disable</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_disable</span>

    <span class="nd">@disable</span><span class="o">.</span><span class="n">setter</span>
    <span class="k">def</span> <span class="nf">disable</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_disable</span> <span class="o">=</span> <span class="n">_checkLevel</span><span class="p">(</span><span class="n">value</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">getLogger</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get a logger with the specified name (channel name), creating it</span>
<span class="sd">        if it doesn&#39;t yet exist. This name is a dot-separated hierarchical</span>
<span class="sd">        name, such as &quot;a&quot;, &quot;a.b&quot;, &quot;a.b.c&quot; or similar.</span>

<span class="sd">        If a PlaceHolder existed for the specified name [i.e. the logger</span>
<span class="sd">        didn&#39;t exist but a child of it did], replace it with the created</span>
<span class="sd">        logger and fix up the parent/child references which pointed to the</span>
<span class="sd">        placeholder to now point to the logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s1">&#39;A logger name must be a string&#39;</span><span class="p">)</span>
        <span class="n">_acquireLock</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">name</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="p">:</span>
                <span class="n">rv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="p">[</span><span class="n">name</span><span class="p">]</span>
                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">rv</span><span class="p">,</span> <span class="n">PlaceHolder</span><span class="p">):</span>
                    <span class="n">ph</span> <span class="o">=</span> <span class="n">rv</span>
                    <span class="n">rv</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loggerClass</span> <span class="ow">or</span> <span class="n">_loggerClass</span><span class="p">)(</span><span class="n">name</span><span class="p">)</span>
                    <span class="n">rv</span><span class="o">.</span><span class="n">manager</span> <span class="o">=</span> <span class="bp">self</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">rv</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">_fixupChildren</span><span class="p">(</span><span class="n">ph</span><span class="p">,</span> <span class="n">rv</span><span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">_fixupParents</span><span class="p">(</span><span class="n">rv</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">rv</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loggerClass</span> <span class="ow">or</span> <span class="n">_loggerClass</span><span class="p">)(</span><span class="n">name</span><span class="p">)</span>
                <span class="n">rv</span><span class="o">.</span><span class="n">manager</span> <span class="o">=</span> <span class="bp">self</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">rv</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_fixupParents</span><span class="p">(</span><span class="n">rv</span><span class="p">)</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">_releaseLock</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">rv</span>

    <span class="k">def</span> <span class="nf">setLoggerClass</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">klass</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set the class to be used when instantiating a logger with this Manager.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">klass</span> <span class="o">!=</span> <span class="n">Logger</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">klass</span><span class="p">,</span> <span class="n">Logger</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;logger not derived from logging.Logger: &quot;</span>
                                <span class="o">+</span> <span class="n">klass</span><span class="o">.</span><span class="vm">__name__</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">loggerClass</span> <span class="o">=</span> <span class="n">klass</span>

    <span class="k">def</span> <span class="nf">setLogRecordFactory</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">factory</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set the factory to be used when instantiating a log record with this</span>
<span class="sd">        Manager.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logRecordFactory</span> <span class="o">=</span> <span class="n">factory</span>

    <span class="k">def</span> <span class="nf">_fixupParents</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">alogger</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Ensure that there are either loggers or placeholders all the way</span>
<span class="sd">        from the specified logger to the root of the logger hierarchy.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">name</span> <span class="o">=</span> <span class="n">alogger</span><span class="o">.</span><span class="n">name</span>
        <span class="n">i</span> <span class="o">=</span> <span class="n">name</span><span class="o">.</span><span class="n">rfind</span><span class="p">(</span><span class="s2">&quot;.&quot;</span><span class="p">)</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">while</span> <span class="p">(</span><span class="n">i</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">rv</span><span class="p">:</span>
            <span class="n">substr</span> <span class="o">=</span> <span class="n">name</span><span class="p">[:</span><span class="n">i</span><span class="p">]</span>
            <span class="k">if</span> <span class="n">substr</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="p">[</span><span class="n">substr</span><span class="p">]</span> <span class="o">=</span> <span class="n">PlaceHolder</span><span class="p">(</span><span class="n">alogger</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">obj</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="p">[</span><span class="n">substr</span><span class="p">]</span>
                <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="n">Logger</span><span class="p">):</span>
                    <span class="n">rv</span> <span class="o">=</span> <span class="n">obj</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="k">assert</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="n">PlaceHolder</span><span class="p">)</span>
                    <span class="n">obj</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">alogger</span><span class="p">)</span>
            <span class="n">i</span> <span class="o">=</span> <span class="n">name</span><span class="o">.</span><span class="n">rfind</span><span class="p">(</span><span class="s2">&quot;.&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">i</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">rv</span><span class="p">:</span>
            <span class="n">rv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">root</span>
        <span class="n">alogger</span><span class="o">.</span><span class="n">parent</span> <span class="o">=</span> <span class="n">rv</span>

    <span class="k">def</span> <span class="nf">_fixupChildren</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ph</span><span class="p">,</span> <span class="n">alogger</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Ensure that children of the placeholder ph are connected to the</span>
<span class="sd">        specified logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">name</span> <span class="o">=</span> <span class="n">alogger</span><span class="o">.</span><span class="n">name</span>
        <span class="n">namelen</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="n">ph</span><span class="o">.</span><span class="n">loggerMap</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
            <span class="c1">#The if means ... if not c.parent.name.startswith(nm)</span>
            <span class="k">if</span> <span class="n">c</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">name</span><span class="p">[:</span><span class="n">namelen</span><span class="p">]</span> <span class="o">!=</span> <span class="n">name</span><span class="p">:</span>
                <span class="n">alogger</span><span class="o">.</span><span class="n">parent</span> <span class="o">=</span> <span class="n">c</span><span class="o">.</span><span class="n">parent</span>
                <span class="n">c</span><span class="o">.</span><span class="n">parent</span> <span class="o">=</span> <span class="n">alogger</span>

    <span class="k">def</span> <span class="nf">_clear_cache</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Clear the cache for all loggers in loggerDict</span>
<span class="sd">        Called when level changes are made</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="n">_acquireLock</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">logger</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">loggerDict</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">logger</span><span class="p">,</span> <span class="n">Logger</span><span class="p">):</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">_cache</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">root</span><span class="o">.</span><span class="n">_cache</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
        <span class="n">_releaseLock</span><span class="p">()</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1">#   Logger classes and functions</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="k">class</span> <span class="nc">Logger</span><span class="p">(</span><span class="n">Filterer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Instances of the Logger class represent a single logging channel. A</span>
<span class="sd">    &quot;logging channel&quot; indicates an area of an application. Exactly how an</span>
<span class="sd">    &quot;area&quot; is defined is up to the application developer. Since an</span>
<span class="sd">    application can have any number of areas, logging channels are identified</span>
<span class="sd">    by a unique string. Application areas can be nested (e.g. an area</span>
<span class="sd">    of &quot;input processing&quot; might include sub-areas &quot;read CSV files&quot;, &quot;read</span>
<span class="sd">    XLS files&quot; and &quot;read Gnumeric files&quot;). To cater for this natural nesting,</span>
<span class="sd">    channel names are organized into a namespace hierarchy where levels are</span>
<span class="sd">    separated by periods, much like the Java or Python package namespace. So</span>
<span class="sd">    in the instance given above, channel names might be &quot;input&quot; for the upper</span>
<span class="sd">    level, and &quot;input.csv&quot;, &quot;input.xls&quot; and &quot;input.gnu&quot; for the sub-levels.</span>
<span class="sd">    There is no arbitrary limit to the depth of nesting.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="o">=</span><span class="n">NOTSET</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the logger with a name and an optional level.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">Filterer</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">=</span> <span class="n">name</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">level</span> <span class="o">=</span> <span class="n">_checkLevel</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">parent</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">propagate</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">handlers</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">disabled</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_cache</span> <span class="o">=</span> <span class="p">{}</span>

    <span class="k">def</span> <span class="nf">setLevel</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set the logging level of this logger.  level must be an int or a str.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">level</span> <span class="o">=</span> <span class="n">_checkLevel</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">_clear_cache</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">debug</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Log &#39;msg % args&#39; with severity &#39;DEBUG&#39;.</span>

<span class="sd">        To pass exception information, use the keyword argument exc_info with</span>
<span class="sd">        a true value, e.g.</span>

<span class="sd">        logger.debug(&quot;Houston, we have a %s&quot;, &quot;thorny problem&quot;, exc_info=1)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">DEBUG</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_log</span><span class="p">(</span><span class="n">DEBUG</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">info</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Log &#39;msg % args&#39; with severity &#39;INFO&#39;.</span>

<span class="sd">        To pass exception information, use the keyword argument exc_info with</span>
<span class="sd">        a true value, e.g.</span>

<span class="sd">        logger.info(&quot;Houston, we have a %s&quot;, &quot;interesting problem&quot;, exc_info=1)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">INFO</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_log</span><span class="p">(</span><span class="n">INFO</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">warning</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Log &#39;msg % args&#39; with severity &#39;WARNING&#39;.</span>

<span class="sd">        To pass exception information, use the keyword argument exc_info with</span>
<span class="sd">        a true value, e.g.</span>

<span class="sd">        logger.warning(&quot;Houston, we have a %s&quot;, &quot;bit of a problem&quot;, exc_info=1)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">WARNING</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_log</span><span class="p">(</span><span class="n">WARNING</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">warn</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;The &#39;warn&#39; method is deprecated, &quot;</span>
            <span class="s2">&quot;use &#39;warning&#39; instead&quot;</span><span class="p">,</span> <span class="ne">DeprecationWarning</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">error</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Log &#39;msg % args&#39; with severity &#39;ERROR&#39;.</span>

<span class="sd">        To pass exception information, use the keyword argument exc_info with</span>
<span class="sd">        a true value, e.g.</span>

<span class="sd">        logger.error(&quot;Houston, we have a %s&quot;, &quot;major problem&quot;, exc_info=1)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">ERROR</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_log</span><span class="p">(</span><span class="n">ERROR</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">exception</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Convenience method for logging an ERROR with exception information.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="n">exc_info</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">critical</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Log &#39;msg % args&#39; with severity &#39;CRITICAL&#39;.</span>

<span class="sd">        To pass exception information, use the keyword argument exc_info with</span>
<span class="sd">        a true value, e.g.</span>

<span class="sd">        logger.critical(&quot;Houston, we have a %s&quot;, &quot;major disaster&quot;, exc_info=1)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">CRITICAL</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_log</span><span class="p">(</span><span class="n">CRITICAL</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="n">fatal</span> <span class="o">=</span> <span class="n">critical</span>

    <span class="k">def</span> <span class="nf">log</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Log &#39;msg % args&#39; with the integer severity &#39;level&#39;.</span>

<span class="sd">        To pass exception information, use the keyword argument exc_info with</span>
<span class="sd">        a true value, e.g.</span>

<span class="sd">        logger.log(level, &quot;We have a %s&quot;, &quot;mysterious problem&quot;, exc_info=1)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">raiseExceptions</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s2">&quot;level must be an integer&quot;</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">return</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">level</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_log</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">findCaller</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">stack_info</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">stacklevel</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Find the stack frame of the caller so that we can note the source</span>
<span class="sd">        file name, line number and function name.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">f</span> <span class="o">=</span> <span class="n">currentframe</span><span class="p">()</span>
        <span class="c1">#On some versions of IronPython, currentframe() returns None if</span>
        <span class="c1">#IronPython isn&#39;t run with -X:Frames.</span>
        <span class="k">if</span> <span class="n">f</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">f</span> <span class="o">=</span> <span class="n">f</span><span class="o">.</span><span class="n">f_back</span>
        <span class="n">orig_f</span> <span class="o">=</span> <span class="n">f</span>
        <span class="k">while</span> <span class="n">f</span> <span class="ow">and</span> <span class="n">stacklevel</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">f</span> <span class="o">=</span> <span class="n">f</span><span class="o">.</span><span class="n">f_back</span>
            <span class="n">stacklevel</span> <span class="o">-=</span> <span class="mi">1</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">f</span><span class="p">:</span>
            <span class="n">f</span> <span class="o">=</span> <span class="n">orig_f</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="s2">&quot;(unknown file)&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;(unknown function)&quot;</span><span class="p">,</span> <span class="kc">None</span>
        <span class="k">while</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="s2">&quot;f_code&quot;</span><span class="p">):</span>
            <span class="n">co</span> <span class="o">=</span> <span class="n">f</span><span class="o">.</span><span class="n">f_code</span>
            <span class="n">filename</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">normcase</span><span class="p">(</span><span class="n">co</span><span class="o">.</span><span class="n">co_filename</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">filename</span> <span class="o">==</span> <span class="n">_srcfile</span><span class="p">:</span>
                <span class="n">f</span> <span class="o">=</span> <span class="n">f</span><span class="o">.</span><span class="n">f_back</span>
                <span class="k">continue</span>
            <span class="n">sinfo</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">if</span> <span class="n">stack_info</span><span class="p">:</span>
                <span class="n">sio</span> <span class="o">=</span> <span class="n">io</span><span class="o">.</span><span class="n">StringIO</span><span class="p">()</span>
                <span class="n">sio</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;Stack (most recent call last):</span><span class="se">\n</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="n">traceback</span><span class="o">.</span><span class="n">print_stack</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="n">file</span><span class="o">=</span><span class="n">sio</span><span class="p">)</span>
                <span class="n">sinfo</span> <span class="o">=</span> <span class="n">sio</span><span class="o">.</span><span class="n">getvalue</span><span class="p">()</span>
                <span class="k">if</span> <span class="n">sinfo</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="s1">&#39;</span><span class="se">\n</span><span class="s1">&#39;</span><span class="p">:</span>
                    <span class="n">sinfo</span> <span class="o">=</span> <span class="n">sinfo</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
                <span class="n">sio</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
            <span class="n">rv</span> <span class="o">=</span> <span class="p">(</span><span class="n">co</span><span class="o">.</span><span class="n">co_filename</span><span class="p">,</span> <span class="n">f</span><span class="o">.</span><span class="n">f_lineno</span><span class="p">,</span> <span class="n">co</span><span class="o">.</span><span class="n">co_name</span><span class="p">,</span> <span class="n">sinfo</span><span class="p">)</span>
            <span class="k">break</span>
        <span class="k">return</span> <span class="n">rv</span>

    <span class="k">def</span> <span class="nf">makeRecord</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">lno</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="p">,</span>
                   <span class="n">func</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">extra</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">sinfo</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        A factory method which can be overridden in subclasses to create</span>
<span class="sd">        specialized LogRecords.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="n">_logRecordFactory</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">lno</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="p">,</span> <span class="n">func</span><span class="p">,</span>
                             <span class="n">sinfo</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">extra</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="n">extra</span><span class="p">:</span>
                <span class="k">if</span> <span class="p">(</span><span class="n">key</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;message&quot;</span><span class="p">,</span> <span class="s2">&quot;asctime&quot;</span><span class="p">])</span> <span class="ow">or</span> <span class="p">(</span><span class="n">key</span> <span class="ow">in</span> <span class="n">rv</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">KeyError</span><span class="p">(</span><span class="s2">&quot;Attempt to overwrite </span><span class="si">%r</span><span class="s2"> in LogRecord&quot;</span> <span class="o">%</span> <span class="n">key</span><span class="p">)</span>
                <span class="n">rv</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">extra</span><span class="p">[</span><span class="n">key</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">rv</span>

    <span class="k">def</span> <span class="nf">_log</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">extra</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">stack_info</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
             <span class="n">stacklevel</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Low-level logging routine which creates a LogRecord and then calls</span>
<span class="sd">        all the handlers of this logger to handle the record.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">sinfo</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="n">_srcfile</span><span class="p">:</span>
            <span class="c1">#IronPython doesn&#39;t track Python frames, so findCaller raises an</span>
            <span class="c1">#exception on some versions of IronPython. We trap it here so that</span>
            <span class="c1">#IronPython can use logging.</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="n">fn</span><span class="p">,</span> <span class="n">lno</span><span class="p">,</span> <span class="n">func</span><span class="p">,</span> <span class="n">sinfo</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">findCaller</span><span class="p">(</span><span class="n">stack_info</span><span class="p">,</span> <span class="n">stacklevel</span><span class="p">)</span>
            <span class="k">except</span> <span class="ne">ValueError</span><span class="p">:</span> <span class="c1"># pragma: no cover</span>
                <span class="n">fn</span><span class="p">,</span> <span class="n">lno</span><span class="p">,</span> <span class="n">func</span> <span class="o">=</span> <span class="s2">&quot;(unknown file)&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;(unknown function)&quot;</span>
        <span class="k">else</span><span class="p">:</span> <span class="c1"># pragma: no cover</span>
            <span class="n">fn</span><span class="p">,</span> <span class="n">lno</span><span class="p">,</span> <span class="n">func</span> <span class="o">=</span> <span class="s2">&quot;(unknown file)&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;(unknown function)&quot;</span>
        <span class="k">if</span> <span class="n">exc_info</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">exc_info</span><span class="p">,</span> <span class="ne">BaseException</span><span class="p">):</span>
                <span class="n">exc_info</span> <span class="o">=</span> <span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">exc_info</span><span class="p">),</span> <span class="n">exc_info</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">.</span><span class="n">__traceback__</span><span class="p">)</span>
            <span class="k">elif</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">exc_info</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
                <span class="n">exc_info</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">exc_info</span><span class="p">()</span>
        <span class="n">record</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">makeRecord</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">fn</span><span class="p">,</span> <span class="n">lno</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span>
                                 <span class="n">exc_info</span><span class="p">,</span> <span class="n">func</span><span class="p">,</span> <span class="n">extra</span><span class="p">,</span> <span class="n">sinfo</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">handle</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Call the handlers for the specified record.</span>

<span class="sd">        This method is used for unpickled records received from a socket, as</span>
<span class="sd">        well as those created locally. Logger-level filtering is applied.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="p">(</span><span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">disabled</span><span class="p">)</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">record</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">callHandlers</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">addHandler</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">hdlr</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Add the specified handler to this logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">_acquireLock</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">hdlr</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">handlers</span><span class="p">):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">handlers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">hdlr</span><span class="p">)</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">_releaseLock</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">removeHandler</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">hdlr</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Remove the specified handler from this logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">_acquireLock</span><span class="p">()</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">hdlr</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">handlers</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">handlers</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">hdlr</span><span class="p">)</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">_releaseLock</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">hasHandlers</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        See if this logger has any handlers configured.</span>

<span class="sd">        Loop through all handlers for this logger and its parents in the</span>
<span class="sd">        logger hierarchy. Return True if a handler was found, else False.</span>
<span class="sd">        Stop searching up the hierarchy whenever a logger with the &quot;propagate&quot;</span>
<span class="sd">        attribute set to zero is found - that will be the last logger which</span>
<span class="sd">        is checked for the existence of handlers.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">c</span> <span class="o">=</span> <span class="bp">self</span>
        <span class="n">rv</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">while</span> <span class="n">c</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">c</span><span class="o">.</span><span class="n">handlers</span><span class="p">:</span>
                <span class="n">rv</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="k">break</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">c</span><span class="o">.</span><span class="n">propagate</span><span class="p">:</span>
                <span class="k">break</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">c</span> <span class="o">=</span> <span class="n">c</span><span class="o">.</span><span class="n">parent</span>
        <span class="k">return</span> <span class="n">rv</span>

    <span class="k">def</span> <span class="nf">callHandlers</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Pass a record to all relevant handlers.</span>

<span class="sd">        Loop through all handlers for this logger and its parents in the</span>
<span class="sd">        logger hierarchy. If no handler was found, output a one-off error</span>
<span class="sd">        message to sys.stderr. Stop searching up the hierarchy whenever a</span>
<span class="sd">        logger with the &quot;propagate&quot; attribute set to zero is found - that</span>
<span class="sd">        will be the last logger whose handlers are called.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">c</span> <span class="o">=</span> <span class="bp">self</span>
        <span class="n">found</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">while</span> <span class="n">c</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">hdlr</span> <span class="ow">in</span> <span class="n">c</span><span class="o">.</span><span class="n">handlers</span><span class="p">:</span>
                <span class="n">found</span> <span class="o">=</span> <span class="n">found</span> <span class="o">+</span> <span class="mi">1</span>
                <span class="k">if</span> <span class="n">record</span><span class="o">.</span><span class="n">levelno</span> <span class="o">&gt;=</span> <span class="n">hdlr</span><span class="o">.</span><span class="n">level</span><span class="p">:</span>
                    <span class="n">hdlr</span><span class="o">.</span><span class="n">handle</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">c</span><span class="o">.</span><span class="n">propagate</span><span class="p">:</span>
                <span class="n">c</span> <span class="o">=</span> <span class="kc">None</span>    <span class="c1">#break out</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">c</span> <span class="o">=</span> <span class="n">c</span><span class="o">.</span><span class="n">parent</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">found</span> <span class="o">==</span> <span class="mi">0</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">lastResort</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">record</span><span class="o">.</span><span class="n">levelno</span> <span class="o">&gt;=</span> <span class="n">lastResort</span><span class="o">.</span><span class="n">level</span><span class="p">:</span>
                    <span class="n">lastResort</span><span class="o">.</span><span class="n">handle</span><span class="p">(</span><span class="n">record</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">raiseExceptions</span> <span class="ow">and</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">emittedNoHandlerWarning</span><span class="p">:</span>
                <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s2">&quot;No handlers could be found for logger&quot;</span>
                                 <span class="s2">&quot; </span><span class="se">\&quot;</span><span class="si">%s</span><span class="se">\&quot;\n</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">emittedNoHandlerWarning</span> <span class="o">=</span> <span class="kc">True</span>

    <span class="k">def</span> <span class="nf">getEffectiveLevel</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get the effective level for this logger.</span>

<span class="sd">        Loop through this logger and its parents in the logger hierarchy,</span>
<span class="sd">        looking for a non-zero logging level. Return the first one found.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">logger</span> <span class="o">=</span> <span class="bp">self</span>
        <span class="k">while</span> <span class="n">logger</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">logger</span><span class="o">.</span><span class="n">level</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">logger</span><span class="o">.</span><span class="n">level</span>
            <span class="n">logger</span> <span class="o">=</span> <span class="n">logger</span><span class="o">.</span><span class="n">parent</span>
        <span class="k">return</span> <span class="n">NOTSET</span>

    <span class="k">def</span> <span class="nf">isEnabledFor</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Is this logger enabled for level &#39;level&#39;?</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">disabled</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">False</span>

        <span class="k">try</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_cache</span><span class="p">[</span><span class="n">level</span><span class="p">]</span>
        <span class="k">except</span> <span class="ne">KeyError</span><span class="p">:</span>
            <span class="n">_acquireLock</span><span class="p">()</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">disable</span> <span class="o">&gt;=</span> <span class="n">level</span><span class="p">:</span>
                    <span class="n">is_enabled</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_cache</span><span class="p">[</span><span class="n">level</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">is_enabled</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_cache</span><span class="p">[</span><span class="n">level</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>
                        <span class="n">level</span> <span class="o">&gt;=</span> <span class="bp">self</span><span class="o">.</span><span class="n">getEffectiveLevel</span><span class="p">()</span>
                    <span class="p">)</span>
            <span class="k">finally</span><span class="p">:</span>
                <span class="n">_releaseLock</span><span class="p">()</span>
            <span class="k">return</span> <span class="n">is_enabled</span>

    <span class="k">def</span> <span class="nf">getChild</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">suffix</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get a logger which is a descendant to this one.</span>

<span class="sd">        This is a convenience method, such that</span>

<span class="sd">        logging.getLogger(&#39;abc&#39;).getChild(&#39;def.ghi&#39;)</span>

<span class="sd">        is the same as</span>

<span class="sd">        logging.getLogger(&#39;abc.def.ghi&#39;)</span>

<span class="sd">        It&#39;s useful, for example, when the parent logger is named using</span>
<span class="sd">        __name__ rather than a literal string.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">root</span> <span class="ow">is</span> <span class="ow">not</span> <span class="bp">self</span><span class="p">:</span>
            <span class="n">suffix</span> <span class="o">=</span> <span class="s1">&#39;.&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="n">suffix</span><span class="p">))</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="n">suffix</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">level</span> <span class="o">=</span> <span class="n">getLevelName</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">getEffectiveLevel</span><span class="p">())</span>
        <span class="k">return</span> <span class="s1">&#39;&lt;</span><span class="si">%s</span><span class="s1"> </span><span class="si">%s</span><span class="s1"> (</span><span class="si">%s</span><span class="s1">)&gt;&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">__reduce__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="c1"># In general, only the root logger will not be accessible via its name.</span>
        <span class="c1"># However, the root logger&#39;s class has its own __reduce__ method.</span>
        <span class="k">if</span> <span class="n">getLogger</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">)</span> <span class="ow">is</span> <span class="ow">not</span> <span class="bp">self</span><span class="p">:</span>
            <span class="kn">import</span> <span class="nn">pickle</span>
            <span class="k">raise</span> <span class="n">pickle</span><span class="o">.</span><span class="n">PicklingError</span><span class="p">(</span><span class="s1">&#39;logger cannot be pickled&#39;</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">getLogger</span><span class="p">,</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,)</span>


<span class="k">class</span> <span class="nc">RootLogger</span><span class="p">(</span><span class="n">Logger</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A root logger is not that different to any other logger, except that</span>
<span class="sd">    it must have a logging level and there is only one instance of it in</span>
<span class="sd">    the hierarchy.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the logger with the name &quot;root&quot;.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">Logger</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;root&quot;</span><span class="p">,</span> <span class="n">level</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">__reduce__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">getLogger</span><span class="p">,</span> <span class="p">()</span>

<span class="n">_loggerClass</span> <span class="o">=</span> <span class="n">Logger</span>

<span class="k">class</span> <span class="nc">LoggerAdapter</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    An adapter for loggers which makes it easier to specify contextual</span>
<span class="sd">    information in logging output.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">logger</span><span class="p">,</span> <span class="n">extra</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Initialize the adapter with a logger and a dict-like object which</span>
<span class="sd">        provides contextual information. This constructor signature allows</span>
<span class="sd">        easy stacking of LoggerAdapters, if so desired.</span>

<span class="sd">        You can effectively pass keyword arguments as shown in the</span>
<span class="sd">        following example:</span>

<span class="sd">        adapter = LoggerAdapter(someLogger, dict(p1=v1, p2=&quot;v2&quot;))</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logger</span> <span class="o">=</span> <span class="n">logger</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">extra</span> <span class="o">=</span> <span class="n">extra</span>

    <span class="k">def</span> <span class="nf">process</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Process the logging message and keyword arguments passed in to</span>
<span class="sd">        a logging call to insert contextual information. You can either</span>
<span class="sd">        manipulate the message itself, the keyword args or both. Return</span>
<span class="sd">        the message and kwargs modified (or not) to suit your needs.</span>

<span class="sd">        Normally, you&#39;ll only need to override this one method in a</span>
<span class="sd">        LoggerAdapter subclass for your specific needs.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;extra&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">extra</span>
        <span class="k">return</span> <span class="n">msg</span><span class="p">,</span> <span class="n">kwargs</span>

    <span class="c1">#</span>
    <span class="c1"># Boilerplate convenience methods</span>
    <span class="c1">#</span>
    <span class="k">def</span> <span class="nf">debug</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delegate a debug call to the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">DEBUG</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">info</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delegate an info call to the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">INFO</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">warning</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delegate a warning call to the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">WARNING</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">warn</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;The &#39;warn&#39; method is deprecated, &quot;</span>
            <span class="s2">&quot;use &#39;warning&#39; instead&quot;</span><span class="p">,</span> <span class="ne">DeprecationWarning</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">error</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delegate an error call to the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">ERROR</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">exception</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delegate an exception call to the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">ERROR</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="n">exc_info</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">critical</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delegate a critical call to the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">CRITICAL</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">log</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delegate a log call to the underlying logger, after adding</span>
<span class="sd">        contextual information from this adapter instance.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">level</span><span class="p">):</span>
            <span class="n">msg</span><span class="p">,</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">process</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">isEnabledFor</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Is this logger enabled for level &#39;level&#39;?</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">isEnabledFor</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">setLevel</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Set the specified level on the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">setLevel</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">getEffectiveLevel</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get the effective level for the underlying logger.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">getEffectiveLevel</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">hasHandlers</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        See if the underlying logger has any handlers.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">hasHandlers</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_log</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">extra</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">stack_info</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Low-level log implementation, proxied to allow nested logger adapters.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">_log</span><span class="p">(</span>
            <span class="n">level</span><span class="p">,</span>
            <span class="n">msg</span><span class="p">,</span>
            <span class="n">args</span><span class="p">,</span>
            <span class="n">exc_info</span><span class="o">=</span><span class="n">exc_info</span><span class="p">,</span>
            <span class="n">extra</span><span class="o">=</span><span class="n">extra</span><span class="p">,</span>
            <span class="n">stack_info</span><span class="o">=</span><span class="n">stack_info</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">manager</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">manager</span>

    <span class="nd">@manager</span><span class="o">.</span><span class="n">setter</span>
    <span class="k">def</span> <span class="nf">manager</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">manager</span> <span class="o">=</span> <span class="n">value</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">name</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">name</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">logger</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">logger</span>
        <span class="n">level</span> <span class="o">=</span> <span class="n">getLevelName</span><span class="p">(</span><span class="n">logger</span><span class="o">.</span><span class="n">getEffectiveLevel</span><span class="p">())</span>
        <span class="k">return</span> <span class="s1">&#39;&lt;</span><span class="si">%s</span><span class="s1"> </span><span class="si">%s</span><span class="s1"> (</span><span class="si">%s</span><span class="s1">)&gt;&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> <span class="n">logger</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="n">level</span><span class="p">)</span>

<span class="n">root</span> <span class="o">=</span> <span class="n">RootLogger</span><span class="p">(</span><span class="n">WARNING</span><span class="p">)</span>
<span class="n">Logger</span><span class="o">.</span><span class="n">root</span> <span class="o">=</span> <span class="n">root</span>
<span class="n">Logger</span><span class="o">.</span><span class="n">manager</span> <span class="o">=</span> <span class="n">Manager</span><span class="p">(</span><span class="n">Logger</span><span class="o">.</span><span class="n">root</span><span class="p">)</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1"># Configuration classes and functions</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="k">def</span> <span class="nf">basicConfig</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Do basic configuration for the logging system.</span>

<span class="sd">    This function does nothing if the root logger already has handlers</span>
<span class="sd">    configured, unless the keyword argument *force* is set to ``True``.</span>
<span class="sd">    It is a convenience method intended for use by simple scripts</span>
<span class="sd">    to do one-shot configuration of the logging package.</span>

<span class="sd">    The default behaviour is to create a StreamHandler which writes to</span>
<span class="sd">    sys.stderr, set a formatter using the BASIC_FORMAT format string, and</span>
<span class="sd">    add the handler to the root logger.</span>

<span class="sd">    A number of optional keyword arguments may be specified, which can alter</span>
<span class="sd">    the default behaviour.</span>

<span class="sd">    filename  Specifies that a FileHandler be created, using the specified</span>
<span class="sd">              filename, rather than a StreamHandler.</span>
<span class="sd">    filemode  Specifies the mode to open the file, if filename is specified</span>
<span class="sd">              (if filemode is unspecified, it defaults to &#39;a&#39;).</span>
<span class="sd">    format    Use the specified format string for the handler.</span>
<span class="sd">    datefmt   Use the specified date/time format.</span>
<span class="sd">    style     If a format string is specified, use this to specify the</span>
<span class="sd">              type of format string (possible values &#39;%&#39;, &#39;{&#39;, &#39;$&#39;, for</span>
<span class="sd">              %-formatting, :meth:`str.format` and :class:`string.Template`</span>
<span class="sd">              - defaults to &#39;%&#39;).</span>
<span class="sd">    level     Set the root logger level to the specified level.</span>
<span class="sd">    stream    Use the specified stream to initialize the StreamHandler. Note</span>
<span class="sd">              that this argument is incompatible with &#39;filename&#39; - if both</span>
<span class="sd">              are present, &#39;stream&#39; is ignored.</span>
<span class="sd">    handlers  If specified, this should be an iterable of already created</span>
<span class="sd">              handlers, which will be added to the root handler. Any handler</span>
<span class="sd">              in the list which does not have a formatter assigned will be</span>
<span class="sd">              assigned the formatter created in this function.</span>
<span class="sd">    force     If this keyword  is specified as true, any existing handlers</span>
<span class="sd">              attached to the root logger are removed and closed, before</span>
<span class="sd">              carrying out the configuration as specified by the other</span>
<span class="sd">              arguments.</span>
<span class="sd">    encoding  If specified together with a filename, this encoding is passed to</span>
<span class="sd">              the created FileHandler, causing it to be used when the file is</span>
<span class="sd">              opened.</span>
<span class="sd">    errors    If specified together with a filename, this value is passed to the</span>
<span class="sd">              created FileHandler, causing it to be used when the file is</span>
<span class="sd">              opened in text mode. If not specified, the default value is</span>
<span class="sd">              `backslashreplace`.</span>

<span class="sd">    Note that you could specify a stream created using open(filename, mode)</span>
<span class="sd">    rather than passing the filename and mode in. However, it should be</span>
<span class="sd">    remembered that StreamHandler does not close its stream (since it may be</span>
<span class="sd">    using sys.stdout or sys.stderr), whereas FileHandler closes its stream</span>
<span class="sd">    when the handler is closed.</span>

<span class="sd">    .. versionchanged:: 3.2</span>
<span class="sd">       Added the ``style`` parameter.</span>

<span class="sd">    .. versionchanged:: 3.3</span>
<span class="sd">       Added the ``handlers`` parameter. A ``ValueError`` is now thrown for</span>
<span class="sd">       incompatible arguments (e.g. ``handlers`` specified together with</span>
<span class="sd">       ``filename``/``filemode``, or ``filename``/``filemode`` specified</span>
<span class="sd">       together with ``stream``, or ``handlers`` specified together with</span>
<span class="sd">       ``stream``.</span>

<span class="sd">    .. versionchanged:: 3.8</span>
<span class="sd">       Added the ``force`` parameter.</span>

<span class="sd">    .. versionchanged:: 3.9</span>
<span class="sd">       Added the ``encoding`` and ``errors`` parameters.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># Add thread safety in case someone mistakenly calls</span>
    <span class="c1"># basicConfig() from multiple threads</span>
    <span class="n">_acquireLock</span><span class="p">()</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">force</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s1">&#39;force&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="n">encoding</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s1">&#39;encoding&#39;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
        <span class="n">errors</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s1">&#39;errors&#39;</span><span class="p">,</span> <span class="s1">&#39;backslashreplace&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">force</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">h</span> <span class="ow">in</span> <span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">[:]:</span>
                <span class="n">root</span><span class="o">.</span><span class="n">removeHandler</span><span class="p">(</span><span class="n">h</span><span class="p">)</span>
                <span class="n">h</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">handlers</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;handlers&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">handlers</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="k">if</span> <span class="s2">&quot;stream&quot;</span> <span class="ow">in</span> <span class="n">kwargs</span> <span class="ow">and</span> <span class="s2">&quot;filename&quot;</span> <span class="ow">in</span> <span class="n">kwargs</span><span class="p">:</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;&#39;stream&#39; and &#39;filename&#39; should not be &quot;</span>
                                     <span class="s2">&quot;specified together&quot;</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">if</span> <span class="s2">&quot;stream&quot;</span> <span class="ow">in</span> <span class="n">kwargs</span> <span class="ow">or</span> <span class="s2">&quot;filename&quot;</span> <span class="ow">in</span> <span class="n">kwargs</span><span class="p">:</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;&#39;stream&#39; or &#39;filename&#39; should not be &quot;</span>
                                     <span class="s2">&quot;specified together with &#39;handlers&#39;&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">handlers</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">filename</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;filename&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
                <span class="n">mode</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;filemode&quot;</span><span class="p">,</span> <span class="s1">&#39;a&#39;</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">filename</span><span class="p">:</span>
                    <span class="k">if</span> <span class="s1">&#39;b&#39;</span><span class="ow">in</span> <span class="n">mode</span><span class="p">:</span>
                        <span class="n">errors</span> <span class="o">=</span> <span class="kc">None</span>
                    <span class="n">h</span> <span class="o">=</span> <span class="n">FileHandler</span><span class="p">(</span><span class="n">filename</span><span class="p">,</span> <span class="n">mode</span><span class="p">,</span>
                                    <span class="n">encoding</span><span class="o">=</span><span class="n">encoding</span><span class="p">,</span> <span class="n">errors</span><span class="o">=</span><span class="n">errors</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">stream</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;stream&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
                    <span class="n">h</span> <span class="o">=</span> <span class="n">StreamHandler</span><span class="p">(</span><span class="n">stream</span><span class="p">)</span>
                <span class="n">handlers</span> <span class="o">=</span> <span class="p">[</span><span class="n">h</span><span class="p">]</span>
            <span class="n">dfs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;datefmt&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
            <span class="n">style</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;style&quot;</span><span class="p">,</span> <span class="s1">&#39;%&#39;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">style</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">_STYLES</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Style must be one of: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="s1">&#39;,&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span>
                                 <span class="n">_STYLES</span><span class="o">.</span><span class="n">keys</span><span class="p">()))</span>
            <span class="n">fs</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;format&quot;</span><span class="p">,</span> <span class="n">_STYLES</span><span class="p">[</span><span class="n">style</span><span class="p">][</span><span class="mi">1</span><span class="p">])</span>
            <span class="n">fmt</span> <span class="o">=</span> <span class="n">Formatter</span><span class="p">(</span><span class="n">fs</span><span class="p">,</span> <span class="n">dfs</span><span class="p">,</span> <span class="n">style</span><span class="p">)</span>
            <span class="k">for</span> <span class="n">h</span> <span class="ow">in</span> <span class="n">handlers</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">h</span><span class="o">.</span><span class="n">formatter</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">h</span><span class="o">.</span><span class="n">setFormatter</span><span class="p">(</span><span class="n">fmt</span><span class="p">)</span>
                <span class="n">root</span><span class="o">.</span><span class="n">addHandler</span><span class="p">(</span><span class="n">h</span><span class="p">)</span>
            <span class="n">level</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;level&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">level</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">root</span><span class="o">.</span><span class="n">setLevel</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">kwargs</span><span class="p">:</span>
                <span class="n">keys</span> <span class="o">=</span> <span class="s1">&#39;, &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">kwargs</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Unrecognised argument(s): </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">keys</span><span class="p">)</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="n">_releaseLock</span><span class="p">()</span>

<span class="c1">#---------------------------------------------------------------------------</span>
<span class="c1"># Utility functions at module level.</span>
<span class="c1"># Basically delegate everything to the root logger.</span>
<span class="c1">#---------------------------------------------------------------------------</span>

<span class="k">def</span> <span class="nf">getLogger</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Return a logger with the specified name, creating it if necessary.</span>

<span class="sd">    If no name is specified, return the root logger.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">name</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="nb">str</span><span class="p">)</span> <span class="ow">and</span> <span class="n">name</span> <span class="o">==</span> <span class="n">root</span><span class="o">.</span><span class="n">name</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">root</span>
    <span class="k">return</span> <span class="n">Logger</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">critical</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Log a message with severity &#39;CRITICAL&#39; on the root logger. If the logger</span>
<span class="sd">    has no handlers, call basicConfig() to add a console handler with a</span>
<span class="sd">    pre-defined format.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">basicConfig</span><span class="p">()</span>
    <span class="n">root</span><span class="o">.</span><span class="n">critical</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="n">fatal</span> <span class="o">=</span> <span class="n">critical</span>

<span class="k">def</span> <span class="nf">error</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Log a message with severity &#39;ERROR&#39; on the root logger. If the logger has</span>
<span class="sd">    no handlers, call basicConfig() to add a console handler with a pre-defined</span>
<span class="sd">    format.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">basicConfig</span><span class="p">()</span>
    <span class="n">root</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">exception</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Log a message with severity &#39;ERROR&#39; on the root logger, with exception</span>
<span class="sd">    information. If the logger has no handlers, basicConfig() is called to add</span>
<span class="sd">    a console handler with a pre-defined format.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">error</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="n">exc_info</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">warning</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Log a message with severity &#39;WARNING&#39; on the root logger. If the logger has</span>
<span class="sd">    no handlers, call basicConfig() to add a console handler with a pre-defined</span>
<span class="sd">    format.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">basicConfig</span><span class="p">()</span>
    <span class="n">root</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">warn</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;The &#39;warn&#39; function is deprecated, &quot;</span>
        <span class="s2">&quot;use &#39;warning&#39; instead&quot;</span><span class="p">,</span> <span class="ne">DeprecationWarning</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
    <span class="n">warning</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">info</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Log a message with severity &#39;INFO&#39; on the root logger. If the logger has</span>
<span class="sd">    no handlers, call basicConfig() to add a console handler with a pre-defined</span>
<span class="sd">    format.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">basicConfig</span><span class="p">()</span>
    <span class="n">root</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">debug</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Log a message with severity &#39;DEBUG&#39; on the root logger. If the logger has</span>
<span class="sd">    no handlers, call basicConfig() to add a console handler with a pre-defined</span>
<span class="sd">    format.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">basicConfig</span><span class="p">()</span>
    <span class="n">root</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">log</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Log &#39;msg % args&#39; with the integer severity &#39;level&#39; on the root logger. If</span>
<span class="sd">    the logger has no handlers, call basicConfig() to add a console handler</span>
<span class="sd">    with a pre-defined format.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">basicConfig</span><span class="p">()</span>
    <span class="n">root</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">msg</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">disable</span><span class="p">(</span><span class="n">level</span><span class="o">=</span><span class="n">CRITICAL</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Disable all logging calls of severity &#39;level&#39; and below.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">root</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">disable</span> <span class="o">=</span> <span class="n">level</span>
    <span class="n">root</span><span class="o">.</span><span class="n">manager</span><span class="o">.</span><span class="n">_clear_cache</span><span class="p">()</span>

<span class="k">def</span> <span class="nf">shutdown</span><span class="p">(</span><span class="n">handlerList</span><span class="o">=</span><span class="n">_handlerList</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Perform any cleanup actions in the logging system (e.g. flushing</span>
<span class="sd">    buffers).</span>

<span class="sd">    Should be called at application exit.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">for</span> <span class="n">wr</span> <span class="ow">in</span> <span class="nb">reversed</span><span class="p">(</span><span class="n">handlerList</span><span class="p">[:]):</span>
        <span class="c1">#errors might occur, for example, if files are locked</span>
        <span class="c1">#we just ignore them if raiseExceptions is not set</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">h</span> <span class="o">=</span> <span class="n">wr</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">h</span><span class="p">:</span>
                <span class="k">try</span><span class="p">:</span>
                    <span class="n">h</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
                    <span class="n">h</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
                    <span class="n">h</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
                <span class="k">except</span> <span class="p">(</span><span class="ne">OSError</span><span class="p">,</span> <span class="ne">ValueError</span><span class="p">):</span>
                    <span class="c1"># Ignore errors which might be caused</span>
                    <span class="c1"># because handlers have been closed but</span>
                    <span class="c1"># references to them are still around at</span>
                    <span class="c1"># application exit.</span>
                    <span class="k">pass</span>
                <span class="k">finally</span><span class="p">:</span>
                    <span class="n">h</span><span class="o">.</span><span class="n">release</span><span class="p">()</span>
        <span class="k">except</span><span class="p">:</span> <span class="c1"># ignore everything, as we&#39;re shutting down</span>
            <span class="k">if</span> <span class="n">raiseExceptions</span><span class="p">:</span>
                <span class="k">raise</span>
            <span class="c1">#else, swallow</span>

<span class="c1">#Let&#39;s try and shutdown automatically on application exit...</span>
<span class="kn">import</span> <span class="nn">atexit</span>
<span class="n">atexit</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="n">shutdown</span><span class="p">)</span>

<span class="c1"># Null handler</span>

<span class="k">class</span> <span class="nc">NullHandler</span><span class="p">(</span><span class="n">Handler</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This handler does nothing. It&#39;s intended to be used to avoid the</span>
<span class="sd">    &quot;No handlers could be found for logger XXX&quot; one-off warning. This is</span>
<span class="sd">    important for library code, which may contain code to log events. If a user</span>
<span class="sd">    of the library does not configure logging, the one-off warning might be</span>
<span class="sd">    produced; to avoid this, the library developer simply needs to instantiate</span>
<span class="sd">    a NullHandler and add it to the top-level logger of the library module or</span>
<span class="sd">    package.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="nf">handle</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Stub.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">emit</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">record</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Stub.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">createLock</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lock</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="nf">_at_fork_reinit</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">pass</span>

<span class="c1"># Warnings integration</span>

<span class="n">_warnings_showwarning</span> <span class="o">=</span> <span class="kc">None</span>

<span class="k">def</span> <span class="nf">_showwarning</span><span class="p">(</span><span class="n">message</span><span class="p">,</span> <span class="n">category</span><span class="p">,</span> <span class="n">filename</span><span class="p">,</span> <span class="n">lineno</span><span class="p">,</span> <span class="n">file</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">line</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Implementation of showwarnings which redirects to logging, which will first</span>
<span class="sd">    check to see if the file parameter is None. If a file is specified, it will</span>
<span class="sd">    delegate to the original warnings implementation of showwarning. Otherwise,</span>
<span class="sd">    it will call warnings.formatwarning and will log the resulting string to a</span>
<span class="sd">    warnings logger named &quot;py.warnings&quot; with level logging.WARNING.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">file</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">_warnings_showwarning</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">_warnings_showwarning</span><span class="p">(</span><span class="n">message</span><span class="p">,</span> <span class="n">category</span><span class="p">,</span> <span class="n">filename</span><span class="p">,</span> <span class="n">lineno</span><span class="p">,</span> <span class="n">file</span><span class="p">,</span> <span class="n">line</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">s</span> <span class="o">=</span> <span class="n">warnings</span><span class="o">.</span><span class="n">formatwarning</span><span class="p">(</span><span class="n">message</span><span class="p">,</span> <span class="n">category</span><span class="p">,</span> <span class="n">filename</span><span class="p">,</span> <span class="n">lineno</span><span class="p">,</span> <span class="n">line</span><span class="p">)</span>
        <span class="n">logger</span> <span class="o">=</span> <span class="n">getLogger</span><span class="p">(</span><span class="s2">&quot;py.warnings&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">logger</span><span class="o">.</span><span class="n">handlers</span><span class="p">:</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">addHandler</span><span class="p">(</span><span class="n">NullHandler</span><span class="p">())</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">s</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">captureWarnings</span><span class="p">(</span><span class="n">capture</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    If capture is true, redirect all warnings to the logging package.</span>
<span class="sd">    If capture is False, ensure that warnings are not redirected to logging</span>
<span class="sd">    but to their original destinations.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">global</span> <span class="n">_warnings_showwarning</span>
    <span class="k">if</span> <span class="n">capture</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">_warnings_showwarning</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">_warnings_showwarning</span> <span class="o">=</span> <span class="n">warnings</span><span class="o">.</span><span class="n">showwarning</span>
            <span class="n">warnings</span><span class="o">.</span><span class="n">showwarning</span> <span class="o">=</span> <span class="n">_showwarning</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">_warnings_showwarning</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">warnings</span><span class="o">.</span><span class="n">showwarning</span> <span class="o">=</span> <span class="n">_warnings_showwarning</span>
            <span class="n">_warnings_showwarning</span> <span class="o">=</span> <span class="kc">None</span>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.common.abstractions.abstract_logger &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.common.abstractions.abstract_logger</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.common.abstractions.abstract_logger</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">logging</span>
<span class="kn">import</span> <span class="nn">logging.config</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.auto_logging</span> <span class="kn">import</span> <span class="n">AutoLoggerConfig</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.environment.environment_config</span> <span class="kn">import</span> <span class="n">DEFAULT_LOGGING_LEVEL</span>


<div class="viewcode-block" id="get_logger"><a class="viewcode-back" href="../../../../super_gradients.common.abstractions.html#super_gradients.common.abstractions.abstract_logger.get_logger">[docs]</a><span class="k">def</span> <span class="nf">get_logger</span><span class="p">(</span>
    <span class="n">logger_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">training_log_path</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">logs_dir_path</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">log_level</span><span class="o">=</span><span class="n">DEFAULT_LOGGING_LEVEL</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">logging</span><span class="o">.</span><span class="n">Logger</span><span class="p">:</span>
    <span class="n">config_dict</span> <span class="o">=</span> <span class="n">AutoLoggerConfig</span><span class="o">.</span><span class="n">generate_config_for_module_name</span><span class="p">(</span>
        <span class="n">module_name</span><span class="o">=</span><span class="n">logger_name</span><span class="p">,</span> <span class="n">training_log_path</span><span class="o">=</span><span class="n">training_log_path</span><span class="p">,</span> <span class="n">logs_dir_path</span><span class="o">=</span><span class="n">logs_dir_path</span><span class="p">,</span> <span class="n">log_level</span><span class="o">=</span><span class="n">log_level</span>
    <span class="p">)</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">config</span><span class="o">.</span><span class="n">dictConfig</span><span class="p">(</span><span class="n">config_dict</span><span class="p">)</span>
    <span class="n">logger</span><span class="p">:</span> <span class="n">logging</span><span class="o">.</span><span class="n">Logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="n">logger_name</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">logger</span></div>


<div class="viewcode-block" id="ILogger"><a class="viewcode-back" href="../../../../super_gradients.common.abstractions.html#super_gradients.common.abstractions.abstract_logger.ILogger">[docs]</a><span class="k">class</span> <span class="nc">ILogger</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Provides logging capabilities to the derived class.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">logger_name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
        <span class="n">logger_name</span> <span class="o">=</span> <span class="n">logger_name</span> <span class="k">if</span> <span class="n">logger_name</span> <span class="k">else</span> <span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__module__</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_logger</span><span class="p">:</span> <span class="n">logging</span><span class="o">.</span><span class="n">Logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="n">logger_name</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.auto_logging.auto_logger &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.auto_logging.auto_logger &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -86,116 +88,130 @@
 
                                        <div itemprop="articleBody">
                
 
                               <h1>Source code for super_gradients.common.auto_logging.auto_logger</h1><div class="highlight"><pre>
                
 
                            -<span></span><span class="kn">import</span> <span class="nn">json</span>
                
 
                            +<span></span><span class="kn">import</span> <span class="nn">logging</span>
                
 
                             <span class="kn">import</span> <span class="nn">os</span>
                
 
                            +<span class="kn">import</span> <span class="nn">sys</span>
                
 
                            +<span class="kn">import</span> <span class="nn">time</span>
                
 
                            +<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>
                
 
                            -<span class="kn">import</span> <span class="nn">pkg_resources</span>
                
 
                            -
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.common.environment.environment_config</span> <span class="kn">import</span> <span class="n">DEFAULT_LOGGING_LEVEL</span>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="AutoLoggerConfig"><a class="viewcode-back" href="../../../../super_gradients.common.auto_logging.html#super_gradients.common.auto_logging.auto_logger.AutoLoggerConfig">[docs]</a><span class="k">class</span> <span class="nc">AutoLoggerConfig</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="AutoLoggerConfig"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AutoLoggerConfig">[docs]</a><span class="k">class</span> <span class="nc">AutoLoggerConfig</span><span class="p">:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">    A Class for the Automated Logging Config that is created from the JSON config file (auto_logging_conf)</span>
                
 
                            +<span class="sd">    A Class for the Automated Logging Config</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -<div class="viewcode-block" id="AutoLoggerConfig.generate_config_for_module_name"><a class="viewcode-back" href="../../../../super_gradients.common.auto_logging.html#super_gradients.common.auto_logging.auto_logger.AutoLoggerConfig.generate_config_for_module_name">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            -    <span class="k">def</span> <span class="nf">generate_config_for_module_name</span><span class="p">(</span>
                
 
                            -        <span class="n">module_name</span><span class="p">,</span>
                
 
                            -        <span class="n">training_log_path</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                
 
                            -        <span class="n">log_level</span><span class="o">=</span><span class="n">DEFAULT_LOGGING_LEVEL</span><span class="p">,</span>
                
 
                            -        <span class="n">max_bytes</span><span class="o">=</span><span class="mi">10485760</span><span class="p">,</span>
                
 
                            -        <span class="n">logs_dir_path</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                
 
                            -        <span class="n">handlers_list</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                
 
                            -    <span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
                
 
                            +    <span class="n">FILE_LOGGING_LEVEL</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;LOG_LEVEL&quot;</span><span class="p">,</span> <span class="s2">&quot;DEBUG&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">upper</span><span class="p">()</span>
                
 
                            +    <span class="n">CONSOLE_LOGGING_LEVEL</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;CONSOLE_LOG_LEVEL&quot;</span><span class="p">,</span> <span class="s2">&quot;INFO&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">upper</span><span class="p">()</span>
                
 
                            +
                
 
                            +    <span class="n">filename</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="kc">None</span><span class="p">]</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_setup_default_logging</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">log_level</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        generate_config_for_module_name - Returns a Config Dict For Logging</span>
                
 
                            -<span class="sd">            :param module_name:     The Python Module name to create auto_logging for</span>
                
 
                            -<span class="sd">            :param log_level:       Minimal log level to set for the new auto_logging</span>
                
 
                            -<span class="sd">            :param max_bytes:       Max size for the log file before rotation starts</span>
                
 
                            -<span class="sd">            :param handlers_list:    A list specifying the handlers (Console, etc..) - Better Leave Empty or None</span>
                
 
                            -<span class="sd">            :param training_log_path: Path to training log file which all modules of super_gradients will write to. Ignored</span>
                
 
                            -<span class="sd">             when set to None.</span>
                
 
                            -<span class="sd">            :param logs_dir_path: Path to sg_logs directory (default=None), where module logs will be saved. When set</span>
                
 
                            -<span class="sd">                to None- module logs will be saved in ~/sg_logs (created if path does not exist). Main use case is for</span>
                
 
                            -<span class="sd">                testing.</span>
                
 
                            -
                
 
                            -
                
 
                            -<span class="sd">            :return: python dict() with the new auto_logging for the module</span>
                
 
                            +<span class="sd">        Setup default logging configuration. Usually happens when app starts, and we don&#39;t have</span>
                
 
                            +<span class="sd">        experiment dir yet.</span>
                
 
                            +<span class="sd">        The default log directory will be `~/sg_logs`</span>
                
 
                            +<span class="sd">        :param log_level: The default log level to use. If None, uses LOG_LEVEL and CONSOLE_LOG_LEVEL environment vars.</span>
                
 
                            +<span class="sd">        :return: None</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="c1"># LOADING THE ORIGINAL ROOT CONFIG FILE</span>
                
 
                            -        <span class="n">conf_file_name</span> <span class="o">=</span> <span class="s2">&quot;auto_logging_conf.json&quot;</span>
                
 
                            -        <span class="n">conf_file_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span>
                
 
                            -            <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s2">&quot;super_gradients&quot;</span><span class="p">,</span> <span class="s2">&quot;/common/auto_logging/&quot;</span><span class="p">),</span> <span class="n">conf_file_name</span>
                
 
                            +        <span class="c1"># There is no _easy_ way to log all events to a single file, when using DDP or DataLoader with num_workers &gt; 1</span>
                
 
                            +        <span class="c1"># on Windows platform. In both these cases a multiple processes will be spawned and multiple logs may be created.</span>
                
 
                            +        <span class="c1"># Therefore the log file will have the parent PID to being able to discriminate the logs corresponding to a single run.</span>
                
 
                            +        <span class="n">timestamp</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">strftime</span><span class="p">(</span><span class="s2">&quot;%Y_%m_</span><span class="si">%d</span><span class="s2">_%H_%M_%S&quot;</span><span class="p">,</span> <span class="n">time</span><span class="o">.</span><span class="n">localtime</span><span class="p">())</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_setup_logging</span><span class="p">(</span>
                
 
                            +            <span class="n">filename</span><span class="o">=</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">expanduser</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;~/sg_logs/sg_logs_</span><span class="si">{</span><span class="n">os</span><span class="o">.</span><span class="n">getppid</span><span class="p">()</span><span class="si">}</span><span class="s2">_</span><span class="si">{</spa
                
 
                            +            <span class="n">copy_already_logged_messages</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                
 
                            +            <span class="n">filemode</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                
 
                            +            <span class="n">log_level</span><span class="o">=</span><span class="n">log_level</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_setup_logging</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">copy_already_logged_messages</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">filemode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span c
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Sets the logging configuration to store messages to specific file</span>
                
 
                            +<span class="sd">        :param filename: Output log file</span>
                
 
                            +<span class="sd">        :param filemode: Open mode for file</span>
                
 
                            +<span class="sd">        :param copy_already_logged_messages: Controls whether messages from previous log configuration should be copied</span>
                
 
                            +<span class="sd">               to new place. This is helpful to transfer diagnostic messages (from the app start) to experiment dir.</span>
                
 
                            +<span class="sd">        :param log_level: The default log level to use. If None, uses LOG_LEVEL and CONSOLE_LOG_LEVEL environment vars.</span>
                
 
                            +<span class="sd">        :return:</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">filename</span><span class="p">),</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">copy_already_logged_messages</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="bp">self</span><span cl
                
 
                            +            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="s2">&quot;r&quot;</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="s2">&quot;utf-8&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">src</span><span class="p">:</span>
                
 
                            +                <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">filename</span><span class="p">,</span> <span class="s2">&quot;w&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">dst</span><span class="p">:</span>
                
 
                            +                    <span class="n">dst</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">src</span><span class="o">.</span><span class="n">read</span><span class="p">())</span>
                
 
                            +
                
 
                            +        <span class="n">file_logging_level</span> <span class="o">=</span> <span class="n">log_level</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">FILE_LOGGING_LEVEL</span>
                
 
                            +        <span class="n">console_logging_level</span> <span class="o">=</span> <span class="n">log_level</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">CONSOLE_LOGGING_LEVEL</span>
                
 
                            +
                
 
                            +        <span class="n">cur_version</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">version_info</span>
                
 
                            +        <span class="n">python_38</span> <span class="o">=</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">8</span><span class="p">)</span>
                
 
                            +        <span class="n">python_39</span> <span class="o">=</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">9</span><span class="p">)</span>
                
 
                            +        <span class="n">manager</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="s2">&quot;&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">manager</span>
                
 
                            +
                
 
                            +        <span class="n">extra_kwargs</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                            +        <span class="k">if</span> <span class="n">cur_version</span> <span class="o">&gt;=</span> <span class="n">python_38</span><span class="p">:</span>
                
 
                            +            <span class="n">extra_kwargs</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span>
                
 
                            +                <span class="n">force</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="c1"># If the logging does not support force=True, we should manually delete handlers</span>
                
 
                            +            <span class="k">del</span> <span class="n">manager</span><span class="o">.</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="p">[:]</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">cur_version</span> <span class="o">&gt;=</span> <span class="n">python_39</span><span class="p">:</span>
                
 
                            +            <span class="n">extra_kwargs</span><span class="p">[</span><span class="s2">&quot;encoding&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;utf-8&quot;</span>
                
 
                            +
                
 
                            +        <span class="n">logging</span><span class="o">.</span><span class="n">basicConfig</span><span class="p">(</span>
                
 
                            +            <span class="n">filename</span><span class="o">=</span><span class="n">filename</span><span class="p">,</span>
                
 
                            +            <span class="n">filemode</span><span class="o">=</span><span class="n">filemode</span><span class="p">,</span>
                
 
                            +            <span class="nb">format</span><span class="o">=</span><span class="s2">&quot;</span><span class="si">%(asctime)s</span><span class="s2"> </span><span class="si">%(levelname)s</span><span class="s2"> - </span><span class="si">%(name)s</span><span class="s2"> - </span><span class="si">%(message)s</span><span class="s2">&quot;</span><span class="p">,</span>
                
 
                            +            <span class="n">datefmt</span><span class="o">=</span><span class="s2">&quot;[%Y-%m-</span><span class="si">%d</span><span class="s2"> %H:%M:%S]&quot;</span><span class="p">,</span>
                
 
                            +            <span class="n">level</span><span class="o">=</span><span class="n">file_logging_level</span><span class="p">,</span>
                
 
                            +            <span class="o">**</span><span class="n">extra_kwargs</span><span class="p">,</span>
                
 
                                     <span class="p">)</span>
                
 
                            -        <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">conf_file_path</span><span class="p">,</span> <span class="s2">&quot;r&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">logging_configuration_file</span><span class="p">:</span>
                
 
                            -            <span class="n">config_dict</span> <span class="o">=</span> <span class="n">json</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">logging_configuration_file</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="c1"># CREATING THE PATH TO THE &quot;HOME&quot; FOLDER WITH THE LOG FILE NAME</span>
                
 
                            -        <span class="k">if</span> <span class="ow">not</span> <span class="n">logs_dir_path</span><span class="p">:</span>
                
 
                            -            <span class="n">log_file_name</span> <span class="o">=</span> <span class="n">module_name</span> <span class="o">+</span> <span class="s2">&quot;.log&quot;</span>
                
 
                            -            <span class="n">user_dir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">expanduser</span><span class="p">(</span><span class="sa">r</span><span class="s2">&quot;~&quot;</span><span class="p">)</span>
                
 
                            -            <span class="n">logs_dir_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">user_dir</span><span class="p">,</span> <span class="s2">&quot;sg_logs&quot;</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">logs_dir_path</span><span class="p">):</span>
                
 
                            -            <span class="k">try</span><span class="p">:</span>
                
 
                            -                <span class="n">os</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">logs_dir_path</span><span class="p">)</span>
                
 
                            -            <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">ex</span><span class="p">:</span>
                
 
                            -                <span class="nb">print</span><span class="p">(</span>
                
 
                            -                    <span class="s2">&quot;[WARNING] - sg_logs folder was not found and couldn&#39;t be created from the code - &quot;</span>
                
 
                            -                    <span class="s2">&quot;All of the Log output will be sent to Console!&quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">ex</span><span class="p">)</span>
                
 
                            -                <span class="p">)</span>
                
 
                            -
                
 
                            -            <span class="c1"># HANDLERS LIST IS EMPTY AS CONSOLE IS ONLY ROOT HANDLER BECAUSE MODULE LOGGERS PROPAGATE THEIR LOGS UP.</span>
                
 
                            -            <span class="n">handlers_list</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                            -            <span class="n">logger</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;level&quot;</span><span class="p">:</span> <span class="n">log_level</span><span class="p">,</span> <span class="s2">&quot;handlers&quot;</span><span class="p">:</span> <span class="n">handlers_list</span><span class="p">,</span> <span class="s2">&quot;propagate&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">}</span>
                
 
                            -            <span class="n">config_dict</span><span class="p">[</span><span class="s2">&quot;loggers&quot;</span><span class="p">][</span><span class="n">module_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">logger</span>
                
 
                            -
                
 
                            -            <span class="k">return</span> <span class="n">config_dict</span>
                
 
                            -
                
 
                            -        <span class="n">log_file_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">logs_dir_path</span><span class="p">,</span> <span class="n">log_file_name</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="c1"># THE ENTRIES TO ADD TO THE ORIGINAL CONFIGURATION</span>
                
 
                            -        <span class="n">handler_name</span> <span class="o">=</span> <span class="n">module_name</span> <span class="o">+</span> <span class="s2">&quot;_file_handler&quot;</span>
                
 
                            -        <span class="n">file_handler</span> <span class="o">=</span> <span class="p">{</span>
                
 
                            -            <span class="s2">&quot;class&quot;</span><span class="p">:</span> <span class="s2">&quot;logging.handlers.RotatingFileHandler&quot;</span><span class="p">,</span>
                
 
                            -            <span class="s2">&quot;level&quot;</span><span class="p">:</span> <span class="n">log_level</span><span class="p">,</span>
                
 
                            -            <span class="s2">&quot;formatter&quot;</span><span class="p">:</span> <span class="s2">&quot;fileFormatter&quot;</span><span class="p">,</span>
                
 
                            -            <span class="s2">&quot;filename&quot;</span><span class="p">:</span> <span class="n">log_file_path</span><span class="p">,</span>
                
 
                            -            <span class="s2">&quot;maxBytes&quot;</span><span class="p">:</span> <span class="n">max_bytes</span><span class="p">,</span>
                
 
                            -            <span class="s2">&quot;backupCount&quot;</span><span class="p">:</span> <span class="mi">20</span><span class="p">,</span>
                
 
                            -            <span class="s2">&quot;encoding&quot;</span><span class="p">:</span> <span class="s2">&quot;utf8&quot;</span><span class="p">,</span>
                
 
                            -        <span class="p">}</span>
                
 
                            -
                
 
                            -        <span class="c1"># CREATING ONLY A FILE HANDLER, CONSOLE IS ONLY ROOT HANDLER AS MODULE LOGGERS PROPAGATE THEIR LOGS UP.</span>
                
 
                            -        <span class="k">if</span> <span class="n">handlers_list</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">handlers_list</span><span class="o">.</span><span class="n">empty</span><span class="p">():</span>
                
 
                            -            <span class="n">handlers_list</span> <span class="o">=</span> <span class="p">[</span><span class="n">handler_name</span><span class="p">]</span>
                
 
                            -
                
 
                            -        <span class="n">logger</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;level&quot;</span><span class="p">:</span> <span class="n">log_level</span><span class="p">,</span> <span class="s2">&quot;handlers&quot;</span><span class="p">:</span> <span class="n">handlers_list</span><span class="p">,</span> <span class="s2">&quot;propagate&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">}</span>
                
 
                            -
                
 
                            -        <span class="c1"># ADDING THE NEW LOGGER ENTRIES TO THE CONFIG DICT</span>
                
 
                            -        <span class="n">config_dict</span><span class="p">[</span><span class="s2">&quot;handlers&quot;</span><span class="p">][</span><span class="n">handler_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">file_handler</span>
                
 
                            -        <span class="n">config_dict</span><span class="p">[</span><span class="s2">&quot;loggers&quot;</span><span class="p">][</span><span class="n">module_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">logger</span>
                
 
                            -        <span class="n">config_dict</span><span class="p">[</span><span class="s2">&quot;root&quot;</span><span class="p">][</span><span class="s2">&quot;handlers&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">handler_name</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="k">if</span> <span class="n">training_log_path</span><span class="p">:</span>
                
 
                            -            <span class="n">training_file_handler</span> <span class="o">=</span> <span class="p">{</span>
                
 
                            -                <span class="s2">&quot;class&quot;</span><span class="p">:</span> <span class="s2">&quot;logging.handlers.RotatingFileHandler&quot;</span><span class="p">,</span>
                
 
                            -                <span class="s2">&quot;level&quot;</span><span class="p">:</span> <span class="n">log_level</span><span class="p">,</span>
                
 
                            -                <span class="s2">&quot;formatter&quot;</span><span class="p">:</span> <span class="s2">&quot;fileFormatter&quot;</span><span class="p">,</span>
                
 
                            -                <span class="s2">&quot;filename&quot;</span><span class="p">:</span> <span class="n">training_log_path</span><span class="p">,</span>
                
 
                            -                <span class="s2">&quot;maxBytes&quot;</span><span class="p">:</span> <span class="n">max_bytes</span><span class="p">,</span>
                
 
                            -                <span class="s2">&quot;backupCount&quot;</span><span class="p">:</span> <span class="mi">20</span><span class="p">,</span>
                
 
                            -                <span class="s2">&quot;encoding&quot;</span><span class="p">:</span> <span class="s2">&quot;utf8&quot;</span><span class="p">,</span>
                
 
                            -            <span class="p">}</span>
                
 
                            -
                
 
                            -            <span class="c1"># ALL OF DECI_TRAINER MODULES LOGGERS PROPAGATE UP TO THE ROOT SO THE ADD TRAIN FILE HANDLER FOR THE ROOT.</span>
                
 
                            -            <span class="n">config_dict</span><span class="p">[</span><span class="s2">&quot;handlers&quot;</span><span class="p">][</span><span class="s2">&quot;training&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">training_file_handler</span>
                
 
                            -            <span class="n">config_dict</span><span class="p">[</span><span class="s2">&quot;root&quot;</span><span class="p">][</span><span class="s2">&quot;handlers&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s2">&quot;training&quot;</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="k">return</span> <span class="n">config_dict</span></div></div>
                
 
                            +        <span class="c1"># Add console handler</span>
                
 
                            +        <span class="n">console_handler</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">StreamHandler</span><span class="p">()</span>
                
 
                            +        <span class="n">console_handler</span><span class="o">.</span><span class="n">setLevel</span><span class="p">(</span><span class="n">console_logging_level</span><span class="p">)</span>
                
 
                            +        <span class="n">console_handler</span><span class="o">.</span><span class="n">setFormatter</span><span class="p">(</span>
                
 
                            +            <span class="n">logging</span><span class="o">.</span><span class="n">Formatter</span><span class="p">(</span>
                
 
                            +                <span class="s2">&quot;</span><span class="si">%(asctime)s</span><span class="s2"> </span><span class="si">%(levelname)s</span><span class="s2"> - </span><span class="si">%(filename)s</span><span class="s2"> - </span><span class="si">%(message)s</span><span class="s2">&quot;</span><span class="p">,</span>
                
 
                            +                <span class="n">datefmt</span><span class="o">=</span><span class="s2">&quot;[%Y-%m-</span><span class="si">%d</span><span class="s2"> %H:%M:%S]&quot;</span><span class="p">,</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +        <span class="n">manager</span><span class="o">.</span><span class="n">root</span><span class="o">.</span><span class="n">handlers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">console_handler</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="n">filename</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="AutoLoggerConfig.get_instance"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AutoLoggerConfig.get_instance">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">get_instance</span><span class="p">(</span><span class="bp">cls</span><span class="p">):</span>
                
 
                            +        <span class="k">global</span> <span class="n">_super_gradients_logger_config</span>
                
 
                            +        <span class="k">if</span> <span class="n">_super_gradients_logger_config</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="n">_super_gradients_logger_config</span> <span class="o">=</span> <span class="bp">cls</span><span class="p">()</span>
                
 
                            +            <span class="n">_super_gradients_logger_config</span><span class="o">.</span><span class="n">_setup_default_logging</span><span class="p">()</span>
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="n">_super_gradients_logger_config</span></div>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="AutoLoggerConfig.get_log_file_path"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AutoLoggerConfig.get_log_file_path">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">get_log_file_path</span><span class="p">(</span><span class="bp">cls</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Return the current log file used to store log messages</span>
                
 
                            +<span class="sd">        :return: Full path to log file</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="bp">self</span> <span class="o">=</span> <span class="bp">cls</span><span class="o">.</span><span class="n">get_instance</span><span class="p">()</span>
                
 
                            +        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">filename</span></div>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="AutoLoggerConfig.setup_logging"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AutoLoggerConfig.setup_logging">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">setup_logging</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">copy_already_logged_messages</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">filemode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span cla
                
 
                            +        <span class="bp">self</span> <span class="o">=</span> <span class="bp">cls</span><span class="o">.</span><span class="n">get_instance</span><span class="p">()</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_setup_logging</span><span class="p">(</span><span class="n">filename</span><span class="p">,</span> <span class="n">copy_already_logged_messages</span><span class="p">,</span> <span class="n">filemode</span><span class="p">,</span> <span class="n">log_level</span><span class="p">)</span></div></div>
                
 
                            +
                
 
                            +
                
 
                            +<span class="n">_super_gradients_logger_config</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -225,4 +241,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.common.auto_logging.console_logging &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.common.auto_logging.console_logging</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.common.auto_logging.console_logging</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">sys</span>
<span class="kn">from</span> <span class="nn">datetime</span> <span class="kn">import</span> <span class="n">datetime</span>
<span class="kn">from</span> <span class="nn">pathlib</span> <span class="kn">import</span> <span class="n">Path</span>
<span class="kn">from</span> <span class="nn">io</span> <span class="kn">import</span> <span class="n">StringIO</span>
<span class="kn">import</span> <span class="nn">atexit</span>
<span class="kn">from</span> <span class="nn">threading</span> <span class="kn">import</span> <span class="n">Lock</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.environment.env_helpers</span> <span class="kn">import</span> <span class="n">multi_process_safe</span><span class="p">,</span> <span class="n">is_main_process</span>


<span class="k">class</span> <span class="nc">BufferWriter</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;File writer buffer that opens a file only when flushing and under the condition that threshold buffersize was reached.&quot;&quot;&quot;</span>

    <span class="n">FILE_BUFFER_SIZE</span> <span class="o">=</span> <span class="mi">10_000</span>  <span class="c1"># Number of chars to be buffered before writing the buffer on disk.</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">buffer</span><span class="p">:</span> <span class="n">StringIO</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">lock</span><span class="p">:</span> <span class="n">Lock</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param filename:         Name of the file where to write the bugger</span>
<span class="sd">        :param buffer:           Buffer object</span>
<span class="sd">        :param buffer_size:      Number of chars to be buffered before writing the buffer on disk.</span>
<span class="sd">        :param lock:             Thread lock to prevent multiple threads to write at the same time</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">buffer</span> <span class="o">=</span> <span class="n">buffer</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="n">filename</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">buffer_size</span> <span class="o">=</span> <span class="n">buffer_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lock</span> <span class="o">=</span> <span class="n">lock</span>

    <span class="k">def</span> <span class="nf">write</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Write to buffer (not on disk).&quot;&quot;&quot;</span>
        <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">lock</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">buffer</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_require_flush</span><span class="p">():</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">flush</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">force</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Write the buffer on disk if relevant.&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="n">force</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">_require_flush</span><span class="p">():</span>
            <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">lock</span><span class="p">:</span>
                <span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">),</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="s2">&quot;a&quot;</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="s2">&quot;utf-8&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
                    <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">buffer</span><span class="o">.</span><span class="n">getvalue</span><span class="p">())</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">buffer</span><span class="o">.</span><span class="n">truncate</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">buffer</span><span class="o">.</span><span class="n">seek</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_require_flush</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;Indicate if a buffer is needed (i.e. if buffer size above threshold)&quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">buffer</span><span class="o">.</span><span class="n">getvalue</span><span class="p">())</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">buffer_size</span>


<span class="k">class</span> <span class="nc">StderrTee</span><span class="p">(</span><span class="n">BufferWriter</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Duplicate the stderr stream to save it into a given file.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">buffer</span><span class="p">:</span> <span class="n">StringIO</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">lock</span><span class="p">:</span> <span class="n">Lock</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param filename:         Name of the file where to write the bugger</span>
<span class="sd">        :param buffer:           Buffer object</span>
<span class="sd">        :param buffer_size:      Number of chars to be buffered before writing the buffer on disk.</span>
<span class="sd">        :param lock:             Thread lock to prevent multiple threads to write at the same time</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">filename</span><span class="p">,</span> <span class="n">buffer</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">,</span> <span class="n">lock</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stderr</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span> <span class="o">=</span> <span class="bp">self</span>

    <span class="k">def</span> <span class="fm">__del__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">stderr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">stderr</span>

    <span class="k">def</span> <span class="nf">write</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__getattr__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">attr</span><span class="p">):</span>
        <span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">stderr</span><span class="p">,</span> <span class="n">attr</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">StdoutTee</span><span class="p">(</span><span class="n">BufferWriter</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Duplicate the stdout stream to save it into a given file.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">buffer</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">lock</span><span class="p">:</span> <span class="n">Lock</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param filename:         Name of the file where to write the bugger</span>
<span class="sd">        :param buffer:           Buffer object</span>
<span class="sd">        :param buffer_size:      Number of chars to be buffered before writing the buffer on disk.</span>
<span class="sd">        :param lock:             Thread lock to prevent multiple threads to write at the same time</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">filename</span><span class="p">,</span> <span class="n">buffer</span><span class="p">,</span> <span class="n">buffer_size</span><span class="p">,</span> <span class="n">lock</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">stdout</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">stdout</span> <span class="o">=</span> <span class="bp">self</span>

    <span class="k">def</span> <span class="fm">__del__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">stdout</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span>

    <span class="k">def</span> <span class="nf">write</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__getattr__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">attr</span><span class="p">):</span>
        <span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">stdout</span><span class="p">,</span> <span class="n">attr</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">copy_file</span><span class="p">(</span><span class="n">src_filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">dest_filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">copy_mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;w&quot;</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Copy a file from source to destination. Also works when the destination folder does not exist.&quot;&quot;&quot;</span>
    <span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="n">dest_filename</span><span class="p">),</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">src_filename</span><span class="p">):</span>
        <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">src_filename</span><span class="p">,</span> <span class="s2">&quot;r&quot;</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="s2">&quot;utf-8&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">src</span><span class="p">:</span>
            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">dest_filename</span><span class="p">,</span> <span class="n">copy_mode</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="s2">&quot;utf-8&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">dst</span><span class="p">:</span>
                <span class="n">dst</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">src</span><span class="o">.</span><span class="n">read</span><span class="p">())</span>


<div class="viewcode-block" id="ConsoleSink"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.auto_logging.ConsoleSink">[docs]</a><span class="k">class</span> <span class="nc">ConsoleSink</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Singleton responsible to sink the console streams (stdout/stderr) into a file.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_setup</span><span class="p">()</span>
        <span class="n">atexit</span><span class="o">.</span><span class="n">register</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_flush</span><span class="p">)</span>  <span class="c1"># Flush at the end of the process</span>

    <span class="nd">@multi_process_safe</span>
    <span class="k">def</span> <span class="nf">_setup</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;On instantiation, setup the default sink file.&quot;&quot;&quot;</span>
        <span class="n">filename</span> <span class="o">=</span> <span class="n">Path</span><span class="o">.</span><span class="n">home</span><span class="p">()</span> <span class="o">/</span> <span class="s2">&quot;sg_logs&quot;</span> <span class="o">/</span> <span class="s2">&quot;console.log&quot;</span>
        <span class="n">filename</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">filename</span><span class="p">)</span>
        <span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">dirname</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">),</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

        <span class="n">buffer</span> <span class="o">=</span> <span class="n">StringIO</span><span class="p">()</span>
        <span class="n">lock</span> <span class="o">=</span> <span class="n">Lock</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span> <span class="o">=</span> <span class="n">StdoutTee</span><span class="p">(</span><span class="n">filename</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="n">buffer</span><span class="o">=</span><span class="n">buffer</span><span class="p">,</span> <span class="n">buffer_size</span><span class="o">=</span><span class="n">BufferWriter</span><span class="o">.</span><span class="n">FILE_BUFFER_SIZE</span><span class="p">,</span> <span class="n">lock</span><span class="o">=</span><span class="n">lock</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stderr</span> <span class="o">=</span> <span class="n">StderrTee</span><span class="p">(</span><span class="n">filename</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="n">buffer</span><span class="o">=</span><span class="n">buffer</span><span class="p">,</span> <span class="n">buffer_size</span><span class="o">=</span><span class="n">BufferWriter</span><span class="o">.</span><span class="n">FILE_BUFFER_SIZE</span><span class="p">,</span> <span class="n">lock</span><span class="o">=</span><span class="n">lock</span><span class="p">)</span>

        <span class="c1"># We don&#39;t want to rewrite this for subprocesses when using DDP.</span>
        <span class="k">if</span> <span class="n">is_main_process</span><span class="p">():</span>
            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="s2">&quot;utf-8&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
                <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s2">&quot;============================================================</span><span class="se">\n</span><span class="s2">&quot;</span><span class="p">)</span>
                <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;New run started at </span><span class="si">{</span><span class="n">datetime</span><span class="o">.</span><span class="n">now</span><span class="p">()</span><span class="o">.</span><span class="n">strftime</span><span class="p">(</span><span class="s2">&quot;%Y-%m-</span><span class="si">%d</span><span class="s2">.%H:%M:%S.</span><span class="si">%f</span><span class="s2">&quot;</span><span class="p">)</span><span class="si">}</span><span class="se">\n</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;sys.argv: &quot;</span><span class="si">{</span><span class="s2">&quot; &quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">sys</span><span class="o">.</span><span class="n">argv</span><span class="p">)</span><span class="si">}</span><span class="s1">&quot;</span><span class="se">\n</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s2">&quot;============================================================</span><span class="se">\n</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The console stream is logged into </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">filename</span><span class="si">}</span><span class="se">\n</span><span class="s2">&quot;</span><span class="p">)</span>

    <span class="nd">@multi_process_safe</span>
    <span class="k">def</span> <span class="nf">_set_location</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Copy and redirect the sink file into another location.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_flush</span><span class="p">()</span>

        <span class="n">prev_filename</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">filename</span>
        <span class="n">copy_file</span><span class="p">(</span><span class="n">src_filename</span><span class="o">=</span><span class="n">prev_filename</span><span class="p">,</span> <span class="n">dest_filename</span><span class="o">=</span><span class="n">filename</span><span class="p">,</span> <span class="n">copy_mode</span><span class="o">=</span><span class="s2">&quot;a&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="n">filename</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="n">filename</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">filename</span> <span class="o">=</span> <span class="n">filename</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The console stream is now moved to </span><span class="si">{</span><span class="n">filename</span><span class="si">}</span><span class="se">\n</span><span class="s2">&quot;</span><span class="p">)</span>

<div class="viewcode-block" id="ConsoleSink.set_location"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.auto_logging.ConsoleSink.set_location">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">set_location</span><span class="p">(</span><span class="n">filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;Copy and redirect the sink file into another location.&quot;&quot;&quot;</span>
        <span class="n">_console_sink</span><span class="o">.</span><span class="n">_set_location</span><span class="p">(</span><span class="n">filename</span><span class="p">)</span></div>

    <span class="nd">@multi_process_safe</span>
    <span class="k">def</span> <span class="nf">_flush</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Force the flush on stdout and stderr.&quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stdout</span><span class="o">.</span><span class="n">flush</span><span class="p">(</span><span class="n">force</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stderr</span><span class="o">.</span><span class="n">flush</span><span class="p">(</span><span class="n">force</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

<div class="viewcode-block" id="ConsoleSink.flush"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.auto_logging.ConsoleSink.flush">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">flush</span><span class="p">():</span>
        <span class="sd">&quot;&quot;&quot;Force the flush on stdout and stderr.&quot;&quot;&quot;</span>
        <span class="n">_console_sink</span><span class="o">.</span><span class="n">_flush</span><span class="p">()</span></div>

<div class="viewcode-block" id="ConsoleSink.get_filename"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.auto_logging.ConsoleSink.get_filename">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">get_filename</span><span class="p">():</span>
        <span class="sd">&quot;&quot;&quot;Get the filename of the sink.&quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">_console_sink</span><span class="o">.</span><span class="n">filename</span></div></div>


<span class="n">_console_sink</span> <span class="o">=</span> <span class="n">ConsoleSink</span><span class="p">()</span>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.aws_connection.aws_connector &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.aws_connection.aws_connector &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -93,7 +95,7 @@
 
                             <span class="kn">from</span> <span class="nn">botocore.exceptions</span> <span class="kn">import</span> <span class="n">ClientError</span><span class="p">,</span> <span class="n">ProfileNotFound</span>
                
 
                            -<div class="viewcode-block" id="AWSConnector"><a class="viewcode-back" href="../../../../super_gradients.common.aws_connection.html#super_gradients.common.AWSConnector">[docs]</a><span class="k">class</span> <span class="nc">AWSConnector</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="AWSConnector"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AWSConnector">[docs]</a><span class="k">class</span> <span class="nc">AWSConnector</span><span class="p">:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    AWSConnector - Connects to AWS using Credentials File or IAM Role</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
@@ -133,7 +135,7 @@
 
                                                 <span class="n">ex</span><span class="p">))</span>
                
 
                                         <span class="k">return</span> <span class="kc">None</span>
                
 
                            -<div class="viewcode-block" id="AWSConnector.get_aws_session"><a class="viewcode-back" href="../../../../super_gradients.common.aws_connection.html#super_gradients.common.AWSConnector.get_aws_session">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="AWSConnector.get_aws_session"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AWSConnector.get_aws_session">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">get_aws_session</span><span class="p">(</span><span class="n">profile_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">boto3</span><span class="o">.</span><span class="n">Session</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        get_aws_session - Connects to AWS to retrieve an AWS Session</span>
                
@@ -149,7 +151,7 @@
 
                                     <span class="k">return</span> <span class="n">aws_session</span></div>
                
 
                            -<div class="viewcode-block" id="AWSConnector.get_aws_client_for_service_name"><a class="viewcode-back" href="../../../../super_gradients.common.aws_connection.html#super_gradients.common.AWSConnector.get_aws_client_for_service_name">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="AWSConnector.get_aws_client_for_service_name"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AWSConnector.get_aws_client_for_service_name">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">get_aws_client_for_service_name</span><span class="p">(</span><span class="n">profile_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">service_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">boto3</span><span class="o">.</span><span class="n">Session</span><span class="o">.</span><span class="n">client</span><
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        get_aws_client_for_service_name - Connects to AWS to retrieve the relevant Client</span>
                
@@ -166,7 +168,7 @@
 
                                     <span class="k">return</span> <span class="n">aws_session</span><span class="o">.</span><span class="n">client</span><span class="p">(</span><span class="n">service_name</span><span class="o">=</span><span class="n">service_name</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="AWSConnector.get_aws_resource_for_service_name"><a class="viewcode-back" href="../../../../super_gradients.common.aws_connection.html#super_gradients.common.AWSConnector.get_aws_resource_for_service_name">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="AWSConnector.get_aws_resource_for_service_name"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AWSConnector.get_aws_resource_for_service_name">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">get_aws_resource_for_service_name</span><span class="p">(</span><span class="n">profile_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">service_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">boto3</span><span class="o">.</span><span class="n">Session</span><span class="o">.</span><span class="n">resource</sp
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Connects to AWS to retrieve the relevant Resource (More functionality then Client)</span>
                
@@ -183,7 +185,7 @@
 
                                     <span class="k">return</span> <span class="n">aws_session</span><span class="o">.</span><span class="n">resource</span><span class="p">(</span><span class="n">service_name</span><span class="o">=</span><span class="n">service_name</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="AWSConnector.is_client_error"><a class="viewcode-back" href="../../../../super_gradients.common.aws_connection.html#super_gradients.common.AWSConnector.is_client_error">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="AWSConnector.is_client_error"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.AWSConnector.is_client_error">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">is_client_error</span><span class="p">(</span><span class="n">code</span><span class="p">):</span>
                
 
                                     <span class="n">e</span> <span class="o">=</span> <span class="n">sys</span><span class="o">.</span><span class="n">exc_info</span><span class="p">()[</span><span class="mi">1</span><span class="p">]</span>
                
 
                                     <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">e</span><span class="p">,</span> <span class="n">ClientError</span><span class="p">)</span> <span class="ow">and</span> <span class="n">e</span><span class="o">.</span><span class="n">response</span><span class="p">[</span><span class="s2">&quot;Error&quot;</span><span class="p">][</span><span class="s2">&quot;Code&quot;</span><span class="p">]</span> <span class="o">==</span> <span class
                
@@ -218,4 +220,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.common.aws_connection.aws_secrets_manager_connector &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.common.aws_connection.aws_secrets_manager_connector</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.common.aws_connection.aws_secrets_manager_connector</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">json</span>
<span class="kn">import</span> <span class="nn">logging</span>

<span class="kn">from</span> <span class="nn">super_gradients.common</span> <span class="kn">import</span> <span class="n">AWSConnector</span>
<span class="kn">from</span> <span class="nn">super_gradients.common</span> <span class="kn">import</span> <span class="n">explicit_params_validation</span>


<div class="viewcode-block" id="AWSSecretsManagerConnector"><a class="viewcode-back" href="../../../../super_gradients.common.aws_connection.html#super_gradients.common.aws_connection.aws_secrets_manager_connector.AWSSecretsManagerConnector">[docs]</a><span class="k">class</span> <span class="nc">AWSSecretsManagerConnector</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    AWSSecretsManagerConnector - This class handles the AWS Secrets Manager connection</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="vm">__slots__</span> <span class="o">=</span> <span class="p">[]</span>  <span class="c1"># Making the class immutable for runtime safety</span>
    <span class="n">current_environment_client</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">DECI_ENVIRONMENTS</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;research&#39;</span><span class="p">,</span> <span class="s1">&#39;development&#39;</span><span class="p">,</span> <span class="s1">&#39;staging&#39;</span><span class="p">,</span> <span class="s1">&#39;production&#39;</span><span class="p">]</span>

    <span class="nd">@staticmethod</span>
    <span class="nd">@explicit_params_validation</span><span class="p">(</span><span class="n">validation_type</span><span class="o">=</span><span class="s1">&#39;NoneOrEmpty&#39;</span><span class="p">)</span>
    <span class="k">def</span> <span class="nf">get_secret_value_for_secret_key</span><span class="p">(</span><span class="n">aws_env</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">secret_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">secret_key</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        get_secret_value_for_secret_key - Gets a Secret Value from AWS Secrets Manager for the Provided Key</span>
<span class="sd">            :param aws_env:         The environment to get the secret for</span>
<span class="sd">            :param secret_name: The Secret Name stored in Secrets Manager</span>
<span class="sd">            :param secret_key:  The Secret Key To retrieve it&#39;s value from AWS</span>
<span class="sd">            :return:  str: The Secret Value</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">current_class_name</span> <span class="o">=</span> <span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span>
        <span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="n">current_class_name</span><span class="p">)</span>
        <span class="n">secret_key</span> <span class="o">=</span> <span class="n">secret_key</span><span class="o">.</span><span class="n">upper</span><span class="p">()</span>
        <span class="n">aws_secrets_dict</span> <span class="o">=</span> <span class="n">AWSSecretsManagerConnector</span><span class="o">.</span><span class="n">__get_secrets_manager_dict_for_secret_name</span><span class="p">(</span>
            <span class="n">aws_env</span><span class="o">=</span><span class="n">aws_env</span><span class="p">,</span> <span class="n">secret_name</span><span class="o">=</span><span class="n">secret_name</span><span class="p">)</span>

        <span class="n">secret_key</span> <span class="o">=</span> <span class="s1">&#39;.&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">aws_env</span><span class="o">.</span><span class="n">upper</span><span class="p">(),</span> <span class="n">secret_key</span><span class="p">])</span>
        <span class="k">if</span> <span class="n">secret_key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">aws_secrets_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
            <span class="n">error</span> <span class="o">=</span> <span class="sa">f</span><span class="s1">&#39;[</span><span class="si">{</span><span class="n">current_class_name</span><span class="si">}</span><span class="s1">] - Secret Key (</span><span class="si">{</span><span class="n">secret_key</span><span class="si">}</span><span class="s1">) not Found in AWS Secret: &#39;</span> <span class="o">+</span> <span class="n">secret_name</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="n">error</span><span class="p">)</span>
            <span class="k">raise</span> <span class="ne">EnvironmentError</span><span class="p">(</span><span class="n">error</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">aws_secrets_dict</span><span class="p">[</span><span class="n">secret_key</span><span class="p">]</span>

    <span class="nd">@staticmethod</span>
    <span class="nd">@explicit_params_validation</span><span class="p">(</span><span class="n">validation_type</span><span class="o">=</span><span class="s1">&#39;NoneOrEmpty&#39;</span><span class="p">)</span>
    <span class="k">def</span> <span class="nf">get_secret_values_dict_for_secret_key_properties</span><span class="p">(</span><span class="n">env</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">secret_key</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">secret_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                                                         <span class="n">db_properties_set</span><span class="p">:</span> <span class="nb">set</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        get_config_dict - Returns the config dict of the properties from the properties dict</span>
<span class="sd">            :param  env:                The environment to open the dict for</span>
<span class="sd">            :param  secret_key:         The Secret Key</span>
<span class="sd">            :param  secret_name:        The Secret to Retrieve to from AWS secrets manager (usually project name)</span>
<span class="sd">            :param  db_properties_set:  The set of the properties to get secrets values for</span>
<span class="sd">            :return:  dict The secrets dict for the requested property</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">current_class_name</span> <span class="o">=</span> <span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span>
        <span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="n">current_class_name</span><span class="p">)</span>
        <span class="n">aws_secrets_dict</span> <span class="o">=</span> <span class="n">AWSSecretsManagerConnector</span><span class="o">.</span><span class="n">__get_secrets_manager_dict_for_secret_name</span><span class="p">(</span>
            <span class="n">aws_env</span><span class="o">=</span><span class="n">env</span><span class="p">,</span> <span class="n">secret_name</span><span class="o">=</span><span class="n">secret_name</span><span class="p">)</span>

        <span class="n">aws_env_safe_secrets</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="c1"># FILL THE DICT VALUES FROM THE AWS SECRETS RESPONSE</span>
        <span class="k">if</span> <span class="n">db_properties_set</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">secret_key_property</span> <span class="ow">in</span> <span class="n">db_properties_set</span><span class="p">:</span>
                <span class="n">secret_key_to_retrieve</span> <span class="o">=</span> <span class="s1">&#39;.&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">env</span><span class="o">.</span><span class="n">upper</span><span class="p">(),</span> <span class="n">secret_key</span><span class="p">,</span> <span class="n">secret_key_property</span><span class="p">])</span>
                <span class="k">if</span> <span class="n">secret_key_to_retrieve</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">aws_secrets_dict</span><span class="p">:</span>
                    <span class="n">error</span> <span class="o">=</span> <span class="sa">f</span><span class="s1">&#39;[</span><span class="si">{</span><span class="n">current_class_name</span><span class="si">}</span><span class="s1">] - Error retrieving data from AWS Secrets Manager for Secret Key &quot;</span><span class="si">{</span><span class="n">secret_name</span><span class="si">}</span><span class="s1">&quot;: &#39;</span> \
                            <span class="sa">f</span><span class="s1">&#39;The secret property &quot;</span><span class="si">{</span><span class="n">secret_key_property</span><span class="si">}</span><span class="s1">&quot; Does Not Exist&#39;</span>
                    <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="n">error</span><span class="p">)</span>
                    <span class="k">raise</span> <span class="ne">EnvironmentError</span><span class="p">(</span><span class="n">error</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">env_stripped_key_name</span> <span class="o">=</span> <span class="n">secret_key_to_retrieve</span><span class="o">.</span><span class="n">lstrip</span><span class="p">(</span><span class="n">env</span><span class="o">.</span><span class="n">upper</span><span class="p">())</span><span class="o">.</span><span class="n">lstrip</span><span class="p">(</span><span class="s1">&#39;.&#39;</span><span class="p">)</span>
                    <span class="n">aws_env_safe_secrets</span><span class="p">[</span><span class="n">env_stripped_key_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">aws_secrets_dict</span><span class="p">[</span><span class="n">secret_key_to_retrieve</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># &quot;db_properties_set&quot; is not specified - validating and returning all the secret keys and values for</span>
            <span class="c1"># the secret name.</span>
            <span class="k">for</span> <span class="n">secret_key_name</span><span class="p">,</span> <span class="n">secret_value</span> <span class="ow">in</span> <span class="n">aws_secrets_dict</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                <span class="n">secret_key_to_retrieve</span> <span class="o">=</span> <span class="s1">&#39;.&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">env</span><span class="o">.</span><span class="n">upper</span><span class="p">(),</span> <span class="n">secret_key</span><span class="p">])</span>
                <span class="k">assert</span> <span class="n">secret_key_name</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span>
                    <span class="n">env</span><span class="o">.</span><span class="n">upper</span><span class="p">()),</span> <span class="sa">f</span><span class="s1">&#39;The secret key property &quot;</span><span class="si">{</span><span class="n">secret_key_name</span><span class="si">}</span><span class="s1">&quot;, found in secret named </span><span class="si">{</span><span class="n">secret_name</span><span class="si">}</span><span class="s1">, is not following the convention of &#39;</span> \
                                  <span class="sa">f</span><span class="s1">&#39;environment prefix. please add the environment prefix &quot;</span><span class="si">{</span><span class="n">env</span><span class="o">.</span><span class="n">upper</span><span class="p">()</span><span class="si">}</span><span class="s1">&quot; to property &quot;</span><span class="si">{</span><span class="n">secret_key_name</span><span class="si">}</span><span class="s1">&quot;&#39;</span>
                <span class="k">if</span> <span class="n">secret_key_name</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="n">secret_key_to_retrieve</span><span class="p">):</span>
                    <span class="n">env_stripped_key_name</span> <span class="o">=</span> <span class="n">secret_key_name</span><span class="o">.</span><span class="n">lstrip</span><span class="p">(</span><span class="n">env</span><span class="o">.</span><span class="n">upper</span><span class="p">())</span><span class="o">.</span><span class="n">lstrip</span><span class="p">(</span><span class="s1">&#39;.&#39;</span><span class="p">)</span>
                    <span class="n">aws_env_safe_secrets</span><span class="p">[</span><span class="n">env_stripped_key_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">secret_value</span>
        <span class="k">return</span> <span class="n">aws_env_safe_secrets</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">__get_secrets_manager_dict_for_secret_name</span><span class="p">(</span><span class="n">aws_env</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">secret_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        __get_secrets_manager_dict_for_secret_name</span>
<span class="sd">            :param  aws_env:                The environment to open the dict for</span>
<span class="sd">            :param  secret_name:        The Secret to Retrieve to from AWS secrets manager (usually project name)</span>
<span class="sd">            :return: python Dictionary with the key/value pairs stored in AWS Secrets Manager</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">current_class_name</span> <span class="o">=</span> <span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span>
        <span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="n">current_class_name</span><span class="p">)</span>

        <span class="n">secrets_path</span> <span class="o">=</span> <span class="n">AWSSecretsManagerConnector</span><span class="o">.</span><span class="n">__get_secrets_path_from_secret_name</span><span class="p">(</span><span class="n">aws_env</span><span class="p">,</span> <span class="n">secret_name</span><span class="p">)</span>

        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">AWSSecretsManagerConnector</span><span class="o">.</span><span class="n">current_environment_client</span><span class="p">:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s1">&#39;Initializing a new secrets manager client...&#39;</span><span class="p">)</span>
                <span class="n">AWSSecretsManagerConnector</span><span class="o">.</span><span class="n">current_environment_client</span> <span class="o">=</span> <span class="n">AWSConnector</span><span class="o">.</span><span class="n">get_aws_client_for_service_name</span><span class="p">(</span>
                    <span class="n">profile_name</span><span class="o">=</span><span class="n">aws_env</span><span class="p">,</span>
                    <span class="n">service_name</span><span class="o">=</span><span class="s1">&#39;secretsmanager&#39;</span><span class="p">)</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Fetching the secret &quot;</span><span class="si">{</span><span class="n">secret_name</span><span class="si">}</span><span class="s1">&quot; in env &quot;</span><span class="si">{</span><span class="n">aws_env</span><span class="si">}</span><span class="s1">&quot;&#39;</span><span class="p">)</span>
            <span class="n">aws_secrets</span> <span class="o">=</span> <span class="n">AWSSecretsManagerConnector</span><span class="o">.</span><span class="n">current_environment_client</span><span class="o">.</span><span class="n">get_secret_value</span><span class="p">(</span><span class="n">SecretId</span><span class="o">=</span><span class="n">secrets_path</span><span class="p">)</span>
            <span class="n">aws_secrets_dict</span> <span class="o">=</span> <span class="n">json</span><span class="o">.</span><span class="n">loads</span><span class="p">(</span><span class="n">aws_secrets</span><span class="p">[</span><span class="s1">&#39;SecretString&#39;</span><span class="p">])</span>
            <span class="k">return</span> <span class="n">aws_secrets_dict</span>

        <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">ex</span><span class="p">:</span>
            <span class="n">error</span> <span class="o">=</span> <span class="s1">&#39;[&#39;</span> <span class="o">+</span> <span class="n">current_class_name</span> <span class="o">+</span> <span class="s1">&#39;] - Caught Exception while trying to connect to aws to get credentials from secrets manager: &#39;</span> <span class="o">+</span> <span class="s1">&#39;&quot;&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span>
                <span class="n">ex</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;&quot;&#39;</span> <span class="o">+</span> <span class="s1">&#39; for &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">secrets_path</span><span class="p">)</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="n">error</span><span class="p">)</span>
            <span class="k">raise</span> <span class="ne">EnvironmentError</span><span class="p">(</span><span class="n">error</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">__get_secrets_path_from_secret_name</span><span class="p">(</span><span class="n">aws_env</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">secret_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        __get_secrets_path_from_secret_name - Extracts the full secret path based on the Environment</span>
<span class="sd">        :param aws_env:         Env</span>
<span class="sd">        :param secret_name: Secret Name</span>
<span class="sd">        :return:     str:   The full secret path</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">current_class_name</span> <span class="o">=</span> <span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span>
        <span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="n">current_class_name</span><span class="p">)</span>

        <span class="c1"># Checking for lowercase exact match, in order to prevent any implicit usage of the environments.</span>
        <span class="k">if</span> <span class="n">aws_env</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">AWSSecretsManagerConnector</span><span class="o">.</span><span class="n">DECI_ENVIRONMENTS</span><span class="p">:</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">critical</span><span class="p">(</span><span class="s1">&#39;[&#39;</span> <span class="o">+</span> <span class="n">current_class_name</span> <span class="o">+</span> <span class="s1">&#39; ] -  wrong environment param... Exiting&#39;</span><span class="p">)</span>
            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s1">&#39;[&#39;</span> <span class="o">+</span> <span class="n">current_class_name</span> <span class="o">+</span> <span class="s1">&#39;] - wrong environment param&#39;</span><span class="p">)</span>

        <span class="n">secrets_path</span> <span class="o">=</span> <span class="s1">&#39;/&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">aws_env</span><span class="p">,</span> <span class="n">secret_name</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">secrets_path</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.common.crash_handler.crash_handler &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.common.crash_handler.crash_handler</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.common.crash_handler.crash_handler</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">sys</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.crash_handler.crash_tips_setup</span> <span class="kn">import</span> <span class="n">setup_crash_tips</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.crash_handler.exception_monitoring_setup</span> <span class="kn">import</span> <span class="n">setup_pro_user_monitoring</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.crash_handler.exception</span> <span class="kn">import</span> <span class="n">register_exceptions</span>


<div class="viewcode-block" id="setup_crash_handler"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.setup_crash_handler">[docs]</a><span class="k">def</span> <span class="nf">setup_crash_handler</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;Setup the environment to handle crashes, with crash tips and more.&quot;&quot;&quot;</span>
    <span class="n">is_setup_crash_tips</span> <span class="o">=</span> <span class="n">setup_crash_tips</span><span class="p">()</span>
    <span class="n">is_setup_pro_user_monitoring</span> <span class="o">=</span> <span class="n">setup_pro_user_monitoring</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">is_setup_crash_tips</span> <span class="ow">or</span> <span class="n">is_setup_pro_user_monitoring</span><span class="p">:</span>  <span class="c1"># We don&#39;t want to wrap sys.excepthook when not required</span>
        <span class="n">sys</span><span class="o">.</span><span class="n">excepthook</span> <span class="o">=</span> <span class="n">register_exceptions</span><span class="p">(</span><span class="n">sys</span><span class="o">.</span><span class="n">excepthook</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.data_connection.s3_connector &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.data_connection.s3_connector &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -98,11 +100,11 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">ILogger</span>
                
 
                            -<div class="viewcode-block" id="KeyNotExistInBucketError"><a class="viewcode-back" href="../../../../super_gradients.common.data_connection.html#super_gradients.common.KeyNotExistInBucketError">[docs]</a><span class="k">class</span> <span class="nc">KeyNotExistInBucketError</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
                
 
                            -    <span class="k">pass</span></div>
                
 
                            +<span class="k">class</span> <span class="nc">KeyNotExistInBucketError</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
                
 
                            +    <span class="k">pass</span>
                
 
                            -<div class="viewcode-block" id="S3Connector"><a class="viewcode-back" href="../../../../super_gradients.common.data_connection.html#super_gradients.common.S3Connector">[docs]</a><span class="k">class</span> <span class="nc">S3Connector</span><span class="p">(</span><span class="n">ILogger</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="S3Connector"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.S3Connector">[docs]</a><span class="k">class</span> <span class="nc">S3Connector</span><span class="p">(</span><span class="n">ILogger</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    S3Connector - S3 Connection Manager</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
@@ -458,7 +460,7 @@
 
                                                                                      <span class="n">ExpiresIn</span><span class="o">=</span><span class="n">expiration</span><span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="n">response</span>
                
 
                            -<div class="viewcode-block" id="S3Connector.convert_content_length_to_mb"><a class="viewcode-back" href="../../../../super_gradients.common.data_connection.html#super_gradients.common.S3Connector.convert_content_length_to_mb">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="S3Connector.convert_content_length_to_mb"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.S3Connector.convert_content_length_to_mb">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">convert_content_length_to_mb</span><span class="p">(</span><span class="n">content_length</span><span class="p">):</span>
                
 
                                     <span class="k">return</span> <span class="nb">round</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">content_length</span> <span class="o">/</span> <span class="p">(</span><span class="mf">1e+6</span><span class="p">)</span><span class="si">:</span><span class="s1">2f</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">),</span> <span 
                
@@ -524,4 +526,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.data_interface.adnn_model_repository_data_interface &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.data_interface.adnn_model_repository_data_interface &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -92,11 +94,11 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">ILogger</span>
                
 
                            -<div class="viewcode-block" id="ModelCheckpointNotFoundException"><a class="viewcode-back" href="../../../../super_gradients.common.data_interface.html#super_gradients.common.ModelCheckpointNotFoundException">[docs]</a><span class="k">class</span> <span class="nc">ModelCheckpointNotFoundException</span><span class="p">(</span><span class="ne">RuntimeError</span><span class="p">):</span>
                
 
                            -    <span class="k">pass</span></div>
                
 
                            +<span class="k">class</span> <span class="nc">ModelCheckpointNotFoundException</span><span class="p">(</span><span class="ne">RuntimeError</span><span class="p">):</span>
                
 
                            +    <span class="k">pass</span>
                
 
                            -<div class="viewcode-block" id="ADNNModelRepositoryDataInterfaces"><a class="viewcode-back" href="../../../../super_gradients.common.data_interface.html#super_gradients.common.ADNNModelRepositoryDataInterfaces">[docs]</a><span class="k">class</span> <span class="nc">ADNNModelRepositoryDataInterfaces</span><span class="p">(</span><span class="n">ILogger</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="ADNNModelRepositoryDataInterfaces"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.ADNNModelRepositoryDataInterfaces">[docs]</a><span class="k">class</span> <span class="nc">ADNNModelRepositoryDataInterfaces</span><span class="p">(</span><span class="n">ILogger</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    ResearchModelRepositoryDataInterface</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
@@ -293,4 +295,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.data_interface.dataset_data_interface &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.data_interface.dataset_data_interface &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -92,7 +94,7 @@
 
                             <span class="kn">import</span> <span class="nn">zipfile</span>
                
 
                            -<div class="viewcode-block" id="DatasetDataInterface"><a class="viewcode-back" href="../../../../super_gradients.common.data_interface.html#super_gradients.common.DatasetDataInterface">[docs]</a><span class="k">class</span> <span class="nc">DatasetDataInterface</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="DatasetDataInterface"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.DatasetDataInterface">[docs]</a><span class="k">class</span> <span class="nc">DatasetDataInterface</span><span class="p">:</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">env</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">data_connection_source</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;s3&#39;</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
@@ -167,4 +169,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.data_types.enum.deep_learning_task &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.data_types.enum.deep_learning_task &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -89,7 +91,7 @@
 
                             <span></span><span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
                
 
                            -<div class="viewcode-block" id="DeepLearningTask"><a class="viewcode-back" href="../../../../../super_gradients.common.data_types.enum.html#super_gradients.common.DeepLearningTask">[docs]</a><span class="k">class</span> <span class="nc">DeepLearningTask</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="DeepLearningTask"><a class="viewcode-back" href="../../../../../super_gradients.common.html#super_gradients.common.DeepLearningTask">[docs]</a><span class="k">class</span> <span class="nc">DeepLearningTask</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
                
 
                                 <span class="n">CLASSIFICATION</span> <span class="o">=</span> <span class="s1">&#39;classification&#39;</span>
                
 
                                 <span class="n">SEMANTIC_SEGMENTATION</span> <span class="o">=</span> <span class="s1">&#39;semantic_segmentation&#39;</span>
                
 
                                 <span class="n">OBJECT_DETECTION</span> <span class="o">=</span> <span class="s1">&#39;object_detection&#39;</span>
                
@@ -126,4 +128,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.data_types.enum.evaluation_type &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.data_types.enum.evaluation_type &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -89,11 +91,11 @@
 
                             <span></span><span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
                
 
                            -<div class="viewcode-block" id="EvaluationType"><a class="viewcode-back" href="../../../../../super_gradients.common.data_types.enum.html#super_gradients.common.EvaluationType">[docs]</a><span class="k">class</span> <span class="nc">EvaluationType</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="EvaluationType"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.common.EvaluationType">[docs]</a><span class="k">class</span> <span class="nc">EvaluationType</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    EvaluationType</span>
                
 
                            -<span class="sd">    Passed to SgModel.evaluate(..), and controls which phase callbacks should be triggered (if at all).</span>
                
 
                            +<span class="sd">    Passed to Trainer.evaluate(..), and controls which phase callbacks should be triggered (if at all).</span>
                
 
                             <span class="sd">        Attributes:</span>
                
 
                             <span class="sd">            TEST</span>
                
@@ -131,4 +133,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.data_types.enum.multi_gpu_mode &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.data_types.enum.multi_gpu_mode &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -87,9 +89,10 @@
 
                               <h1>Source code for super_gradients.common.data_types.enum.multi_gpu_mode</h1><div class="highlight"><pre>
                
 
                             <span></span><span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
                
 
                            +<span class="kn">import</span> <span class="nn">stringcase</span>
                
 
                            -<div class="viewcode-block" id="MultiGPUMode"><a class="viewcode-back" href="../../../../../super_gradients.training.sg_model.html#super_gradients.common.MultiGPUMode">[docs]</a><span class="k">class</span> <span class="nc">MultiGPUMode</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="MultiGPUMode"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.common.MultiGPUMode">[docs]</a><span class="k">class</span> <span class="nc">MultiGPUMode</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    MultiGPUMode</span>
                
@@ -98,10 +101,26 @@
 
                             <span class="sd">            DATA_PARALLEL             - Multiple GPUs, Synchronous</span>
                
 
                             <span class="sd">            DISTRIBUTED_DATA_PARALLEL - Multiple GPUs, Asynchronous</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="n">OFF</span> <span class="o">=</span> <span class="s1">&#39;Off&#39;</span>
                
 
                            -    <span class="n">DATA_PARALLEL</span> <span class="o">=</span> <span class="s1">&#39;DP&#39;</span>
                
 
                            -    <span class="n">DISTRIBUTED_DATA_PARALLEL</span> <span class="o">=</span> <span class="s1">&#39;DDP&#39;</span>
                
 
                            -    <span class="n">AUTO</span> <span class="o">=</span> <span class="s2">&quot;AUTO&quot;</span></div>
                
 
                            +
                
 
                            +    <span class="n">OFF</span> <span class="o">=</span> <span class="s2">&quot;Off&quot;</span>
                
 
                            +    <span class="n">DATA_PARALLEL</span> <span class="o">=</span> <span class="s2">&quot;DP&quot;</span>
                
 
                            +    <span class="n">DISTRIBUTED_DATA_PARALLEL</span> <span class="o">=</span> <span class="s2">&quot;DDP&quot;</span>
                
 
                            +    <span class="n">AUTO</span> <span class="o">=</span> <span class="s2">&quot;AUTO&quot;</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="MultiGPUMode.dict"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.common.MultiGPUMode.dict">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">dict</span><span class="p">(</span><span class="bp">cls</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        return dictionary mapping from the mode name (in call string cases) to the enum value</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">out_dict</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
                
 
                            +        <span class="k">for</span> <span class="n">mode</span> <span class="ow">in</span> <span class="n">MultiGPUMode</span><span class="p">:</span>
                
 
                            +            <span class="n">out_dict</span><span class="p">[</span><span class="n">mode</span><span class="o">.</span><span class="n">value</span><span class="p">]</span> <span class="o">=</span> <span class="n">mode</span>
                
 
                            +            <span class="n">out_dict</span><span class="p">[</span><span class="n">mode</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">mode</span>
                
 
                            +            <span class="n">out_dict</span><span class="p">[</span><span class="n">stringcase</span><span class="o">.</span><span class="n">capitalcase</span><span class="p">(</span><span class="n">mode</span><span class="o">.</span><span class="n">name</span><span class="p">)]</span> <span class="o">=</span> <span class="n">mode</span>
                
 
                            +            <span class="n">out_dict</span><span class="p">[</span><span class="n">stringcase</span><span class="o">.</span><span class="n">camelcase</span><span class="p">(</span><span class="n">mode</span><span class="o">.</span><span class="n">name</span><span class="p">)]</span> <span class="o">=</span> <span class="n">mode</span>
                
 
                            +            <span class="n">out_dict</span><span class="p">[</span><span class="n">stringcase</span><span class="o">.</span><span class="n">lowercase</span><span class="p">(</span><span class="n">mode</span><span class="o">.</span><span class="n">name</span><span class="p">)]</span> <span class="o">=</span> <span class="n">mode</span>
                
 
                            +        <span class="n">out_dict</span><span class="p">[</span><span class="kc">False</span><span class="p">]</span> <span class="o">=</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">OFF</span>
                
 
                            +        <span class="k">return</span> <span class="n">out_dict</span></div></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -131,4 +150,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.data_types.enum.strict_load &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.data_types.enum.strict_load &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -89,18 +91,19 @@
 
                             <span></span><span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
                
 
                            -<div class="viewcode-block" id="StrictLoad"><a class="viewcode-back" href="../../../../../super_gradients.training.sg_model.html#super_gradients.common.StrictLoad">[docs]</a><span class="k">class</span> <span class="nc">StrictLoad</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="StrictLoad"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.common.StrictLoad">[docs]</a><span class="k">class</span> <span class="nc">StrictLoad</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Wrapper for adding more functionality to torch&#39;s strict_load parameter in load_state_dict().</span>
                
 
                            -<span class="sd">        Attributes:</span>
                
 
                            -<span class="sd">            OFF              - Native torch &quot;strict_load = off&quot; behaviour. See nn.Module.load_state_dict() documentation for more details.</span>
                
 
                            -<span class="sd">            ON               - Native torch &quot;strict_load = on&quot; behaviour. See nn.Module.load_state_dict() documentation for more details.</span>
                
 
                            -<span class="sd">            NO_KEY_MATCHING  - Allows the usage of SuperGradient&#39;s adapt_checkpoint function, which loads a checkpoint by matching each</span>
                
 
                            -<span class="sd">                               layer&#39;s shapes (and bypasses the strict matching of the names of each layer (ie: disregards the state_dict key matching)).</span>
                
 
                            +<span class="sd">    Attributes:</span>
                
 
                            +<span class="sd">        OFF              - Native torch &quot;strict_load = off&quot; behaviour. See nn.Module.load_state_dict() documentation for more details.</span>
                
 
                            +<span class="sd">        ON               - Native torch &quot;strict_load = on&quot; behaviour. See nn.Module.load_state_dict() documentation for more details.</span>
                
 
                            +<span class="sd">        NO_KEY_MATCHING  - Allows the usage of SuperGradient&#39;s adapt_checkpoint function, which loads a checkpoint by matching each</span>
                
 
                            +<span class="sd">                           layer&#39;s shapes (and bypasses the strict matching of the names of each layer (ie: disregards the state_dict key matching)).</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                                 <span class="n">OFF</span> <span class="o">=</span> <span class="kc">False</span>
                
 
                                 <span class="n">ON</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                            -    <span class="n">NO_KEY_MATCHING</span> <span class="o">=</span> <span class="s1">&#39;no_key_matching&#39;</span></div>
                
 
                            +    <span class="n">NO_KEY_MATCHING</span> <span class="o">=</span> <span class="s2">&quot;no_key_matching&quot;</span></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -130,4 +133,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.common.data_types.enum.upsample_mode &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
        <script src="../../../../../_static/jquery.js"></script>
        <script src="../../../../../_static/underscore.js"></script>
        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../../_static/doctools.js"></script>
        <script src="../../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.common.data_types.enum.upsample_mode</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.common.data_types.enum.upsample_mode</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>


<div class="viewcode-block" id="UpsampleMode"><a class="viewcode-back" href="../../../../../super_gradients.common.html#super_gradients.common.UpsampleMode">[docs]</a><span class="k">class</span> <span class="nc">UpsampleMode</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
    <span class="n">NEAREST</span> <span class="o">=</span> <span class="s2">&quot;nearest&quot;</span>
    <span class="n">BILINEAR</span> <span class="o">=</span> <span class="s2">&quot;bilinear&quot;</span>
    <span class="n">BICUBIC</span> <span class="o">=</span> <span class="s2">&quot;bicubic&quot;</span>
    <span class="n">SNPE_BILINEAR</span> <span class="o">=</span> <span class="s2">&quot;snpe_bilinear&quot;</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.common.decorators.deci_logger &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.common.decorators.deci_logger</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.common.decorators.deci_logger</h1><div class="highlight"><pre>
<div class="viewcode-block" id="deci_func_logger"><a class="viewcode-back" href="../../../../super_gradients.common.decorators.html#super_gradients.common.decorators.deci_logger.deci_func_logger">[docs]</a><span></span><span class="k">def</span> <span class="nf">deci_func_logger</span><span class="p">(</span><span class="n">_func</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="n">name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;abstract_decorator&#39;</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This decorator is used to wrap our functions with logs.</span>
<span class="sd">    It will log every enter and exit of the functon with the equivalent parameters as extras.</span>
<span class="sd">    It will also log exceptions that raises in the function.</span>
<span class="sd">    It will also log the exception time of the function.</span>

<span class="sd">    How it works:`</span>
<span class="sd">        First it will check if the decorator called with name keyword.</span>
<span class="sd">        If so it will return a new decorator that its logger is the name parameter.</span>
<span class="sd">        If not it will return a new decorator that its logger is the wrapped function name.</span>
<span class="sd">        Then the return decorator will return a new function that warps the original function with the new logs.</span>
<span class="sd">        For further understanding advise real-python &quot;fancy decorators documentation&quot;</span>
<span class="sd">    Args:</span>
<span class="sd">        _func (): used when called without name specify. dont pass it directly</span>
<span class="sd">        name (): The name of the logger to save logs by.</span>

<span class="sd">    Returns:</span>
<span class="sd">        a decorator that wraps function with logs logic.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># TODO: Not Working - Breaks the code, tests does not pass (s3 connector, platform...)</span>
    <span class="c1"># TODO: Fix problem with ExplicitParamValidation error (arguments not passed)</span>
    <span class="c1"># TODO: Run ALL test suite of deci2 (NOT circieCI test suite, but ALL the tests under tests folders)</span>
    <span class="c1"># TODO: Delete/Update all failing tests.</span>

    <span class="c1"># def deci_logger_decorator(fn):</span>
    <span class="c1">#</span>
    <span class="c1"># @functools.wraps(fn)</span>
    <span class="c1"># def wrapper_func(*args, **kwargs):</span>
    <span class="c1">#     try:</span>
    <span class="c1">#</span>
    <span class="c1">#         try:</span>
    <span class="c1">#             logger.debug(f&quot;Start: {fn.__name__}&quot;, extra={&quot;args&quot;: args, &quot;kwargs&quot;: kwargs})</span>
    <span class="c1">#             time1 = time.perf_counter()</span>
    <span class="c1">#         except Exception:</span>
    <span class="c1">#             # failed to write log - continue.</span>
    <span class="c1">#             pass</span>
    <span class="c1">#</span>
    <span class="c1">#         result = fn(*args, **kwargs)</span>
    <span class="c1">#</span>
    <span class="c1">#         try:</span>
    <span class="c1">#             time2 = time.perf_counter()</span>
    <span class="c1">#             logger.debug(f&quot;End: {fn.__name__}&quot;,</span>
    <span class="c1">#                          extra={&#39;duration&#39;: (time2 - time1) * 1000.0, &#39;return_value&#39;: result})</span>
    <span class="c1">#         except Exception:</span>
    <span class="c1">#             # failed to write log - continue.</span>
    <span class="c1">#             pass</span>
    <span class="c1">#</span>
    <span class="c1">#         return result</span>
    <span class="c1">#</span>
    <span class="c1">#     except Exception as ex:</span>
    <span class="c1">#         # This exception was raised from inside the function call</span>
    <span class="c1">#         logger.error(f&quot;Exception: {ex}&quot;, exc_info=ex)</span>
    <span class="c1">#         raise ex</span>
    <span class="c1">#</span>
    <span class="c1"># return wrapper_func</span>

    <span class="c1"># if _func is None:</span>
    <span class="c1">#     logger = get_logger(name)</span>
    <span class="c1">#     return deci_logger_decorator</span>
    <span class="c1"># else:</span>
    <span class="c1">#     logger = get_logger(_func.__name__)</span>
    <span class="c1">#     return deci_logger_decorator(_func)</span>

    <span class="k">return</span> <span class="n">_func</span></div>


<div class="viewcode-block" id="deci_class_logger"><a class="viewcode-back" href="../../../../super_gradients.common.decorators.html#super_gradients.common.decorators.deci_logger.deci_class_logger">[docs]</a><span class="k">def</span> <span class="nf">deci_class_logger</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This decorator wraps every class method with deci_func_logger decorator.</span>
<span class="sd">    It works by checking if class method is callable and if so it will set a new decorated method as the same method name.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">wrapper</span><span class="p">(</span><span class="bp">cls</span><span class="p">):</span>
        <span class="c1"># TODO: Not Working - Breaks the code, tests does not pass (s3 connector, platform...)</span>
        <span class="c1"># TODO: Fix problem with ExplicitParamValidation error (arguments not passed)</span>
        <span class="c1"># TODO: Run ALL test suite of deci2 (NOT circieCI test suite, but ALL the tests under tests folders)</span>
        <span class="c1"># TODO: Delete/Update all failing tests.</span>

        <span class="c1"># for attr in cls.__dict__:</span>
        <span class="c1">#     if callable(getattr(cls, attr)) and attr != &#39;__init__&#39;:</span>
        <span class="c1">#         decorated_function = deci_func_logger(name=cls.__name__)(getattr(cls, attr))</span>
        <span class="c1">#         if type(cls.__dict__[attr]) is staticmethod:</span>
        <span class="c1">#             decorated_function = staticmethod(decorated_function)</span>
        <span class="c1">#         setattr(cls, attr, decorated_function)</span>
        <span class="k">return</span> <span class="bp">cls</span>

    <span class="k">return</span> <span class="n">wrapper</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.decorators.explicit_params_validator &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.decorators.explicit_params_validator &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -160,7 +162,7 @@
 
                             <span class="c1"># WRAPS THE RETRY DECORATOR CLASS TO ENABLE CALLING WITHOUT PARAMS</span>
                
 
                            -<div class="viewcode-block" id="explicit_params_validation"><a class="viewcode-back" href="../../../../super_gradients.common.decorators.html#super_gradients.common.explicit_params_validation">[docs]</a><span class="k">def</span> <span class="nf">explicit_params_validation</span><span class="p">(</span><span class="n">function</span><span class="p">:</span> <span class="n">Callable</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">validation_ty
                
 
                            +<div class="viewcode-block" id="explicit_params_validation"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.explicit_params_validation">[docs]</a><span class="k">def</span> <span class="nf">explicit_params_validation</span><span class="p">(</span><span class="n">function</span><span class="p">:</span> <span class="n">Callable</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">validation_type</span><s
                
 
                                 <span class="k">if</span> <span class="n">function</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                     <span class="k">return</span> <span class="n">_ExplicitParamsValidator</span><span class="p">(</span><span class="n">function</span><span class="o">=</span><span class="n">function</span><span class="p">)</span>
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
@@ -197,4 +199,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.decorators.singleton &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.decorators.singleton &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -86,7 +88,7 @@
 
                                        <div itemprop="articleBody">
                
 
                               <h1>Source code for super_gradients.common.decorators.singleton</h1><div class="highlight"><pre>
                
 
                            -<div class="viewcode-block" id="SingletonMeta"><a class="viewcode-back" href="../../../../super_gradients.common.decorators.html#super_gradients.common.SingletonMeta">[docs]</a><span></span><span class="k">class</span> <span class="nc">SingletonMeta</span><span class="p">(</span><span class="nb">type</span><span class="p">):</span>
                
 
                            +<span></span><span class="k">class</span> <span class="nc">SingletonMeta</span><span class="p">(</span><span class="nb">type</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    A Singleton meta class.</span>
                
 
                             <span class="sd">    A class that derives from this class will have only 1 instance of that type for the process.</span>
                
@@ -97,7 +99,7 @@
 
                                 <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                                     <span class="k">if</span> <span class="bp">cls</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">cls</span><span class="o">.</span><span class="n">_instances</span><span class="p">:</span>
                
 
                                         <span class="bp">cls</span><span class="o">.</span><span class="n">_instances</span><span class="p">[</span><span class="bp">cls</span><span class="p">]</span> <span class="o">=</span> <span class="nb">super</span><span class="p">(</span><span class="n">SingletonMeta</span><span class="p">,</span> <span class="bp">cls</span><span class="p">)</span><span class="o">.</span><span class="fm">__call__</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span 
                
 
                            -        <span class="k">return</span> <span class="bp">cls</span><span class="o">.</span><span class="n">_instances</span><span class="p">[</span><span class="bp">cls</span><span class="p">]</span></div>
                
 
                            +        <span class="k">return</span> <span class="bp">cls</span><span class="o">.</span><span class="n">_instances</span><span class="p">[</span><span class="bp">cls</span><span class="p">]</span>
                
 
                             <span class="k">class</span> <span class="nc">_SingletonWrapper</span><span class="p">:</span>
                
@@ -117,7 +119,7 @@
 
                                     <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_instance</span>
                
 
                            -<div class="viewcode-block" id="singleton"><a class="viewcode-back" href="../../../../super_gradients.common.decorators.html#super_gradients.common.singleton">[docs]</a><span class="k">def</span> <span class="nf">singleton</span><span class="p">(</span><span class="bp">cls</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="singleton"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.singleton">[docs]</a><span class="k">def</span> <span class="nf">singleton</span><span class="p">(</span><span class="bp">cls</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    A singleton decorator. Returns a wrapper objects. A call on that object</span>
                
 
                             <span class="sd">    returns a single instance object of decorated class. Use the __wrapped__</span>
                
@@ -153,4 +155,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.common.environment.env_helpers &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.common.environment.env_helpers &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -87,38 +89,56 @@
 
                               <h1>Source code for super_gradients.common.environment.env_helpers</h1><div class="highlight"><pre>
                
 
                             <span></span><span class="kn">import</span> <span class="nn">argparse</span>
                
 
                            +<span class="kn">import</span> <span class="nn">importlib</span>
                
 
                             <span class="kn">import</span> <span class="nn">os</span>
                
 
                            +<span class="kn">import</span> <span class="nn">socket</span>
                
 
                             <span class="kn">import</span> <span class="nn">sys</span>
                
 
                             <span class="kn">from</span> <span class="nn">functools</span> <span class="kn">import</span> <span class="n">wraps</span>
                
 
                            +<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span>
                
 
                            +
                
 
                            +<span class="kn">from</span> <span class="nn">omegaconf</span> <span class="kn">import</span> <span class="n">OmegaConf</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.environment</span> <span class="kn">import</span> <span class="n">environment_config</span>
                
 
                            -<div class="viewcode-block" id="TerminalColours"><a class="viewcode-back" href="../../../../super_gradients.common.environment.html#super_gradients.common.TerminalColours">[docs]</a><span class="k">class</span> <span class="nc">TerminalColours</span><span class="p">:</span>
                
 
                            +<span class="k">class</span> <span class="nc">TerminalColours</span><span class="p">:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Usage: https://stackoverflow.com/questions/287871/how-to-print-colored-text-in-python?page=1&amp;tab=votes#tab-top</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="n">HEADER</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[95m&#39;</span>
                
 
                            -    <span class="n">OKBLUE</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[94m&#39;</span>
                
 
                            -    <span class="n">OKCYAN</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[96m&#39;</span>
                
 
                            -    <span class="n">OKGREEN</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[92m&#39;</span>
                
 
                            -    <span class="n">WARNING</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[93m&#39;</span>
                
 
                            -    <span class="n">FAIL</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[91m&#39;</span>
                
 
                            -    <span class="n">ENDC</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[0m&#39;</span>
                
 
                            -    <span class="n">BOLD</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[1m&#39;</span>
                
 
                            -    <span class="n">UNDERLINE</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="se">\033</span><span class="s1">[4m&#39;</span></div>
                
 
                            -
                
 
                            -
                
 
                            -<div class="viewcode-block" id="ColouredTextFormatter"><a class="viewcode-back" href="../../../../super_gradients.common.environment.html#super_gradients.common.ColouredTextFormatter">[docs]</a><span class="k">class</span> <span class="nc">ColouredTextFormatter</span><span class="p">:</span>
                
 
                            -<div class="viewcode-block" id="ColouredTextFormatter.print_coloured_text"><a class="viewcode-back" href="../../../../super_gradients.common.environment.html#super_gradients.common.ColouredTextFormatter.print_coloured_text">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +
                
 
                            +    <span class="n">HEADER</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[95m&quot;</span>
                
 
                            +    <span class="n">OKBLUE</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[94m&quot;</span>
                
 
                            +    <span class="n">OKCYAN</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[96m&quot;</span>
                
 
                            +    <span class="n">OKGREEN</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[92m&quot;</span>
                
 
                            +    <span class="n">WARNING</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[93m&quot;</span>
                
 
                            +    <span class="n">FAIL</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[91m&quot;</span>
                
 
                            +    <span class="n">ENDC</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[0m&quot;</span>
                
 
                            +    <span class="n">BOLD</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[1m&quot;</span>
                
 
                            +    <span class="n">UNDERLINE</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\033</span><span class="s2">[4m&quot;</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">class</span> <span class="nc">ColouredTextFormatter</span><span class="p">:</span>
                
 
                            +    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">print_coloured_text</span><span class="p">(</span><span class="n">text</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">colour</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Prints a text with colour ascii characters.</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">return</span> <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">colour</span><span class="p">,</span> <span class="n">text</span><span class="p">,</span> <span class="n">TerminalColours</span><span class="o">.</span><span class="n">ENDC</span><span class="p">]))</span></div></div>
                
 
                            +        <span class="k">return</span> <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">colour</span><span class="p">,</span> <span class="n">text</span><span class="p">,</span> <span class="n">TerminalColours</span><span class="o">.</span><span class="n">ENDC</span><span class="p">]))</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">get_cls</span><span class="p">(</span><span class="n">cls_path</span><span class="p">):</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">    A resolver for Hydra/OmegaConf to allow getting a class instead on an instance.</span>
                
 
                            +<span class="sd">    usage:</span>
                
 
                            +<span class="sd">    class_of_optimizer: ${class:torch.optim.Adam}</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +    <span class="n">module</span> <span class="o">=</span> <span class="s2">&quot;.&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">cls_path</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;.&quot;</span><span class="p">)[:</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>
                
 
                            +    <span class="n">name</span> <span class="o">=</span> <span class="n">cls_path</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;.&quot;</span><span class="p">)[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
                
 
                            +    <span class="n">importlib</span><span class="o">.</span><span class="n">import_module</span><span class="p">(</span><span class="n">module</span><span class="p">)</span>
                
 
                            +    <span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">sys</span><span class="o">.</span><span class="n">modules</span><span class="p">[</span><span class="n">module</span><span class="p">],</span> <span class="n">name</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="get_environ_as_type"><a class="viewcode-back" href="../../../../super_gradients.common.environment.html#super_gradients.common.get_environ_as_type">[docs]</a><span class="k">def</span> <span class="nf">get_environ_as_type</span><span class="p">(</span><span class="n">environment_variable_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="kc">None</span><span cl
                
 
                            +<span class="k">def</span> <span class="nf">get_environ_as_type</span><span class="p">(</span><span class="n">environment_variable_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cast_to_type</span><span class="p">:</span> <span class="nb">type</span> <span class="o">=</span> <span class="nb">str</span><span class="p">)</span> <span c
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Tries to get an environment variable and cast it into a requested type.</span>
                
 
                             <span class="sd">    :return: cast_to_type object, or None if failed.</span>
                
@@ -131,41 +151,101 @@
 
                                     <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
                
 
                                         <span class="nb">print</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
                
 
                                         <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                
 
                            -                <span class="sa">f</span><span class="s1">&#39;Failed to cast environment variable </span><span class="si">{</span><span class="n">environment_variable_name</span><span class="si">}</span><span class="s1"> to type </span><span class="si">{</span><span class="n">cast_to_type</span><span class="si">}</span><span class="s1">: the value </span><span class="si">{</span><span class="n">value</span><span class="si">}</span><span class="s1"> is not a valid </span><span class="si">{</spa
                
 
                            -    <span class="k">return</span></div>
                
 
                            +                <span class="sa">f</span><span class="s2">&quot;Failed to cast environment variable </span><span class="si">{</span><span class="n">environment_variable_name</span><span class="si">}</span><span class="s2"> to type </span><span class="si">{</span><span class="n">cast_to_type</span><span class="si">}</span><span class="s2">: the value </span><span class="si">{</span><span class="n">value</span><span class="si">}</span><span class="s2"> is not a valid </span><span class="si">{</sp
                
 
                            +            <span class="p">)</span>
                
 
                            +    <span class="k">return</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">hydra_output_dir_resolver</span><span class="p">(</span><span class="n">ckpt_root_dir</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">):</span>
                
 
                            +    <span class="k">if</span> <span class="n">ckpt_root_dir</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +        <span class="n">output_dir_path</span> <span class="o">=</span> <span class="n">environment_config</span><span class="o">.</span><span class="n">PKG_CHECKPOINTS_DIR</span> <span class="o">+</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">sep</span> <span class="o">+</span> <span class="n">experiment_name</span>
                
 
                            +    <span class="k">else</span><span class="p">:</span>
                
 
                            +        <span class="n">output_dir_path</span> <span class="o">=</span> <span class="n">ckpt_root_dir</span> <span class="o">+</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">sep</span> <span class="o">+</span> <span class="n">experiment_name</span>
                
 
                            +    <span class="k">return</span> <span class="n">output_dir_path</span>
                
 
                            -<div class="viewcode-block" id="init_trainer"><a class="viewcode-back" href="../../../../super_gradients.common.environment.html#super_gradients.common.init_trainer">[docs]</a><span class="k">def</span> <span class="nf">init_trainer</span><span class="p">():</span>
                
 
                            +<div class="viewcode-block" id="init_trainer"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.init_trainer">[docs]</a><span class="k">def</span> <span class="nf">init_trainer</span><span class="p">():</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">    a function to initialize the super_gradients environment. This function should be the first thing to be called</span>
                
 
                            -<span class="sd">    by any code running super_gradients. It resolves conflicts between the different tools, packages and environments used</span>
                
 
                            -<span class="sd">    and prepares the super_gradients environment.</span>
                
 
                            +<span class="sd">    Initialize the super_gradients environment.</span>
                
 
                            +
                
 
                            +<span class="sd">    This function should be the first thing to be called by any code running super_gradients.</span>
                
 
                            +<span class="sd">    It resolves conflicts between the different tools, packages and environments used and prepares the super_gradients environment.</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +    <span class="k">if</span> <span class="ow">not</span> <span class="n">environment_config</span><span class="o">.</span><span class="n">INIT_TRAINER</span><span class="p">:</span>
                
 
                            +        <span class="n">register_hydra_resolvers</span><span class="p">()</span>
                
 
                            +
                
 
                            +        <span class="c1"># We pop local_rank if it was specified in the args, because it would break</span>
                
 
                            +        <span class="n">args_local_rank</span> <span class="o">=</span> <span class="n">pop_arg</span><span class="p">(</span><span class="s2">&quot;local_rank&quot;</span><span class="p">,</span> <span class="n">default_value</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="c1"># Set local_rank with priority order (env variable &gt; args.local_rank &gt; args.default_value)</span>
                
 
                            +        <span class="n">environment_config</span><span class="o">.</span><span class="n">DDP_LOCAL_RANK</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">getenv</span><span class="p">(</span><span class="s2">&quot;LOCAL_RANK&quot;</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="n">args_local_rank</span><span class="p">))</span>
                
 
                            +        <span class="n">environment_config</span><span class="o">.</span><span class="n">INIT_TRAINER</span> <span class="o">=</span> <span class="kc">True</span></div>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">register_hydra_resolvers</span><span class="p">():</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Register all the hydra resolvers required for the super-gradients recipes.&quot;&quot;&quot;</span>
                
 
                            +    <span class="n">OmegaConf</span><span class="o">.</span><span class="n">register_new_resolver</span><span class="p">(</span><span class="s2">&quot;hydra_output_dir&quot;</span><span class="p">,</span> <span class="n">hydra_output_dir_resolver</span><span class="p">,</span> <span class="n">replace</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +    <span class="n">OmegaConf</span><span class="o">.</span><span class="n">register_new_resolver</span><span class="p">(</span><span class="s2">&quot;class&quot;</span><span class="p">,</span> <span class="k">lambda</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">get_cls</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">),</span> <span class="n">replace</span><span class="o">=</span><span class="k
                
 
                            +    <span class="n">OmegaConf</span><span class="o">.</span><span class="n">register_new_resolver</span><span class="p">(</span><span class="s2">&quot;add&quot;</span><span class="p">,</span> <span class="k">lambda</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="nb">sum</span><span class="p">(</span><span class="n">args</span><span class="p">),</span> <span class="n">replace</span><span class="o">=</span><span class="kc">True</span><span class="p"
                
 
                            +    <span class="n">OmegaConf</span><span class="o">.</span><span class="n">register_new_resolver</span><span class="p">(</span><span class="s2">&quot;cond&quot;</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">boolean</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">x</span> <span class="k">if</span> <span class="n">boolean</span> <span class="k">else</span> <span class="
                
 
                            +    <span class="n">OmegaConf</span><span class="o">.</span><span class="n">register_new_resolver</span><span class="p">(</span><span class="s2">&quot;getitem&quot;</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">container</span><span class="p">,</span> <span class="n">key</span><span class="p">:</span> <span class="n">container</span><span class="p">[</span><span class="n">key</span><span class="p">],</span> <span class="n">replace</span><span class="o">=</span><sp
                
 
                            +    <span class="n">OmegaConf</span><span class="o">.</span><span class="n">register_new_resolver</span><span class="p">(</span><span class="s2">&quot;first&quot;</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">lst</span><span class="p">:</span> <span class="n">lst</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">replace</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>  <span class="c1
                
 
                            +    <span class="n">OmegaConf</span><span class="o">.</span><span class="n">register_new_resolver</span><span class="p">(</span><span class="s2">&quot;last&quot;</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">lst</span><span class="p">:</span> <span class="n">lst</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">replace</span><span class="o">=</span><span class="kc">True</span><span class="p">)<
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">pop_arg</span><span class="p">(</span><span class="n">arg_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">default_value</span><span class="p">:</span> <span class="n">Any</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Get the specified args and remove them from argv&quot;&quot;&quot;</span>
                
 
                                 <span class="n">parser</span> <span class="o">=</span> <span class="n">argparse</span><span class="o">.</span><span class="n">ArgumentParser</span><span class="p">()</span>
                
 
                            -    <span class="n">parser</span><span class="o">.</span><span class="n">add_argument</span><span class="p">(</span><span class="s2">&quot;--local_rank&quot;</span><span class="p">,</span> <span class="nb">type</span><span class="o">=</span><span class="nb">int</span><span class="p">,</span> <span class="n">default</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>  <span class="c1"># used by DDP</span>
                
 
                            +    <span class="n">parser</span><span class="o">.</span><span class="n">add_argument</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;--</span><span class="si">{</span><span class="n">arg_name</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="n">default_value</span><span class="p">)</span>
                
 
                                 <span class="n">args</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">parser</span><span class="o">.</span><span class="n">parse_known_args</span><span class="p">()</span>
                
 
                            -    <span class="c1"># remove any flags starting with --local_rank from the argv list</span>
                
 
                            -    <span class="n">to_remove</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">filter</span><span class="p">(</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">x</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;--local_rank&#39;</span><span class="p">),</span> <span class="n">sys</span><span class="o">.</span><span class="n">argv</spa
                
 
                            -    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">to_remove</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                
 
                            -        <span class="k">for</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">to_remove</span><span class="p">:</span>
                
 
                            -            <span class="n">sys</span><span class="o">.</span><span class="n">argv</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
                
 
                            -
                
 
                            -    <span class="n">environment_config</span><span class="o">.</span><span class="n">DDP_LOCAL_RANK</span> <span class="o">=</span> <span class="n">args</span><span class="o">.</span><span class="n">local_rank</span></div>
                
 
                            +    <span class="c1"># Remove the ddp args to not have a conflict with the use of hydra</span>
                
 
                            +    <span class="k">for</span> <span class="n">val</span> <span class="ow">in</span> <span class="nb">filter</span><span class="p">(</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">x</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;--</span><span class="si">{</span><span class="n">arg_name</span><span class="si">}</span><span class="s2">&quot;</span><span cl
                
 
                            +        <span class="n">environment_config</span><span class="o">.</span><span class="n">EXTRA_ARGS</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
                
 
                            +        <span class="n">sys</span><span class="o">.</span><span class="n">argv</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
                
 
                            +    <span class="k">return</span> <span class="nb">vars</span><span class="p">(</span><span class="n">args</span><span class="p">)[</span><span class="n">arg_name</span><span class="p">]</span>
                
 
                            -<div class="viewcode-block" id="is_distributed"><a class="viewcode-back" href="../../../../super_gradients.common.environment.html#super_gradients.common.is_distributed">[docs]</a><span class="k">def</span> <span class="nf">is_distributed</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="is_distributed"><a class="viewcode-back" href="../../../../super_gradients.common.html#super_gradients.common.is_distributed">[docs]</a><span class="k">def</span> <span class="nf">is_distributed</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
                
 
                                 <span class="k">return</span> <span class="n">environment_config</span><span class="o">.</span><span class="n">DDP_LOCAL_RANK</span> <span class="o">&gt;=</span> <span class="mi">0</span></div>
                
 
                            -<div class="viewcode-block" id="multi_process_safe"><a class="viewcode-back" href="../../../../super_gradients.common.environment.html#super_gradients.common.multi_process_safe">[docs]</a><span class="k">def</span> <span class="nf">multi_process_safe</span><span class="p">(</span><span class="n">func</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">is_rank_0</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Check if the node was launched with torch.distributed.launch and if the node is of rank 0&quot;&quot;&quot;</span>
                
 
                            +    <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">getenv</span><span class="p">(</span><span class="s2">&quot;LOCAL_RANK&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;0&quot;</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">is_launched_using_sg</span><span class="p">():</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Check if the current process is a subprocess launched using SG restart_script_with_ddp&quot;&quot;&quot;</span>
                
 
                            +    <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;TORCHELASTIC_RUN_ID&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;sg_initiated&quot;</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">is_main_process</span><span class="p">():</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Check if current process is considered as the main process (i.e. is responsible for sanity check, atexit upload, ...).</span>
                
 
                            +<span class="sd">    The definition ensures that 1 and only 1 process follows this condition, regardless of how the run was started.</span>
                
 
                            +
                
 
                            +<span class="sd">    The rule is as follow:</span>
                
 
                            +<span class="sd">        - If not DDP: main process is current process</span>
                
 
                            +<span class="sd">        - If DDP launched using SuperGradients: main process is the launching process (rank=-1)</span>
                
 
                            +<span class="sd">        - If DDP launched with torch: main process is rank 0</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_distributed</span><span class="p">():</span>  <span class="c1"># If no DDP, or DDP launching process</span>
                
 
                            +        <span class="k">return</span> <span class="kc">True</span>
                
 
                            +    <span class="k">elif</span> <span class="n">is_rank_0</span><span class="p">()</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">is_launched_using_sg</span><span class="p">():</span>  <span class="c1"># If DDP launched using torch.distributed.launch or torchrun, we need to run the check on rank 0</span>
                
 
                            +        <span class="k">return</span> <span class="kc">True</span>
                
 
                            +    <span class="k">else</span><span class="p">:</span>
                
 
                            +        <span class="k">return</span> <span class="kc">False</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">multi_process_safe</span><span class="p">(</span><span class="n">func</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    A decorator for making sure a function runs only in main process.</span>
                
 
                             <span class="sd">    If not in DDP mode (local_rank = -1), the function will run.</span>
                
 
                             <span class="sd">    If in DDP mode, the function will run only in the main process (local_rank = 0)</span>
                
 
                             <span class="sd">    This works only for functions with no return value</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                                 <span class="k">def</span> <span class="nf">do_nothing</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                                     <span class="k">pass</span>
                
@@ -176,7 +256,18 @@
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="k">return</span> <span class="n">do_nothing</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            -    <span class="k">return</span> <span class="n">wrapper</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">wrapper</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">find_free_port</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Find an available port of current machine/node.</span>
                
 
                            +<span class="sd">    Note: there is still a chance the port could be taken by other processes.&quot;&quot;&quot;</span>
                
 
                            +
                
 
                            +    <span class="k">with</span> <span class="n">socket</span><span class="o">.</span><span class="n">socket</span><span class="p">(</span><span class="n">socket</span><span class="o">.</span><span class="n">AF_INET</span><span class="p">,</span> <span class="n">socket</span><span class="o">.</span><span class="n">SOCK_STREAM</span><span class="p">)</span> <span class="k">as</span> <span class="n">sock</span><span class="p">:</span>
                
 
                            +        <span class="c1"># Binding to port 0 will cause the OS to find an available port for us</span>
                
 
                            +        <span class="n">sock</span><span class="o">.</span><span class="n">bind</span><span class="p">((</span><span class="s2">&quot;&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span>
                
 
                            +        <span class="n">_ip</span><span class="p">,</span> <span class="n">port</span> <span class="o">=</span> <span class="n">sock</span><span class="o">.</span><span class="n">getsockname</span><span class="p">()</span>
                
 
                            +    <span class="k">return</span> <span class="n">port</span>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -206,4 +297,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.common.object_names &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../" id="documentation_options" src="../../../_static/documentation_options.js"></script>
        <script src="../../../_static/jquery.js"></script>
        <script src="../../../_static/underscore.js"></script>
        <script src="../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../_static/doctools.js"></script>
        <script src="../../../_static/sphinx_highlight.js"></script>
    <script src="../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../genindex.html" />
    <link rel="search" title="Search" href="../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.common.object_names</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.common.object_names</h1><div class="highlight"><pre>
<div class="viewcode-block" id="Losses"><a class="viewcode-back" href="../../../super_gradients.training.html#super_gradients.training.losses.Losses">[docs]</a><span></span><span class="k">class</span> <span class="nc">Losses</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class holding all the supported loss names&quot;&quot;&quot;</span>

    <span class="n">CROSS_ENTROPY</span> <span class="o">=</span> <span class="s2">&quot;cross_entropy&quot;</span>
    <span class="n">MSE</span> <span class="o">=</span> <span class="s2">&quot;mse&quot;</span>
    <span class="n">R_SQUARED_LOSS</span> <span class="o">=</span> <span class="s2">&quot;r_squared_loss&quot;</span>
    <span class="n">SHELFNET_OHEM_LOSS</span> <span class="o">=</span> <span class="s2">&quot;shelfnet_ohem_loss&quot;</span>
    <span class="n">SHELFNET_SE_LOSS</span> <span class="o">=</span> <span class="s2">&quot;shelfnet_se_loss&quot;</span>
    <span class="n">YOLOX_LOSS</span> <span class="o">=</span> <span class="s2">&quot;yolox_loss&quot;</span>
    <span class="n">YOLOX_FAST_LOSS</span> <span class="o">=</span> <span class="s2">&quot;yolox_fast_loss&quot;</span>
    <span class="n">SSD_LOSS</span> <span class="o">=</span> <span class="s2">&quot;ssd_loss&quot;</span>
    <span class="n">STDC_LOSS</span> <span class="o">=</span> <span class="s2">&quot;stdc_loss&quot;</span>
    <span class="n">BCE_DICE_LOSS</span> <span class="o">=</span> <span class="s2">&quot;bce_dice_loss&quot;</span>
    <span class="n">KD_LOSS</span> <span class="o">=</span> <span class="s2">&quot;kd_loss&quot;</span>
    <span class="n">DICE_CE_EDGE_LOSS</span> <span class="o">=</span> <span class="s2">&quot;dice_ce_edge_loss&quot;</span></div>


<div class="viewcode-block" id="Metrics"><a class="viewcode-back" href="../../../super_gradients.training.html#super_gradients.training.losses.Metrics">[docs]</a><span class="k">class</span> <span class="nc">Metrics</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class holding all the supported metric names&quot;&quot;&quot;</span>

    <span class="n">ACCURACY</span> <span class="o">=</span> <span class="s2">&quot;Accuracy&quot;</span>
    <span class="n">TOP5</span> <span class="o">=</span> <span class="s2">&quot;Top5&quot;</span>
    <span class="n">DETECTION_METRICS</span> <span class="o">=</span> <span class="s2">&quot;DetectionMetrics&quot;</span>
    <span class="n">DETECTION_METRICS_050_095</span> <span class="o">=</span> <span class="s2">&quot;DetectionMetrics_050_095&quot;</span>
    <span class="n">DETECTION_METRICS_050</span> <span class="o">=</span> <span class="s2">&quot;DetectionMetrics_050&quot;</span>
    <span class="n">DETECTION_METRICS_075</span> <span class="o">=</span> <span class="s2">&quot;DetectionMetrics_075&quot;</span>
    <span class="n">IOU</span> <span class="o">=</span> <span class="s2">&quot;IoU&quot;</span>
    <span class="n">BINARY_IOU</span> <span class="o">=</span> <span class="s2">&quot;BinaryIOU&quot;</span>
    <span class="n">DICE</span> <span class="o">=</span> <span class="s2">&quot;Dice&quot;</span>
    <span class="n">BINARY_DICE</span> <span class="o">=</span> <span class="s2">&quot;BinaryDice&quot;</span>
    <span class="n">PIXEL_ACCURACY</span> <span class="o">=</span> <span class="s2">&quot;PixelAccuracy&quot;</span></div>


<div class="viewcode-block" id="Transforms"><a class="viewcode-back" href="../../../super_gradients.training.html#super_gradients.training.losses.Transforms">[docs]</a><span class="k">class</span> <span class="nc">Transforms</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class holding all the supported transform names&quot;&quot;&quot;</span>

    <span class="c1"># From SG</span>
    <span class="n">SegRandomFlip</span> <span class="o">=</span> <span class="s2">&quot;SegRandomFlip&quot;</span>
    <span class="n">SegResize</span> <span class="o">=</span> <span class="s2">&quot;SegResize&quot;</span>
    <span class="n">SegRescale</span> <span class="o">=</span> <span class="s2">&quot;SegRescale&quot;</span>
    <span class="n">SegRandomRescale</span> <span class="o">=</span> <span class="s2">&quot;SegRandomRescale&quot;</span>
    <span class="n">SegRandomRotate</span> <span class="o">=</span> <span class="s2">&quot;SegRandomRotate&quot;</span>
    <span class="n">SegCropImageAndMask</span> <span class="o">=</span> <span class="s2">&quot;SegCropImageAndMask&quot;</span>
    <span class="n">SegRandomGaussianBlur</span> <span class="o">=</span> <span class="s2">&quot;SegRandomGaussianBlur&quot;</span>
    <span class="n">SegPadShortToCropSize</span> <span class="o">=</span> <span class="s2">&quot;SegPadShortToCropSize&quot;</span>
    <span class="n">SegColorJitter</span> <span class="o">=</span> <span class="s2">&quot;SegColorJitter&quot;</span>
    <span class="n">DetectionMosaic</span> <span class="o">=</span> <span class="s2">&quot;DetectionMosaic&quot;</span>
    <span class="n">DetectionRandomAffine</span> <span class="o">=</span> <span class="s2">&quot;DetectionRandomAffine&quot;</span>
    <span class="n">DetectionMixup</span> <span class="o">=</span> <span class="s2">&quot;DetectionMixup&quot;</span>
    <span class="n">DetectionHSV</span> <span class="o">=</span> <span class="s2">&quot;DetectionHSV&quot;</span>
    <span class="n">DetectionHorizontalFlip</span> <span class="o">=</span> <span class="s2">&quot;DetectionHorizontalFlip&quot;</span>
    <span class="n">DetectionPaddedRescale</span> <span class="o">=</span> <span class="s2">&quot;DetectionPaddedRescale&quot;</span>
    <span class="n">DetectionTargetsFormat</span> <span class="o">=</span> <span class="s2">&quot;DetectionTargetsFormat&quot;</span>
    <span class="n">DetectionTargetsFormatTransform</span> <span class="o">=</span> <span class="s2">&quot;DetectionTargetsFormatTransform&quot;</span>
    <span class="n">RandomResizedCropAndInterpolation</span> <span class="o">=</span> <span class="s2">&quot;RandomResizedCropAndInterpolation&quot;</span>
    <span class="n">RandAugmentTransform</span> <span class="o">=</span> <span class="s2">&quot;RandAugmentTransform&quot;</span>
    <span class="n">Lighting</span> <span class="o">=</span> <span class="s2">&quot;Lighting&quot;</span>
    <span class="n">RandomErase</span> <span class="o">=</span> <span class="s2">&quot;RandomErase&quot;</span>

    <span class="c1"># From torch</span>
    <span class="n">Compose</span> <span class="o">=</span> <span class="s2">&quot;Compose&quot;</span>
    <span class="n">ToTensor</span> <span class="o">=</span> <span class="s2">&quot;ToTensor&quot;</span>
    <span class="n">PILToTensor</span> <span class="o">=</span> <span class="s2">&quot;PILToTensor&quot;</span>
    <span class="n">ConvertImageDtype</span> <span class="o">=</span> <span class="s2">&quot;ConvertImageDtype&quot;</span>
    <span class="n">ToPILImage</span> <span class="o">=</span> <span class="s2">&quot;ToPILImage&quot;</span>
    <span class="n">Normalize</span> <span class="o">=</span> <span class="s2">&quot;Normalize&quot;</span>
    <span class="n">Resize</span> <span class="o">=</span> <span class="s2">&quot;Resize&quot;</span>
    <span class="n">CenterCrop</span> <span class="o">=</span> <span class="s2">&quot;CenterCrop&quot;</span>
    <span class="n">Pad</span> <span class="o">=</span> <span class="s2">&quot;Pad&quot;</span>
    <span class="n">Lambda</span> <span class="o">=</span> <span class="s2">&quot;Lambda&quot;</span>
    <span class="n">RandomApply</span> <span class="o">=</span> <span class="s2">&quot;RandomApply&quot;</span>
    <span class="n">RandomChoice</span> <span class="o">=</span> <span class="s2">&quot;RandomChoice&quot;</span>
    <span class="n">RandomOrder</span> <span class="o">=</span> <span class="s2">&quot;RandomOrder&quot;</span>
    <span class="n">RandomCrop</span> <span class="o">=</span> <span class="s2">&quot;RandomCrop&quot;</span>
    <span class="n">RandomHorizontalFlip</span> <span class="o">=</span> <span class="s2">&quot;RandomHorizontalFlip&quot;</span>
    <span class="n">RandomVerticalFlip</span> <span class="o">=</span> <span class="s2">&quot;RandomVerticalFlip&quot;</span>
    <span class="n">RandomResizedCrop</span> <span class="o">=</span> <span class="s2">&quot;RandomResizedCrop&quot;</span>
    <span class="n">FiveCrop</span> <span class="o">=</span> <span class="s2">&quot;FiveCrop&quot;</span>
    <span class="n">TenCrop</span> <span class="o">=</span> <span class="s2">&quot;TenCrop&quot;</span>
    <span class="n">LinearTransformation</span> <span class="o">=</span> <span class="s2">&quot;LinearTransformation&quot;</span>
    <span class="n">ColorJitter</span> <span class="o">=</span> <span class="s2">&quot;ColorJitter&quot;</span>
    <span class="n">RandomRotation</span> <span class="o">=</span> <span class="s2">&quot;RandomRotation&quot;</span>
    <span class="n">RandomAffine</span> <span class="o">=</span> <span class="s2">&quot;RandomAffine&quot;</span>
    <span class="n">Grayscale</span> <span class="o">=</span> <span class="s2">&quot;Grayscale&quot;</span>
    <span class="n">RandomGrayscale</span> <span class="o">=</span> <span class="s2">&quot;RandomGrayscale&quot;</span>
    <span class="n">RandomPerspective</span> <span class="o">=</span> <span class="s2">&quot;RandomPerspective&quot;</span>
    <span class="n">RandomErasing</span> <span class="o">=</span> <span class="s2">&quot;RandomErasing&quot;</span>
    <span class="n">GaussianBlur</span> <span class="o">=</span> <span class="s2">&quot;GaussianBlur&quot;</span>
    <span class="n">InterpolationMode</span> <span class="o">=</span> <span class="s2">&quot;InterpolationMode&quot;</span>
    <span class="n">RandomInvert</span> <span class="o">=</span> <span class="s2">&quot;RandomInvert&quot;</span>
    <span class="n">RandomPosterize</span> <span class="o">=</span> <span class="s2">&quot;RandomPosterize&quot;</span>
    <span class="n">RandomSolarize</span> <span class="o">=</span> <span class="s2">&quot;RandomSolarize&quot;</span>
    <span class="n">RandomAdjustSharpness</span> <span class="o">=</span> <span class="s2">&quot;RandomAdjustSharpness&quot;</span>
    <span class="n">RandomAutocontrast</span> <span class="o">=</span> <span class="s2">&quot;RandomAutocontrast&quot;</span>
    <span class="n">RandomEqualize</span> <span class="o">=</span> <span class="s2">&quot;RandomEqualize&quot;</span></div>


<span class="k">class</span> <span class="nc">Optimizers</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class holding all the supported optimizer names&quot;&quot;&quot;</span>

    <span class="n">SGD</span> <span class="o">=</span> <span class="s2">&quot;SGD&quot;</span>
    <span class="n">ADAM</span> <span class="o">=</span> <span class="s2">&quot;Adam&quot;</span>
    <span class="n">RMS_PROP</span> <span class="o">=</span> <span class="s2">&quot;RMSprop&quot;</span>
    <span class="n">RMS_PROP_TF</span> <span class="o">=</span> <span class="s2">&quot;RMSpropTF&quot;</span>
    <span class="n">LAMB</span> <span class="o">=</span> <span class="s2">&quot;Lamb&quot;</span>


<span class="k">class</span> <span class="nc">Callbacks</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class holding all the supported callback names&quot;&quot;&quot;</span>

    <span class="n">DECI_LAB_UPLOAD</span> <span class="o">=</span> <span class="s2">&quot;DeciLabUploadCallback&quot;</span>
    <span class="n">LR_CALLBACK_BASE</span> <span class="o">=</span> <span class="s2">&quot;LRCallbackBase&quot;</span>
    <span class="n">LR_SCHEDULER</span> <span class="o">=</span> <span class="s2">&quot;LRSchedulerCallback&quot;</span>
    <span class="n">METRICS_UPDATE</span> <span class="o">=</span> <span class="s2">&quot;MetricsUpdateCallback&quot;</span>
    <span class="n">MODEL_CONVERSION_CHECK</span> <span class="o">=</span> <span class="s2">&quot;ModelConversionCheckCallback&quot;</span>
    <span class="n">EARLY_STOP</span> <span class="o">=</span> <span class="s2">&quot;EarlyStop&quot;</span>
    <span class="n">DETECTION_MULTISCALE_PREPREDICTION</span> <span class="o">=</span> <span class="s2">&quot;DetectionMultiscalePrePredictionCallback&quot;</span>
    <span class="n">YOLOX_TRAINING_STAGE_SWITCH</span> <span class="o">=</span> <span class="s2">&quot;YoloXTrainingStageSwitchCallback&quot;</span>


<span class="k">class</span> <span class="nc">LRSchedulers</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class to hold all the supported LR Scheduler names&quot;&quot;&quot;</span>

    <span class="n">STEP</span> <span class="o">=</span> <span class="s2">&quot;step&quot;</span>
    <span class="n">POLY</span> <span class="o">=</span> <span class="s2">&quot;poly&quot;</span>
    <span class="n">COSINE</span> <span class="o">=</span> <span class="s2">&quot;cosine&quot;</span>
    <span class="n">EXP</span> <span class="o">=</span> <span class="s2">&quot;exp&quot;</span>
    <span class="n">FUNCTION</span> <span class="o">=</span> <span class="s2">&quot;function&quot;</span>


<span class="k">class</span> <span class="nc">LRWarmups</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class to hold all the supported LR Warmup names&quot;&quot;&quot;</span>

    <span class="n">LINEAR_STEP</span> <span class="o">=</span> <span class="s2">&quot;linear_step&quot;</span>


<span class="k">class</span> <span class="nc">Samplers</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class to hold all the supported Samplers names&quot;&quot;&quot;</span>

    <span class="n">INFINITE</span> <span class="o">=</span> <span class="s2">&quot;InfiniteSampler&quot;</span>
    <span class="n">REPEAT_AUG</span> <span class="o">=</span> <span class="s2">&quot;RepeatAugSampler&quot;</span>
    <span class="n">DISTRIBUTED</span> <span class="o">=</span> <span class="s2">&quot;DistributedSampler&quot;</span>


<span class="k">class</span> <span class="nc">ContextModules</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class to hold all the segmentation context module names&quot;&quot;&quot;</span>

    <span class="n">ASPP</span> <span class="o">=</span> <span class="s2">&quot;ASPP&quot;</span>
    <span class="n">SPPM</span> <span class="o">=</span> <span class="s2">&quot;SPPM&quot;</span>


<span class="k">class</span> <span class="nc">Models</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Static class to hold all the available model names&quot;&quot;&quot;</span>

    <span class="n">RESNET18</span> <span class="o">=</span> <span class="s2">&quot;resnet18&quot;</span>
    <span class="n">RESNET34</span> <span class="o">=</span> <span class="s2">&quot;resnet34&quot;</span>
    <span class="n">RESNET50_3343</span> <span class="o">=</span> <span class="s2">&quot;resnet50_3343&quot;</span>
    <span class="n">RESNET50</span> <span class="o">=</span> <span class="s2">&quot;resnet50&quot;</span>
    <span class="n">RESNET101</span> <span class="o">=</span> <span class="s2">&quot;resnet101&quot;</span>
    <span class="n">RESNET152</span> <span class="o">=</span> <span class="s2">&quot;resnet152&quot;</span>
    <span class="n">RESNET18_CIFAR</span> <span class="o">=</span> <span class="s2">&quot;resnet18_cifar&quot;</span>
    <span class="n">CUSTOM_RESNET</span> <span class="o">=</span> <span class="s2">&quot;custom_resnet&quot;</span>
    <span class="n">CUSTOM_RESNET50</span> <span class="o">=</span> <span class="s2">&quot;custom_resnet50&quot;</span>
    <span class="n">CUSTOM_RESNET_CIFAR</span> <span class="o">=</span> <span class="s2">&quot;custom_resnet_cifar&quot;</span>
    <span class="n">CUSTOM_RESNET50_CIFAR</span> <span class="o">=</span> <span class="s2">&quot;custom_resnet50_cifar&quot;</span>
    <span class="n">MOBILENET_V2</span> <span class="o">=</span> <span class="s2">&quot;mobilenet_v2&quot;</span>
    <span class="n">MOBILE_NET_V2_135</span> <span class="o">=</span> <span class="s2">&quot;mobile_net_v2_135&quot;</span>
    <span class="n">CUSTOM_MOBILENET_V2</span> <span class="o">=</span> <span class="s2">&quot;custom_mobilenet_v2&quot;</span>
    <span class="n">MOBILENET_V3_LARGE</span> <span class="o">=</span> <span class="s2">&quot;mobilenet_v3_large&quot;</span>
    <span class="n">MOBILENET_V3_SMALL</span> <span class="o">=</span> <span class="s2">&quot;mobilenet_v3_small&quot;</span>
    <span class="n">MOBILENET_V3_CUSTOM</span> <span class="o">=</span> <span class="s2">&quot;mobilenet_v3_custom&quot;</span>
    <span class="n">CUSTOM_DENSENET</span> <span class="o">=</span> <span class="s2">&quot;custom_densenet&quot;</span>
    <span class="n">DENSENET121</span> <span class="o">=</span> <span class="s2">&quot;densenet121&quot;</span>
    <span class="n">DENSENET161</span> <span class="o">=</span> <span class="s2">&quot;densenet161&quot;</span>
    <span class="n">DENSENET169</span> <span class="o">=</span> <span class="s2">&quot;densenet169&quot;</span>
    <span class="n">DENSENET201</span> <span class="o">=</span> <span class="s2">&quot;densenet201&quot;</span>
    <span class="n">SHELFNET18_LW</span> <span class="o">=</span> <span class="s2">&quot;shelfnet18_lw&quot;</span>
    <span class="n">SHELFNET34_LW</span> <span class="o">=</span> <span class="s2">&quot;shelfnet34_lw&quot;</span>
    <span class="n">SHELFNET50_3343</span> <span class="o">=</span> <span class="s2">&quot;shelfnet50_3343&quot;</span>
    <span class="n">SHELFNET50</span> <span class="o">=</span> <span class="s2">&quot;shelfnet50&quot;</span>
    <span class="n">SHELFNET101</span> <span class="o">=</span> <span class="s2">&quot;shelfnet101&quot;</span>
    <span class="n">SHUFFLENET_V2_X0_5</span> <span class="o">=</span> <span class="s2">&quot;shufflenet_v2_x0_5&quot;</span>
    <span class="n">SHUFFLENET_V2_X1_0</span> <span class="o">=</span> <span class="s2">&quot;shufflenet_v2_x1_0&quot;</span>
    <span class="n">SHUFFLENET_V2_X1_5</span> <span class="o">=</span> <span class="s2">&quot;shufflenet_v2_x1_5&quot;</span>
    <span class="n">SHUFFLENET_V2_X2_0</span> <span class="o">=</span> <span class="s2">&quot;shufflenet_v2_x2_0&quot;</span>
    <span class="n">SHUFFLENET_V2_CUSTOM5</span> <span class="o">=</span> <span class="s2">&quot;shufflenet_v2_custom5&quot;</span>
    <span class="n">DARKNET53</span> <span class="o">=</span> <span class="s2">&quot;darknet53&quot;</span>
    <span class="n">CSP_DARKNET53</span> <span class="o">=</span> <span class="s2">&quot;csp_darknet53&quot;</span>
    <span class="n">RESNEXT50</span> <span class="o">=</span> <span class="s2">&quot;resnext50&quot;</span>
    <span class="n">RESNEXT101</span> <span class="o">=</span> <span class="s2">&quot;resnext101&quot;</span>
    <span class="n">GOOGLENET_V1</span> <span class="o">=</span> <span class="s2">&quot;googlenet_v1&quot;</span>
    <span class="n">EFFICIENTNET_B0</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b0&quot;</span>
    <span class="n">EFFICIENTNET_B1</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b1&quot;</span>
    <span class="n">EFFICIENTNET_B2</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b2&quot;</span>
    <span class="n">EFFICIENTNET_B3</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b3&quot;</span>
    <span class="n">EFFICIENTNET_B4</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b4&quot;</span>
    <span class="n">EFFICIENTNET_B5</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b5&quot;</span>
    <span class="n">EFFICIENTNET_B6</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b6&quot;</span>
    <span class="n">EFFICIENTNET_B7</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b7&quot;</span>
    <span class="n">EFFICIENTNET_B8</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_b8&quot;</span>
    <span class="n">EFFICIENTNET_L2</span> <span class="o">=</span> <span class="s2">&quot;efficientnet_l2&quot;</span>
    <span class="n">CUSTOMIZEDEFFICIENTNET</span> <span class="o">=</span> <span class="s2">&quot;CustomizedEfficientnet&quot;</span>
    <span class="n">REGNETY200</span> <span class="o">=</span> <span class="s2">&quot;regnetY200&quot;</span>
    <span class="n">REGNETY400</span> <span class="o">=</span> <span class="s2">&quot;regnetY400&quot;</span>
    <span class="n">REGNETY600</span> <span class="o">=</span> <span class="s2">&quot;regnetY600&quot;</span>
    <span class="n">REGNETY800</span> <span class="o">=</span> <span class="s2">&quot;regnetY800&quot;</span>
    <span class="n">CUSTOM_REGNET</span> <span class="o">=</span> <span class="s2">&quot;custom_regnet&quot;</span>
    <span class="n">CUSTOM_ANYNET</span> <span class="o">=</span> <span class="s2">&quot;custom_anynet&quot;</span>
    <span class="n">NAS_REGNET</span> <span class="o">=</span> <span class="s2">&quot;nas_regnet&quot;</span>
    <span class="n">YOLOX_N</span> <span class="o">=</span> <span class="s2">&quot;yolox_n&quot;</span>
    <span class="n">YOLOX_T</span> <span class="o">=</span> <span class="s2">&quot;yolox_t&quot;</span>
    <span class="n">YOLOX_S</span> <span class="o">=</span> <span class="s2">&quot;yolox_s&quot;</span>
    <span class="n">YOLOX_M</span> <span class="o">=</span> <span class="s2">&quot;yolox_m&quot;</span>
    <span class="n">YOLOX_L</span> <span class="o">=</span> <span class="s2">&quot;yolox_l&quot;</span>
    <span class="n">YOLOX_X</span> <span class="o">=</span> <span class="s2">&quot;yolox_x&quot;</span>
    <span class="n">CUSTOM_YOLO_X</span> <span class="o">=</span> <span class="s2">&quot;custom_yolox&quot;</span>
    <span class="n">SSD_MOBILENET_V1</span> <span class="o">=</span> <span class="s2">&quot;ssd_mobilenet_v1&quot;</span>
    <span class="n">SSD_LITE_MOBILENET_V2</span> <span class="o">=</span> <span class="s2">&quot;ssd_lite_mobilenet_v2&quot;</span>
    <span class="n">REPVGG_A0</span> <span class="o">=</span> <span class="s2">&quot;repvgg_a0&quot;</span>
    <span class="n">REPVGG_A1</span> <span class="o">=</span> <span class="s2">&quot;repvgg_a1&quot;</span>
    <span class="n">REPVGG_A2</span> <span class="o">=</span> <span class="s2">&quot;repvgg_a2&quot;</span>
    <span class="n">REPVGG_B0</span> <span class="o">=</span> <span class="s2">&quot;repvgg_b0&quot;</span>
    <span class="n">REPVGG_B1</span> <span class="o">=</span> <span class="s2">&quot;repvgg_b1&quot;</span>
    <span class="n">REPVGG_B2</span> <span class="o">=</span> <span class="s2">&quot;repvgg_b2&quot;</span>
    <span class="n">REPVGG_B3</span> <span class="o">=</span> <span class="s2">&quot;repvgg_b3&quot;</span>
    <span class="n">REPVGG_D2SE</span> <span class="o">=</span> <span class="s2">&quot;repvgg_d2se&quot;</span>
    <span class="n">REPVGG_CUSTOM</span> <span class="o">=</span> <span class="s2">&quot;repvgg_custom&quot;</span>
    <span class="n">DDRNET_23</span> <span class="o">=</span> <span class="s2">&quot;ddrnet_23&quot;</span>
    <span class="n">DDRNET_23_SLIM</span> <span class="o">=</span> <span class="s2">&quot;ddrnet_23_slim&quot;</span>
    <span class="n">CUSTOM_DDRNET_23</span> <span class="o">=</span> <span class="s2">&quot;custom_ddrnet_23&quot;</span>
    <span class="n">STDC1_CLASSIFICATION</span> <span class="o">=</span> <span class="s2">&quot;stdc1_classification&quot;</span>
    <span class="n">STDC2_CLASSIFICATION</span> <span class="o">=</span> <span class="s2">&quot;stdc2_classification&quot;</span>
    <span class="n">STDC1_SEG</span> <span class="o">=</span> <span class="s2">&quot;stdc1_seg&quot;</span>
    <span class="n">STDC1_SEG50</span> <span class="o">=</span> <span class="s2">&quot;stdc1_seg50&quot;</span>
    <span class="n">STDC1_SEG75</span> <span class="o">=</span> <span class="s2">&quot;stdc1_seg75&quot;</span>
    <span class="n">STDC2_SEG</span> <span class="o">=</span> <span class="s2">&quot;stdc2_seg&quot;</span>
    <span class="n">STDC2_SEG50</span> <span class="o">=</span> <span class="s2">&quot;stdc2_seg50&quot;</span>
    <span class="n">STDC2_SEG75</span> <span class="o">=</span> <span class="s2">&quot;stdc2_seg75&quot;</span>
    <span class="n">CUSTOM_STDC</span> <span class="o">=</span> <span class="s2">&quot;custom_stdc&quot;</span>
    <span class="n">REGSEG48</span> <span class="o">=</span> <span class="s2">&quot;regseg48&quot;</span>
    <span class="n">KD_MODULE</span> <span class="o">=</span> <span class="s2">&quot;kd_module&quot;</span>
    <span class="n">VIT_BASE</span> <span class="o">=</span> <span class="s2">&quot;vit_base&quot;</span>
    <span class="n">VIT_LARGE</span> <span class="o">=</span> <span class="s2">&quot;vit_large&quot;</span>
    <span class="n">VIT_HUGE</span> <span class="o">=</span> <span class="s2">&quot;vit_huge&quot;</span>
    <span class="n">BEIT_BASE_PATCH16_224</span> <span class="o">=</span> <span class="s2">&quot;beit_base_patch16_224&quot;</span>
    <span class="n">BEIT_LARGE_PATCH16_224</span> <span class="o">=</span> <span class="s2">&quot;beit_large_patch16_224&quot;</span>
    <span class="n">PP_LITE_T_SEG</span> <span class="o">=</span> <span class="s2">&quot;pp_lite_t_seg&quot;</span>
    <span class="n">PP_LITE_T_SEG50</span> <span class="o">=</span> <span class="s2">&quot;pp_lite_t_seg50&quot;</span>
    <span class="n">PP_LITE_T_SEG75</span> <span class="o">=</span> <span class="s2">&quot;pp_lite_t_seg75&quot;</span>
    <span class="n">PP_LITE_B_SEG</span> <span class="o">=</span> <span class="s2">&quot;pp_lite_b_seg&quot;</span>
    <span class="n">PP_LITE_B_SEG50</span> <span class="o">=</span> <span class="s2">&quot;pp_lite_b_seg50&quot;</span>
    <span class="n">PP_LITE_B_SEG75</span> <span class="o">=</span> <span class="s2">&quot;pp_lite_b_seg75&quot;</span>
    <span class="n">UNET_CUSTOM</span> <span class="o">=</span> <span class="s2">&quot;unet_custom&quot;</span>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.dataloaders.dataloaders &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.dataloaders.dataloaders</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.dataloaders.dataloaders</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os.path</span>
<span class="kn">import</span> <span class="nn">pkg_resources</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span>

<span class="kn">import</span> <span class="nn">hydra</span>
<span class="kn">from</span> <span class="nn">hydra</span> <span class="kn">import</span> <span class="n">compose</span><span class="p">,</span> <span class="n">initialize_config_dir</span>
<span class="kn">from</span> <span class="nn">hydra.core.global_hydra</span> <span class="kn">import</span> <span class="n">GlobalHydra</span>

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">BatchSampler</span><span class="p">,</span> <span class="n">DataLoader</span><span class="p">,</span> <span class="n">TensorDataset</span>

<span class="kn">import</span> <span class="nn">super_gradients</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.detection_datasets.pascal_voc_detection</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">PascalVOCUnifiedDetectionTrainDataset</span><span class="p">,</span>
    <span class="n">PascalVOCDetectionDataset</span><span class="p">,</span>
<span class="p">)</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">get_param</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.hydra_utils</span> <span class="kn">import</span> <span class="n">normalize_path</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets</span> <span class="kn">import</span> <span class="n">ImageNetDataset</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.detection_datasets</span> <span class="kn">import</span> <span class="n">COCODetectionDataset</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.classification_datasets.cifar</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">Cifar10</span><span class="p">,</span>
    <span class="n">Cifar100</span><span class="p">,</span>
<span class="p">)</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">CityscapesDataset</span><span class="p">,</span>
    <span class="n">CoCoSegmentationDataSet</span><span class="p">,</span>
    <span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">,</span>
    <span class="n">PascalVOCAndAUGUnifiedDataset</span><span class="p">,</span>
    <span class="n">SuperviselyPersonsDataset</span><span class="p">,</span>
<span class="p">)</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.factories.samplers_factory</span> <span class="kn">import</span> <span class="n">SamplersFactory</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.distributed_training_utils</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">wait_for_the_master</span><span class="p">,</span>
    <span class="n">get_local_rank</span><span class="p">,</span>
<span class="p">)</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">override_default_params_without_nones</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>


<div class="viewcode-block" id="get_data_loader"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.get_data_loader">[docs]</a><span class="k">def</span> <span class="nf">get_data_loader</span><span class="p">(</span><span class="n">config_name</span><span class="p">,</span> <span class="n">dataset_cls</span><span class="p">,</span> <span class="n">train</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Class for creating dataloaders for taking defaults from yaml files in src/super_gradients/recipes.</span>

<span class="sd">    :param config_name: yaml config filename in recipes (for example coco2017_yolox).</span>
<span class="sd">    :param dataset_cls: torch dataset uninitialized class.</span>
<span class="sd">    :param train: controls whether to take</span>
<span class="sd">        cfg.dataset_params.train_dataloader_params or cfg.dataset_params.valid_dataloader_params as defaults for the dataset constructor</span>
<span class="sd">     and</span>
<span class="sd">        cfg.dataset_params.train_dataset_params or cfg.dataset_params.valid_dataset_params as defaults for DataLoader contructor.</span>

<span class="sd">    :param dataset_params: dataset params that override the yaml configured defaults, then passed to the dataset_cls.__init__.</span>
<span class="sd">    :param dataloader_params: DataLoader params that override the yaml configured defaults, then passed to the DataLoader.__init__</span>
<span class="sd">    :return: DataLoader</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">dataloader_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">dataloader_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">dataset_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">dataset_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>

    <span class="n">GlobalHydra</span><span class="o">.</span><span class="n">instance</span><span class="p">()</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
    <span class="n">sg_recipes_dir</span> <span class="o">=</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s2">&quot;super_gradients.recipes&quot;</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
    <span class="n">dataset_config</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s2">&quot;dataset_params&quot;</span><span class="p">,</span> <span class="n">config_name</span><span class="p">)</span>
    <span class="k">with</span> <span class="n">initialize_config_dir</span><span class="p">(</span><span class="n">config_dir</span><span class="o">=</span><span class="n">normalize_path</span><span class="p">(</span><span class="n">sg_recipes_dir</span><span class="p">),</span> <span class="n">version_base</span><span class="o">=</span><span class="s2">&quot;1.2&quot;</span><span class="p">):</span>
        <span class="c1"># config is relative to a module</span>
        <span class="n">cfg</span> <span class="o">=</span> <span class="n">compose</span><span class="p">(</span><span class="n">config_name</span><span class="o">=</span><span class="n">normalize_path</span><span class="p">(</span><span class="n">dataset_config</span><span class="p">))</span>

        <span class="n">dataset_params</span> <span class="o">=</span> <span class="n">_process_dataset_params</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">dataset_params</span><span class="p">,</span> <span class="n">train</span><span class="p">)</span>

        <span class="n">local_rank</span> <span class="o">=</span> <span class="n">get_local_rank</span><span class="p">()</span>
        <span class="k">with</span> <span class="n">wait_for_the_master</span><span class="p">(</span><span class="n">local_rank</span><span class="p">):</span>
            <span class="n">dataset</span> <span class="o">=</span> <span class="n">dataset_cls</span><span class="p">(</span><span class="o">**</span><span class="n">dataset_params</span><span class="p">)</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="s2">&quot;dataset_params&quot;</span><span class="p">):</span>
                <span class="n">dataset</span><span class="o">.</span><span class="n">dataset_params</span> <span class="o">=</span> <span class="n">dataset_params</span>

        <span class="n">dataloader_params</span> <span class="o">=</span> <span class="n">_process_dataloader_params</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">,</span> <span class="n">dataset</span><span class="p">,</span> <span class="n">train</span><span class="p">)</span>

        <span class="n">dataloader</span> <span class="o">=</span> <span class="n">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="o">=</span><span class="n">dataset</span><span class="p">,</span> <span class="o">**</span><span class="n">dataloader_params</span><span class="p">)</span>
        <span class="n">dataloader</span><span class="o">.</span><span class="n">dataloader_params</span> <span class="o">=</span> <span class="n">dataloader_params</span>
        <span class="k">return</span> <span class="n">dataloader</span></div>


<span class="k">def</span> <span class="nf">_process_dataset_params</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">dataset_params</span><span class="p">,</span> <span class="n">train</span><span class="p">):</span>
    <span class="n">default_dataset_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_dataset_params</span> <span class="k">if</span> <span class="n">train</span> <span class="k">else</span> <span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_dataset_params</span>
    <span class="n">default_dataset_params</span> <span class="o">=</span> <span class="n">hydra</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">instantiate</span><span class="p">(</span><span class="n">default_dataset_params</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">default_dataset_params</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">dataset_params</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> <span class="ow">or</span> <span class="n">dataset_params</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">dataset_params</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">val</span>

    <span class="k">return</span> <span class="n">dataset_params</span>


<span class="k">def</span> <span class="nf">_process_dataloader_params</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">,</span> <span class="n">dataset</span><span class="p">,</span> <span class="n">train</span><span class="p">):</span>
    <span class="n">default_dataloader_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_dataloader_params</span> <span class="k">if</span> <span class="n">train</span> <span class="k">else</span> <span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_dataloader_params</span>
    <span class="n">default_dataloader_params</span> <span class="o">=</span> <span class="n">hydra</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">instantiate</span><span class="p">(</span><span class="n">default_dataloader_params</span><span class="p">)</span>
    <span class="n">dataloader_params</span> <span class="o">=</span> <span class="n">_process_sampler_params</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">,</span> <span class="n">dataset</span><span class="p">,</span> <span class="n">default_dataloader_params</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">dataloader_params</span>


<span class="k">def</span> <span class="nf">_process_sampler_params</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">,</span> <span class="n">dataset</span><span class="p">,</span> <span class="n">default_dataloader_params</span><span class="p">):</span>
    <span class="n">is_dist</span> <span class="o">=</span> <span class="n">super_gradients</span><span class="o">.</span><span class="n">is_distributed</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">get_param</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">,</span> <span class="s2">&quot;sampler&quot;</span><span class="p">)</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">dataloader_params</span> <span class="o">=</span> <span class="n">_instantiate_sampler</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">get_param</span><span class="p">(</span><span class="n">default_dataloader_params</span><span class="p">,</span> <span class="s2">&quot;sampler&quot;</span><span class="p">)</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">default_dataloader_params</span> <span class="o">=</span> <span class="n">_instantiate_sampler</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">default_dataloader_params</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">is_dist</span><span class="p">:</span>
        <span class="n">default_dataloader_params</span><span class="p">[</span><span class="s2">&quot;sampler&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;DistributedSampler&quot;</span><span class="p">:</span> <span class="p">{}}</span>
        <span class="n">default_dataloader_params</span> <span class="o">=</span> <span class="n">_instantiate_sampler</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">default_dataloader_params</span><span class="p">)</span>
    <span class="n">dataloader_params</span> <span class="o">=</span> <span class="n">override_default_params_without_nones</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">,</span> <span class="n">default_dataloader_params</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">get_param</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">,</span> <span class="s2">&quot;batch_sampler&quot;</span><span class="p">):</span>
        <span class="n">sampler</span> <span class="o">=</span> <span class="n">dataloader_params</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;sampler&quot;</span><span class="p">)</span>
        <span class="n">batch_size</span> <span class="o">=</span> <span class="n">dataloader_params</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;batch_size&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="s2">&quot;drop_last&quot;</span> <span class="ow">in</span> <span class="n">dataloader_params</span><span class="p">:</span>
            <span class="n">drop_last</span> <span class="o">=</span> <span class="n">dataloader_params</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;drop_last&quot;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">drop_last</span> <span class="o">=</span> <span class="n">default_dataloader_params</span><span class="p">[</span><span class="s2">&quot;drop_last&quot;</span><span class="p">]</span>
        <span class="n">dataloader_params</span><span class="p">[</span><span class="s2">&quot;batch_sampler&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">BatchSampler</span><span class="p">(</span><span class="n">sampler</span><span class="o">=</span><span class="n">sampler</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">drop_last</span><span class="o">=</span><span class="n">drop_last</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">dataloader_params</span>


<span class="k">def</span> <span class="nf">_instantiate_sampler</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">):</span>
    <span class="n">sampler_name</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">[</span><span class="s2">&quot;sampler&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">keys</span><span class="p">())[</span><span class="mi">0</span><span class="p">]</span>
    <span class="n">dataloader_params</span><span class="p">[</span><span class="s2">&quot;sampler&quot;</span><span class="p">][</span><span class="n">sampler_name</span><span class="p">][</span><span class="s2">&quot;dataset&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">dataset</span>
    <span class="n">dataloader_params</span><span class="p">[</span><span class="s2">&quot;sampler&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">SamplersFactory</span><span class="p">()</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">[</span><span class="s2">&quot;sampler&quot;</span><span class="p">])</span>
    <span class="k">return</span> <span class="n">dataloader_params</span>


<div class="viewcode-block" id="coco2017_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco2017_train">[docs]</a><span class="k">def</span> <span class="nf">coco2017_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;coco_detection_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">COCODetectionDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="coco2017_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco2017_val">[docs]</a><span class="k">def</span> <span class="nf">coco2017_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;coco_detection_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">COCODetectionDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="coco2017_train_yolox"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco2017_train_yolox">[docs]</a><span class="k">def</span> <span class="nf">coco2017_train_yolox</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">coco2017_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="coco2017_val_yolox"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco2017_val_yolox">[docs]</a><span class="k">def</span> <span class="nf">coco2017_val_yolox</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">coco2017_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="coco2017_train_ssd_lite_mobilenet_v2"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco2017_train_ssd_lite_mobilenet_v2">[docs]</a><span class="k">def</span> <span class="nf">coco2017_train_ssd_lite_mobilenet_v2</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;coco_detection_ssd_lite_mobilenet_v2_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">COCODetectionDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="coco2017_val_ssd_lite_mobilenet_v2"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco2017_val_ssd_lite_mobilenet_v2">[docs]</a><span class="k">def</span> <span class="nf">coco2017_val_ssd_lite_mobilenet_v2</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;coco_detection_ssd_lite_mobilenet_v2_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">COCODetectionDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_dataset_params&quot;</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="n">config_name</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">ImageNetDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_dataset_params&quot;</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="n">config_name</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">ImageNetDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_efficientnet_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_efficientnet_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_efficientnet_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_train</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_efficientnet_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_efficientnet_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_efficientnet_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_efficientnet_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_val</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_efficientnet_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv2_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_mobilenetv2_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv2_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_train</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_mobilenetv2_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv2_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_mobilenetv2_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv2_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_val</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_mobilenetv2_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv3_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_mobilenetv3_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv3_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_train</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_mobilenetv3_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv3_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_mobilenetv3_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv3_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_val</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_mobilenetv3_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_regnetY_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_regnetY_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_regnetY_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">,</span> <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_regnetY_dataset_params&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_regnetY_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_regnetY_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_regnetY_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">,</span> <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_regnetY_dataset_params&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_resnet50_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_resnet50_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_resnet50_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_train</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_resnet50_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_resnet50_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_resnet50_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_resnet50_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_val</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_resnet50_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_resnet50_kd_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_resnet50_kd_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_resnet50_kd_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_train</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_resnet50_kd_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_resnet50_kd_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_resnet50_kd_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_resnet50_kd_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_val</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_resnet50_kd_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_vit_base_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_vit_base_train">[docs]</a><span class="k">def</span> <span class="nf">imagenet_vit_base_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_train</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_vit_base_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_vit_base_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.imagenet_vit_base_val">[docs]</a><span class="k">def</span> <span class="nf">imagenet_vit_base_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">imagenet_val</span><span class="p">(</span>
        <span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="p">,</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;imagenet_vit_base_dataset_params&quot;</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="tiny_imagenet_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.tiny_imagenet_train">[docs]</a><span class="k">def</span> <span class="nf">tiny_imagenet_train</span><span class="p">(</span>
    <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;tiny_imagenet_dataset_params&quot;</span><span class="p">,</span>
<span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="n">config_name</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">ImageNetDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="tiny_imagenet_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.tiny_imagenet_val">[docs]</a><span class="k">def</span> <span class="nf">tiny_imagenet_val</span><span class="p">(</span>
    <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">dataloader_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;tiny_imagenet_dataset_params&quot;</span><span class="p">,</span>
<span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="n">config_name</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">ImageNetDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cifar10_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cifar10_train">[docs]</a><span class="k">def</span> <span class="nf">cifar10_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cifar10_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">Cifar10</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cifar10_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cifar10_val">[docs]</a><span class="k">def</span> <span class="nf">cifar10_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cifar10_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">Cifar10</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cifar100_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cifar100_train">[docs]</a><span class="k">def</span> <span class="nf">cifar100_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cifar100_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">Cifar100</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cifar100_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cifar100_val">[docs]</a><span class="k">def</span> <span class="nf">cifar100_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cifar100_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">Cifar100</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<span class="k">def</span> <span class="nf">classification_test_dataloader</span><span class="p">(</span><span class="n">batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">5</span><span class="p">,</span> <span class="n">image_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">32</span><span class="p">,</span> <span class="n">dataset_size</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">DataLoader</span><span class="p">:</span>
    <span class="n">dataset_size</span> <span class="o">=</span> <span class="n">dataset_size</span> <span class="ow">or</span> <span class="n">batch_size</span>
    <span class="n">images</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">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">dataset_size</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">)))</span>
    <span class="n">ground_truth</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">LongTensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">dataset_size</span><span class="p">)))</span>
    <span class="n">dataset</span> <span class="o">=</span> <span class="n">TensorDataset</span><span class="p">(</span><span class="n">images</span><span class="p">,</span> <span class="n">ground_truth</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="o">=</span><span class="n">dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="n">batch_size</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">detection_test_dataloader</span><span class="p">(</span><span class="n">batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">5</span><span class="p">,</span> <span class="n">image_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">320</span><span class="p">,</span> <span class="n">dataset_size</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">DataLoader</span><span class="p">:</span>
    <span class="n">dataset_size</span> <span class="o">=</span> <span class="n">dataset_size</span> <span class="ow">or</span> <span class="n">batch_size</span>
    <span class="n">images</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">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">dataset_size</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">)))</span>
    <span class="n">ground_truth</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">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">dataset_size</span><span class="p">,</span> <span class="mi">6</span><span class="p">)))</span>
    <span class="n">dataset</span> <span class="o">=</span> <span class="n">TensorDataset</span><span class="p">(</span><span class="n">images</span><span class="p">,</span> <span class="n">ground_truth</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="o">=</span><span class="n">dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="n">batch_size</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">segmentation_test_dataloader</span><span class="p">(</span><span class="n">batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">5</span><span class="p">,</span> <span class="n">image_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">512</span><span class="p">,</span> <span class="n">dataset_size</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">DataLoader</span><span class="p">:</span>
    <span class="n">dataset_size</span> <span class="o">=</span> <span class="n">dataset_size</span> <span class="ow">or</span> <span class="n">batch_size</span>
    <span class="n">images</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">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">dataset_size</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">)))</span>
    <span class="n">ground_truth</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">LongTensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">dataset_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">)))</span>
    <span class="n">dataset</span> <span class="o">=</span> <span class="n">TensorDataset</span><span class="p">(</span><span class="n">images</span><span class="p">,</span> <span class="n">ground_truth</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="o">=</span><span class="n">dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="n">batch_size</span><span class="p">)</span>


<div class="viewcode-block" id="cityscapes_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_train">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_val">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_stdc_seg50_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_stdc_seg50_train">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_stdc_seg50_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_stdc_seg50_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_stdc_seg50_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_stdc_seg50_val">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_stdc_seg50_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_stdc_seg50_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_stdc_seg75_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_stdc_seg75_train">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_stdc_seg75_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_stdc_seg75_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_stdc_seg75_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_stdc_seg75_val">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_stdc_seg75_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_stdc_seg75_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_regseg48_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_regseg48_train">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_regseg48_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_regseg48_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_regseg48_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_regseg48_val">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_regseg48_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_regseg48_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_ddrnet_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_ddrnet_train">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_ddrnet_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_ddrnet_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_ddrnet_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.cityscapes_ddrnet_val">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_ddrnet_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;cityscapes_ddrnet_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CityscapesDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="coco_segmentation_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco_segmentation_train">[docs]</a><span class="k">def</span> <span class="nf">coco_segmentation_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;coco_segmentation_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CoCoSegmentationDataSet</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="coco_segmentation_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.coco_segmentation_val">[docs]</a><span class="k">def</span> <span class="nf">coco_segmentation_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;coco_segmentation_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">CoCoSegmentationDataSet</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="pascal_aug_segmentation_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.pascal_aug_segmentation_train">[docs]</a><span class="k">def</span> <span class="nf">pascal_aug_segmentation_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;pascal_aug_segmentation_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">PascalVOCAndAUGUnifiedDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="pascal_aug_segmentation_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.pascal_aug_segmentation_val">[docs]</a><span class="k">def</span> <span class="nf">pascal_aug_segmentation_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">pascal_voc_segmentation_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="pascal_voc_segmentation_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.pascal_voc_segmentation_train">[docs]</a><span class="k">def</span> <span class="nf">pascal_voc_segmentation_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;pascal_voc_segmentation_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="pascal_voc_segmentation_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.pascal_voc_segmentation_val">[docs]</a><span class="k">def</span> <span class="nf">pascal_voc_segmentation_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;pascal_voc_segmentation_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="supervisely_persons_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.supervisely_persons_train">[docs]</a><span class="k">def</span> <span class="nf">supervisely_persons_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;supervisely_persons_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">SuperviselyPersonsDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="supervisely_persons_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.supervisely_persons_val">[docs]</a><span class="k">def</span> <span class="nf">supervisely_persons_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;supervisely_persons_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">SuperviselyPersonsDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="pascal_voc_detection_train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.pascal_voc_detection_train">[docs]</a><span class="k">def</span> <span class="nf">pascal_voc_detection_train</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;pascal_voc_detection_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">PascalVOCUnifiedDetectionTrainDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="pascal_voc_detection_val"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.pascal_voc_detection_val">[docs]</a><span class="k">def</span> <span class="nf">pascal_voc_detection_val</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get_data_loader</span><span class="p">(</span>
        <span class="n">config_name</span><span class="o">=</span><span class="s2">&quot;pascal_voc_detection_dataset_params&quot;</span><span class="p">,</span>
        <span class="n">dataset_cls</span><span class="o">=</span><span class="n">PascalVOCDetectionDataset</span><span class="p">,</span>
        <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
        <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">,</span>
    <span class="p">)</span></div>


<span class="n">ALL_DATALOADERS</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;coco2017_train&quot;</span><span class="p">:</span> <span class="n">coco2017_train</span><span class="p">,</span>
    <span class="s2">&quot;coco2017_val&quot;</span><span class="p">:</span> <span class="n">coco2017_val</span><span class="p">,</span>
    <span class="s2">&quot;coco2017_train_yolox&quot;</span><span class="p">:</span> <span class="n">coco2017_train_yolox</span><span class="p">,</span>
    <span class="s2">&quot;coco2017_val_yolox&quot;</span><span class="p">:</span> <span class="n">coco2017_val_yolox</span><span class="p">,</span>
    <span class="s2">&quot;coco2017_train_ssd_lite_mobilenet_v2&quot;</span><span class="p">:</span> <span class="n">coco2017_train_ssd_lite_mobilenet_v2</span><span class="p">,</span>
    <span class="s2">&quot;coco2017_val_ssd_lite_mobilenet_v2&quot;</span><span class="p">:</span> <span class="n">coco2017_val_ssd_lite_mobilenet_v2</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_train&quot;</span><span class="p">:</span> <span class="n">imagenet_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_val&quot;</span><span class="p">:</span> <span class="n">imagenet_val</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_efficientnet_train&quot;</span><span class="p">:</span> <span class="n">imagenet_efficientnet_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_efficientnet_val&quot;</span><span class="p">:</span> <span class="n">imagenet_efficientnet_val</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_mobilenetv2_train&quot;</span><span class="p">:</span> <span class="n">imagenet_mobilenetv2_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_mobilenetv2_val&quot;</span><span class="p">:</span> <span class="n">imagenet_mobilenetv2_val</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_mobilenetv3_train&quot;</span><span class="p">:</span> <span class="n">imagenet_mobilenetv3_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_mobilenetv3_val&quot;</span><span class="p">:</span> <span class="n">imagenet_mobilenetv3_val</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_regnetY_train&quot;</span><span class="p">:</span> <span class="n">imagenet_regnetY_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_regnetY_val&quot;</span><span class="p">:</span> <span class="n">imagenet_regnetY_val</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_resnet50_train&quot;</span><span class="p">:</span> <span class="n">imagenet_resnet50_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_resnet50_val&quot;</span><span class="p">:</span> <span class="n">imagenet_resnet50_val</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_resnet50_kd_train&quot;</span><span class="p">:</span> <span class="n">imagenet_resnet50_kd_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_resnet50_kd_val&quot;</span><span class="p">:</span> <span class="n">imagenet_resnet50_kd_val</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_vit_base_train&quot;</span><span class="p">:</span> <span class="n">imagenet_vit_base_train</span><span class="p">,</span>
    <span class="s2">&quot;imagenet_vit_base_val&quot;</span><span class="p">:</span> <span class="n">imagenet_vit_base_val</span><span class="p">,</span>
    <span class="s2">&quot;tiny_imagenet_train&quot;</span><span class="p">:</span> <span class="n">tiny_imagenet_train</span><span class="p">,</span>
    <span class="s2">&quot;tiny_imagenet_val&quot;</span><span class="p">:</span> <span class="n">tiny_imagenet_val</span><span class="p">,</span>
    <span class="s2">&quot;cifar10_train&quot;</span><span class="p">:</span> <span class="n">cifar10_train</span><span class="p">,</span>
    <span class="s2">&quot;cifar10_val&quot;</span><span class="p">:</span> <span class="n">cifar10_val</span><span class="p">,</span>
    <span class="s2">&quot;cifar100_train&quot;</span><span class="p">:</span> <span class="n">cifar100_train</span><span class="p">,</span>
    <span class="s2">&quot;cifar100_val&quot;</span><span class="p">:</span> <span class="n">cifar100_val</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_train&quot;</span><span class="p">:</span> <span class="n">cityscapes_train</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_val&quot;</span><span class="p">:</span> <span class="n">cityscapes_val</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_stdc_seg50_train&quot;</span><span class="p">:</span> <span class="n">cityscapes_stdc_seg50_train</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_stdc_seg50_val&quot;</span><span class="p">:</span> <span class="n">cityscapes_stdc_seg50_val</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_stdc_seg75_train&quot;</span><span class="p">:</span> <span class="n">cityscapes_stdc_seg75_train</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_stdc_seg75_val&quot;</span><span class="p">:</span> <span class="n">cityscapes_stdc_seg75_val</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_regseg48_train&quot;</span><span class="p">:</span> <span class="n">cityscapes_regseg48_train</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_regseg48_val&quot;</span><span class="p">:</span> <span class="n">cityscapes_regseg48_val</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_ddrnet_train&quot;</span><span class="p">:</span> <span class="n">cityscapes_ddrnet_train</span><span class="p">,</span>
    <span class="s2">&quot;cityscapes_ddrnet_val&quot;</span><span class="p">:</span> <span class="n">cityscapes_ddrnet_val</span><span class="p">,</span>
    <span class="s2">&quot;coco_segmentation_train&quot;</span><span class="p">:</span> <span class="n">coco_segmentation_train</span><span class="p">,</span>
    <span class="s2">&quot;coco_segmentation_val&quot;</span><span class="p">:</span> <span class="n">coco_segmentation_val</span><span class="p">,</span>
    <span class="s2">&quot;pascal_aug_segmentation_train&quot;</span><span class="p">:</span> <span class="n">pascal_aug_segmentation_train</span><span class="p">,</span>
    <span class="s2">&quot;pascal_aug_segmentation_val&quot;</span><span class="p">:</span> <span class="n">pascal_aug_segmentation_val</span><span class="p">,</span>
    <span class="s2">&quot;pascal_voc_segmentation_train&quot;</span><span class="p">:</span> <span class="n">pascal_voc_segmentation_train</span><span class="p">,</span>
    <span class="s2">&quot;pascal_voc_segmentation_val&quot;</span><span class="p">:</span> <span class="n">pascal_voc_segmentation_val</span><span class="p">,</span>
    <span class="s2">&quot;supervisely_persons_train&quot;</span><span class="p">:</span> <span class="n">supervisely_persons_train</span><span class="p">,</span>
    <span class="s2">&quot;supervisely_persons_val&quot;</span><span class="p">:</span> <span class="n">supervisely_persons_val</span><span class="p">,</span>
    <span class="s2">&quot;pascal_voc_detection_train&quot;</span><span class="p">:</span> <span class="n">pascal_voc_detection_train</span><span class="p">,</span>
    <span class="s2">&quot;pascal_voc_detection_val&quot;</span><span class="p">:</span> <span class="n">pascal_voc_detection_val</span><span class="p">,</span>
<span class="p">}</span>


<div class="viewcode-block" id="get"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.dataloaders.get">[docs]</a><span class="k">def</span> <span class="nf">get</span><span class="p">(</span><span class="n">name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataset_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dataset</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">Dataset</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">DataLoader</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Get DataLoader of the recipe-configured dataset defined by name in ALL_DATALOADERS.</span>

<span class="sd">    :param name: dataset name in ALL_DATALOADERS.</span>
<span class="sd">    :param dataset_params: dataset params that override the yaml configured defaults, then passed to the dataset_cls.__init__.</span>
<span class="sd">    :param dataloader_params: DataLoader params that override the yaml configured defaults, then passed to the DataLoader.__init__</span>
<span class="sd">    :param dataset: torch.utils.data.Dataset to be used instead of passing &quot;name&quot; (i.e for external dataset objects).</span>
<span class="sd">    :return: initialized DataLoader.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">dataset</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">name</span> <span class="ow">or</span> <span class="n">dataset_params</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;&#39;name&#39; and &#39;dataset_params&#39; cannot be passed with initialized dataset.&quot;</span><span class="p">)</span>
        <span class="n">dataloader_params</span> <span class="o">=</span> <span class="n">_process_sampler_params</span><span class="p">(</span><span class="n">dataloader_params</span><span class="p">,</span> <span class="n">dataset</span><span class="p">,</span> <span class="p">{})</span>
        <span class="n">dataloader</span> <span class="o">=</span> <span class="n">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="o">=</span><span class="n">dataset</span><span class="p">,</span> <span class="o">**</span><span class="n">dataloader_params</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">name</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">ALL_DATALOADERS</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Unsupported dataloader: &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">name</span><span class="p">))</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">dataloader_cls</span> <span class="o">=</span> <span class="n">ALL_DATALOADERS</span><span class="p">[</span><span class="n">name</span><span class="p">]</span>
        <span class="n">dataloader</span> <span class="o">=</span> <span class="n">dataloader_cls</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="n">dataloader_params</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">dataloader</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.all_datasets &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.all_datasets</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.all_datasets</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">collections</span> <span class="kn">import</span> <span class="n">defaultdict</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Type</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.dataset_interfaces</span> <span class="kn">import</span> <span class="n">DatasetInterface</span><span class="p">,</span> <span class="n">TestDatasetInterface</span><span class="p">,</span> \
    <span class="n">LibraryDatasetInterface</span><span class="p">,</span> \
    <span class="n">ClassificationDatasetInterface</span><span class="p">,</span> <span class="n">Cifar10DatasetInterface</span><span class="p">,</span> <span class="n">Cifar100DatasetInterface</span><span class="p">,</span> \
    <span class="n">ImageNetDatasetInterface</span><span class="p">,</span> <span class="n">TinyImageNetDatasetInterface</span><span class="p">,</span> <span class="n">CoCoSegmentationDatasetInterface</span><span class="p">,</span>\
    <span class="n">PascalAUG2012SegmentationDataSetInterface</span><span class="p">,</span> <span class="n">PascalVOC2012SegmentationDataSetInterface</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.data_types.enum.deep_learning_task</span> <span class="kn">import</span> <span class="n">DeepLearningTask</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.dataset_interfaces.dataset_interface</span> <span class="kn">import</span> <span class="n">CoCoDetectionDatasetInterface</span>

<span class="n">CLASSIFICATION_DATASETS</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;test_dataset&quot;</span><span class="p">:</span> <span class="n">TestDatasetInterface</span><span class="p">,</span>
    <span class="s2">&quot;library_dataset&quot;</span><span class="p">:</span> <span class="n">LibraryDatasetInterface</span><span class="p">,</span>
    <span class="s2">&quot;classification_dataset&quot;</span><span class="p">:</span> <span class="n">ClassificationDatasetInterface</span><span class="p">,</span>
    <span class="s2">&quot;cifar_10&quot;</span><span class="p">:</span> <span class="n">Cifar10DatasetInterface</span><span class="p">,</span>
    <span class="s2">&quot;cifar_100&quot;</span><span class="p">:</span> <span class="n">Cifar100DatasetInterface</span><span class="p">,</span>
    <span class="s2">&quot;imagenet&quot;</span><span class="p">:</span> <span class="n">ImageNetDatasetInterface</span><span class="p">,</span>
    <span class="s2">&quot;tiny_imagenet&quot;</span><span class="p">:</span> <span class="n">TinyImageNetDatasetInterface</span>
<span class="p">}</span>

<span class="n">OBJECT_DETECTION_DATASETS</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;coco&quot;</span><span class="p">:</span> <span class="n">CoCoDetectionDatasetInterface</span><span class="p">,</span>
<span class="p">}</span>

<span class="n">SEMANTIC_SEGMENTATION_DATASETS</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;coco&quot;</span><span class="p">:</span> <span class="n">CoCoSegmentationDatasetInterface</span><span class="p">,</span>
    <span class="s2">&quot;pascal_voc&quot;</span><span class="p">:</span> <span class="n">PascalVOC2012SegmentationDataSetInterface</span><span class="p">,</span>
    <span class="s2">&quot;pascal_aug&quot;</span><span class="p">:</span> <span class="n">PascalAUG2012SegmentationDataSetInterface</span>
<span class="p">}</span>


<div class="viewcode-block" id="DataSetDoesNotExistException"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.all_datasets.DataSetDoesNotExistException">[docs]</a><span class="k">class</span> <span class="nc">DataSetDoesNotExistException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The requested dataset does not exist, or is not implemented.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">pass</span></div>


<div class="viewcode-block" id="SgLibraryDatasets"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.all_datasets.SgLibraryDatasets">[docs]</a><span class="k">class</span> <span class="nc">SgLibraryDatasets</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Holds all of the different library dataset dictionaries, by DL Task mapping</span>

<span class="sd">        Attributes:</span>
<span class="sd">            CLASSIFICATION          Dictionary of Classification Data sets</span>
<span class="sd">            OBJECT_DETECTION        Dictionary of Object Detection Data sets</span>
<span class="sd">            SEMANTIC_SEGMENTATION   Dictionary of Semantic Segmentation Data sets</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">CLASSIFICATION</span> <span class="o">=</span> <span class="n">CLASSIFICATION_DATASETS</span>
    <span class="n">OBJECT_DETECTION</span> <span class="o">=</span> <span class="n">OBJECT_DETECTION_DATASETS</span>
    <span class="n">SEMANTIC_SEGMENTATION</span> <span class="o">=</span> <span class="n">SEMANTIC_SEGMENTATION_DATASETS</span>

    <span class="n">_datasets_mapping</span> <span class="o">=</span> <span class="p">{</span>
        <span class="n">DeepLearningTask</span><span class="o">.</span><span class="n">CLASSIFICATION</span><span class="p">:</span> <span class="n">CLASSIFICATION</span><span class="p">,</span>
        <span class="n">DeepLearningTask</span><span class="o">.</span><span class="n">SEMANTIC_SEGMENTATION</span><span class="p">:</span> <span class="n">SEMANTIC_SEGMENTATION</span><span class="p">,</span>
        <span class="n">DeepLearningTask</span><span class="o">.</span><span class="n">OBJECT_DETECTION</span><span class="p">:</span> <span class="n">OBJECT_DETECTION</span><span class="p">,</span>
    <span class="p">}</span>

<div class="viewcode-block" id="SgLibraryDatasets.get_all_available_datasets"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.all_datasets.SgLibraryDatasets.get_all_available_datasets">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">get_all_available_datasets</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Gets all the available datasets.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">all_datasets</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">list</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">dl_task</span><span class="p">,</span> <span class="n">task_datasets</span> <span class="ow">in</span> <span class="n">SgLibraryDatasets</span><span class="o">.</span><span class="n">_datasets_mapping</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="k">for</span> <span class="n">dataset_name</span><span class="p">,</span> <span class="n">dataset_interface</span> <span class="ow">in</span> <span class="n">task_datasets</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                <span class="n">all_datasets</span><span class="p">[</span><span class="n">dl_task</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dataset_name</span><span class="p">)</span>

        <span class="c1"># TODO: Return Dataset Metadata list from the dataset interfaces objects</span>
        <span class="c1"># TODO: Transform DatasetInterface -&gt; DataSetMetadata</span>
        <span class="k">return</span> <span class="n">all_datasets</span></div>

<div class="viewcode-block" id="SgLibraryDatasets.get_dataset"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.all_datasets.SgLibraryDatasets.get_dataset">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">get_dataset</span><span class="p">(</span><span class="n">dl_task</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">dataset_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Type</span><span class="p">[</span><span class="n">DatasetInterface</span><span class="p">]:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get&#39;s a dataset with a given name for a given deep learning task.</span>
<span class="sd">        examp:</span>
<span class="sd">        &gt;&gt;&gt; SgLibraryDatasets.get_dataset(dl_task=&#39;classification&#39;, dataset_name=&#39;cifar_100&#39;)</span>
<span class="sd">        &gt;&gt;&gt; &lt;Cifar100DatasetInterface instance&gt;</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">task_datasets</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">DatasetInterface</span><span class="p">]</span> <span class="o">=</span> <span class="n">SgLibraryDatasets</span><span class="o">.</span><span class="n">_datasets_mapping</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">dl_task</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">task_datasets</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Invalid Deep Learining Task: </span><span class="si">{</span><span class="n">dl_task</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>

        <span class="n">dataset</span><span class="p">:</span> <span class="n">DatasetInterface</span> <span class="o">=</span> <span class="n">task_datasets</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">dataset_name</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">dataset</span><span class="p">:</span>
            <span class="k">raise</span> <span class="n">DataSetDoesNotExistException</span><span class="p">(</span><span class="n">dataset_name</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">dataset</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.auto_augment &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.auto_augment</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.auto_augment</h1><div class="highlight"><pre>
<span></span><span class="sd">&quot;&quot;&quot; RandAugment</span>
<span class="sd">RandAugment is a variant of AutoAugment which randomly selects transformations</span>
<span class="sd"> from AutoAugment to be applied on an image.</span>

<span class="sd">RandomAugmentation Implementation adapted from:</span>
<span class="sd">    https://github.com/rwightman/pytorch-image-models/blob/master/timm/data/auto_augment.py</span>

<span class="sd">Papers:</span>
<span class="sd">    RandAugment: Practical automated data augmentation... - https://arxiv.org/abs/1909.13719</span>

<span class="sd">&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">random</span>
<span class="kn">import</span> <span class="nn">re</span>
<span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span><span class="p">,</span> <span class="n">ImageOps</span><span class="p">,</span> <span class="n">ImageEnhance</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>


<span class="n">_FILL</span> <span class="o">=</span> <span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">)</span>

<span class="c1"># to unify the calls of the different augmentations in terms of params, all augmentations are set to work with a single</span>
<span class="c1"># magnitude params, normalized according to _MAX_MAGNITUDE</span>
<span class="n">_MAX_MAGNITUDE</span> <span class="o">=</span> <span class="mf">10.</span>

<span class="n">_HPARAMS_DEFAULT</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span>
    <span class="n">translate_const</span><span class="o">=</span><span class="mi">250</span><span class="p">,</span>
    <span class="n">img_mean</span><span class="o">=</span><span class="n">_FILL</span><span class="p">,</span>
<span class="p">)</span>

<span class="c1"># Define the interpolation types</span>
<span class="n">_RANDOM_INTERPOLATION</span> <span class="o">=</span> <span class="n">Image</span><span class="o">.</span><span class="n">BILINEAR</span>


<span class="k">def</span> <span class="nf">_interpolation</span><span class="p">(</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Performs Bi-Linear interpolation</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">interpolation</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s1">&#39;resample&#39;</span><span class="p">,</span> <span class="n">Image</span><span class="o">.</span><span class="n">BILINEAR</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">interpolation</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
        <span class="k">return</span> <span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="n">interpolation</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">interpolation</span>


<div class="viewcode-block" id="shear_x"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.shear_x">[docs]</a><span class="k">def</span> <span class="nf">shear_x</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">,</span> <span class="n">Image</span><span class="o">.</span><span class="n">AFFINE</span><span class="p">,</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="shear_y"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.shear_y">[docs]</a><span class="k">def</span> <span class="nf">shear_y</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">,</span> <span class="n">Image</span><span class="o">.</span><span class="n">AFFINE</span><span class="p">,</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="translate_x_rel"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.translate_x_rel">[docs]</a><span class="k">def</span> <span class="nf">translate_x_rel</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">pct</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="n">pixels</span> <span class="o">=</span> <span class="n">pct</span> <span class="o">*</span> <span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">,</span> <span class="n">Image</span><span class="o">.</span><span class="n">AFFINE</span><span class="p">,</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">pixels</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="translate_y_rel"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.translate_y_rel">[docs]</a><span class="k">def</span> <span class="nf">translate_y_rel</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">pct</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="n">pixels</span> <span class="o">=</span> <span class="n">pct</span> <span class="o">*</span> <span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">,</span> <span class="n">Image</span><span class="o">.</span><span class="n">AFFINE</span><span class="p">,</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">pixels</span><span class="p">),</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="translate_x_abs"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.translate_x_abs">[docs]</a><span class="k">def</span> <span class="nf">translate_x_abs</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">pixels</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">,</span> <span class="n">Image</span><span class="o">.</span><span class="n">AFFINE</span><span class="p">,</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">pixels</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="translate_y_abs"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.translate_y_abs">[docs]</a><span class="k">def</span> <span class="nf">translate_y_abs</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">pixels</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">transform</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">size</span><span class="p">,</span> <span class="n">Image</span><span class="o">.</span><span class="n">AFFINE</span><span class="p">,</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">pixels</span><span class="p">),</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="rotate"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.rotate">[docs]</a><span class="k">def</span> <span class="nf">rotate</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">degrees</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">rotate</span><span class="p">(</span><span class="n">degrees</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>


<div class="viewcode-block" id="auto_contrast"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.auto_contrast">[docs]</a><span class="k">def</span> <span class="nf">auto_contrast</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageOps</span><span class="o">.</span><span class="n">autocontrast</span><span class="p">(</span><span class="n">img</span><span class="p">)</span></div>


<div class="viewcode-block" id="invert"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.invert">[docs]</a><span class="k">def</span> <span class="nf">invert</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageOps</span><span class="o">.</span><span class="n">invert</span><span class="p">(</span><span class="n">img</span><span class="p">)</span></div>


<div class="viewcode-block" id="equalize"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.equalize">[docs]</a><span class="k">def</span> <span class="nf">equalize</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageOps</span><span class="o">.</span><span class="n">equalize</span><span class="p">(</span><span class="n">img</span><span class="p">)</span></div>


<div class="viewcode-block" id="solarize"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.solarize">[docs]</a><span class="k">def</span> <span class="nf">solarize</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">thresh</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageOps</span><span class="o">.</span><span class="n">solarize</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">thresh</span><span class="p">)</span></div>


<div class="viewcode-block" id="solarize_add"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.solarize_add">[docs]</a><span class="k">def</span> <span class="nf">solarize_add</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">add</span><span class="p">,</span> <span class="n">thresh</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="n">lut</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">256</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">i</span> <span class="o">&lt;</span> <span class="n">thresh</span><span class="p">:</span>
            <span class="n">lut</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">min</span><span class="p">(</span><span class="mi">255</span><span class="p">,</span> <span class="n">i</span> <span class="o">+</span> <span class="n">add</span><span class="p">))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">lut</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">i</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">img</span><span class="o">.</span><span class="n">mode</span> <span class="ow">in</span> <span class="p">(</span><span class="s2">&quot;L&quot;</span><span class="p">,</span> <span class="s2">&quot;RGB&quot;</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">img</span><span class="o">.</span><span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;RGB&quot;</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">lut</span><span class="p">)</span> <span class="o">==</span> <span class="mi">256</span><span class="p">:</span>
            <span class="n">lut</span> <span class="o">=</span> <span class="n">lut</span> <span class="o">+</span> <span class="n">lut</span> <span class="o">+</span> <span class="n">lut</span>
        <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">point</span><span class="p">(</span><span class="n">lut</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">img</span></div>


<div class="viewcode-block" id="posterize"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.posterize">[docs]</a><span class="k">def</span> <span class="nf">posterize</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">bits_to_keep</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">if</span> <span class="n">bits_to_keep</span> <span class="o">&gt;=</span> <span class="mi">8</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">img</span>
    <span class="k">return</span> <span class="n">ImageOps</span><span class="o">.</span><span class="n">posterize</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">bits_to_keep</span><span class="p">)</span></div>


<div class="viewcode-block" id="contrast"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.contrast">[docs]</a><span class="k">def</span> <span class="nf">contrast</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageEnhance</span><span class="o">.</span><span class="n">Contrast</span><span class="p">(</span><span class="n">img</span><span class="p">)</span><span class="o">.</span><span class="n">enhance</span><span class="p">(</span><span class="n">factor</span><span class="p">)</span></div>


<div class="viewcode-block" id="color"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.color">[docs]</a><span class="k">def</span> <span class="nf">color</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageEnhance</span><span class="o">.</span><span class="n">Color</span><span class="p">(</span><span class="n">img</span><span class="p">)</span><span class="o">.</span><span class="n">enhance</span><span class="p">(</span><span class="n">factor</span><span class="p">)</span></div>


<div class="viewcode-block" id="brightness"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.brightness">[docs]</a><span class="k">def</span> <span class="nf">brightness</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageEnhance</span><span class="o">.</span><span class="n">Brightness</span><span class="p">(</span><span class="n">img</span><span class="p">)</span><span class="o">.</span><span class="n">enhance</span><span class="p">(</span><span class="n">factor</span><span class="p">)</span></div>


<div class="viewcode-block" id="sharpness"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.sharpness">[docs]</a><span class="k">def</span> <span class="nf">sharpness</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">factor</span><span class="p">,</span> <span class="o">**</span><span class="n">__</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">ImageEnhance</span><span class="o">.</span><span class="n">Sharpness</span><span class="p">(</span><span class="n">img</span><span class="p">)</span><span class="o">.</span><span class="n">enhance</span><span class="p">(</span><span class="n">factor</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_randomly_negate</span><span class="p">(</span><span class="n">v</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;With 50% prob, negate the value&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="o">-</span><span class="n">v</span> <span class="k">if</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&gt;</span> <span class="mf">0.5</span> <span class="k">else</span> <span class="n">v</span>


<span class="k">def</span> <span class="nf">_rotate_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># range [-30, 30]</span>
    <span class="n">level</span> <span class="o">=</span> <span class="p">(</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mf">30.</span>
    <span class="n">level</span> <span class="o">=</span> <span class="n">_randomly_negate</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">level</span><span class="p">,</span>


<span class="k">def</span> <span class="nf">_enhance_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># range [0.1, 1.9]</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mf">1.8</span> <span class="o">+</span> <span class="mf">0.1</span><span class="p">,</span>


<span class="k">def</span> <span class="nf">_enhance_increasing_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># range [0.1, 1.9]</span>
    <span class="n">level</span> <span class="o">=</span> <span class="p">(</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mf">.9</span>
    <span class="n">level</span> <span class="o">=</span> <span class="mf">1.0</span> <span class="o">+</span> <span class="n">_randomly_negate</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">level</span><span class="p">,</span>


<span class="k">def</span> <span class="nf">_shear_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># range [-0.3, 0.3]</span>
    <span class="n">level</span> <span class="o">=</span> <span class="p">(</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.3</span>
    <span class="n">level</span> <span class="o">=</span> <span class="n">_randomly_negate</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">level</span><span class="p">,</span>


<span class="k">def</span> <span class="nf">_translate_abs_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">hparams</span><span class="p">):</span>
    <span class="n">translate_const</span> <span class="o">=</span> <span class="n">hparams</span><span class="p">[</span><span class="s1">&#39;translate_const&#39;</span><span class="p">]</span>
    <span class="n">level</span> <span class="o">=</span> <span class="p">(</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="nb">float</span><span class="p">(</span><span class="n">translate_const</span><span class="p">)</span>
    <span class="n">level</span> <span class="o">=</span> <span class="n">_randomly_negate</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">level</span><span class="p">,</span>


<span class="k">def</span> <span class="nf">_translate_rel_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">hparams</span><span class="p">):</span>
    <span class="c1"># default range [-0.45, 0.45]</span>
    <span class="n">translate_pct</span> <span class="o">=</span> <span class="n">hparams</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;translate_pct&#39;</span><span class="p">,</span> <span class="mf">0.45</span><span class="p">)</span>
    <span class="n">level</span> <span class="o">=</span> <span class="p">(</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="n">translate_pct</span>
    <span class="n">level</span> <span class="o">=</span> <span class="n">_randomly_negate</span><span class="p">(</span><span class="n">level</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">level</span><span class="p">,</span>


<span class="k">def</span> <span class="nf">_posterize_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># As per Tensorflow TPU EfficientNet impl</span>
    <span class="c1"># range [0, 4], &#39;keep 0 up to 4 MSB of original image&#39;</span>
    <span class="c1"># intensity/severity of augmentation decreases with level</span>
    <span class="k">return</span> <span class="nb">int</span><span class="p">((</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mi">4</span><span class="p">),</span>


<span class="k">def</span> <span class="nf">_posterize_increasing_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">hparams</span><span class="p">):</span>
    <span class="c1"># As per Tensorflow models research and UDA impl</span>
    <span class="c1"># range [4, 0], &#39;keep 4 down to 0 MSB of original image&#39;,</span>
    <span class="c1"># intensity/severity of augmentation increases with level</span>
    <span class="k">return</span> <span class="mi">4</span> <span class="o">-</span> <span class="n">_posterize_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">hparams</span><span class="p">)[</span><span class="mi">0</span><span class="p">],</span>


<span class="k">def</span> <span class="nf">_posterize_original_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># As per original AutoAugment paper description</span>
    <span class="c1"># range [4, 8], &#39;keep 4 up to 8 MSB of image&#39;</span>
    <span class="c1"># intensity/severity of augmentation decreases with level</span>
    <span class="k">return</span> <span class="nb">int</span><span class="p">((</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mi">4</span><span class="p">)</span> <span class="o">+</span> <span class="mi">4</span><span class="p">,</span>


<span class="k">def</span> <span class="nf">_solarize_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># range [0, 256]</span>
    <span class="c1"># intensity/severity of augmentation decreases with level</span>
    <span class="k">return</span> <span class="nb">int</span><span class="p">((</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mi">256</span><span class="p">),</span>


<span class="k">def</span> <span class="nf">_solarize_increasing_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># range [0, 256]</span>
    <span class="c1"># intensity/severity of augmentation increases with level</span>
    <span class="k">return</span> <span class="mi">256</span> <span class="o">-</span> <span class="n">_solarize_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">)[</span><span class="mi">0</span><span class="p">],</span>


<span class="k">def</span> <span class="nf">_solarize_add_level_to_arg</span><span class="p">(</span><span class="n">level</span><span class="p">,</span> <span class="n">_hparams</span><span class="p">):</span>
    <span class="c1"># range [0, 110]</span>
    <span class="k">return</span> <span class="nb">int</span><span class="p">((</span><span class="n">level</span> <span class="o">/</span> <span class="n">_MAX_MAGNITUDE</span><span class="p">)</span> <span class="o">*</span> <span class="mi">110</span><span class="p">),</span>


<span class="n">LEVEL_TO_ARG</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s1">&#39;AutoContrast&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
    <span class="s1">&#39;Equalize&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
    <span class="s1">&#39;Invert&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
    <span class="s1">&#39;Rotate&#39;</span><span class="p">:</span> <span class="n">_rotate_level_to_arg</span><span class="p">,</span>
    <span class="c1"># There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers</span>
    <span class="s1">&#39;Posterize&#39;</span><span class="p">:</span> <span class="n">_posterize_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;PosterizeIncreasing&#39;</span><span class="p">:</span> <span class="n">_posterize_increasing_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;PosterizeOriginal&#39;</span><span class="p">:</span> <span class="n">_posterize_original_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;Solarize&#39;</span><span class="p">:</span> <span class="n">_solarize_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeIncreasing&#39;</span><span class="p">:</span> <span class="n">_solarize_increasing_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeAdd&#39;</span><span class="p">:</span> <span class="n">_solarize_add_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;Color&#39;</span><span class="p">:</span> <span class="n">_enhance_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;ColorIncreasing&#39;</span><span class="p">:</span> <span class="n">_enhance_increasing_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;Contrast&#39;</span><span class="p">:</span> <span class="n">_enhance_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;ContrastIncreasing&#39;</span><span class="p">:</span> <span class="n">_enhance_increasing_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;Brightness&#39;</span><span class="p">:</span> <span class="n">_enhance_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;BrightnessIncreasing&#39;</span><span class="p">:</span> <span class="n">_enhance_increasing_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;Sharpness&#39;</span><span class="p">:</span> <span class="n">_enhance_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;SharpnessIncreasing&#39;</span><span class="p">:</span> <span class="n">_enhance_increasing_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;ShearX&#39;</span><span class="p">:</span> <span class="n">_shear_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;ShearY&#39;</span><span class="p">:</span> <span class="n">_shear_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;TranslateX&#39;</span><span class="p">:</span> <span class="n">_translate_abs_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;TranslateY&#39;</span><span class="p">:</span> <span class="n">_translate_abs_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;TranslateXRel&#39;</span><span class="p">:</span> <span class="n">_translate_rel_level_to_arg</span><span class="p">,</span>
    <span class="s1">&#39;TranslateYRel&#39;</span><span class="p">:</span> <span class="n">_translate_rel_level_to_arg</span><span class="p">,</span>
<span class="p">}</span>


<span class="n">NAME_TO_OP</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s1">&#39;AutoContrast&#39;</span><span class="p">:</span> <span class="n">auto_contrast</span><span class="p">,</span>
    <span class="s1">&#39;Equalize&#39;</span><span class="p">:</span> <span class="n">equalize</span><span class="p">,</span>
    <span class="s1">&#39;Invert&#39;</span><span class="p">:</span> <span class="n">invert</span><span class="p">,</span>
    <span class="s1">&#39;Rotate&#39;</span><span class="p">:</span> <span class="n">rotate</span><span class="p">,</span>
    <span class="s1">&#39;Posterize&#39;</span><span class="p">:</span> <span class="n">posterize</span><span class="p">,</span>
    <span class="s1">&#39;PosterizeIncreasing&#39;</span><span class="p">:</span> <span class="n">posterize</span><span class="p">,</span>
    <span class="s1">&#39;PosterizeOriginal&#39;</span><span class="p">:</span> <span class="n">posterize</span><span class="p">,</span>
    <span class="s1">&#39;Solarize&#39;</span><span class="p">:</span> <span class="n">solarize</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeIncreasing&#39;</span><span class="p">:</span> <span class="n">solarize</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeAdd&#39;</span><span class="p">:</span> <span class="n">solarize_add</span><span class="p">,</span>
    <span class="s1">&#39;Color&#39;</span><span class="p">:</span> <span class="n">color</span><span class="p">,</span>
    <span class="s1">&#39;ColorIncreasing&#39;</span><span class="p">:</span> <span class="n">color</span><span class="p">,</span>
    <span class="s1">&#39;Contrast&#39;</span><span class="p">:</span> <span class="n">contrast</span><span class="p">,</span>
    <span class="s1">&#39;ContrastIncreasing&#39;</span><span class="p">:</span> <span class="n">contrast</span><span class="p">,</span>
    <span class="s1">&#39;Brightness&#39;</span><span class="p">:</span> <span class="n">brightness</span><span class="p">,</span>
    <span class="s1">&#39;BrightnessIncreasing&#39;</span><span class="p">:</span> <span class="n">brightness</span><span class="p">,</span>
    <span class="s1">&#39;Sharpness&#39;</span><span class="p">:</span> <span class="n">sharpness</span><span class="p">,</span>
    <span class="s1">&#39;SharpnessIncreasing&#39;</span><span class="p">:</span> <span class="n">sharpness</span><span class="p">,</span>
    <span class="s1">&#39;ShearX&#39;</span><span class="p">:</span> <span class="n">shear_x</span><span class="p">,</span>
    <span class="s1">&#39;ShearY&#39;</span><span class="p">:</span> <span class="n">shear_y</span><span class="p">,</span>
    <span class="s1">&#39;TranslateX&#39;</span><span class="p">:</span> <span class="n">translate_x_abs</span><span class="p">,</span>
    <span class="s1">&#39;TranslateY&#39;</span><span class="p">:</span> <span class="n">translate_y_abs</span><span class="p">,</span>
    <span class="s1">&#39;TranslateXRel&#39;</span><span class="p">:</span> <span class="n">translate_x_rel</span><span class="p">,</span>
    <span class="s1">&#39;TranslateYRel&#39;</span><span class="p">:</span> <span class="n">translate_y_rel</span><span class="p">,</span>
<span class="p">}</span>


<div class="viewcode-block" id="AugmentOp"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.AugmentOp">[docs]</a><span class="k">class</span> <span class="nc">AugmentOp</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    single auto augment operations</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">prob</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span> <span class="n">magnitude</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">hparams</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">hparams</span> <span class="o">=</span> <span class="n">hparams</span> <span class="ow">or</span> <span class="n">_HPARAMS_DEFAULT</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">aug_fn</span> <span class="o">=</span> <span class="n">NAME_TO_OP</span><span class="p">[</span><span class="n">name</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">level_fn</span> <span class="o">=</span> <span class="n">LEVEL_TO_ARG</span><span class="p">[</span><span class="n">name</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">=</span> <span class="n">prob</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">magnitude</span> <span class="o">=</span> <span class="n">magnitude</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">hparams</span> <span class="o">=</span> <span class="n">hparams</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">kwargs</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span>
            <span class="n">fillcolor</span><span class="o">=</span><span class="n">hparams</span><span class="p">[</span><span class="s1">&#39;img_mean&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;img_mean&#39;</span> <span class="ow">in</span> <span class="n">hparams</span> <span class="k">else</span> <span class="n">_FILL</span><span class="p">,</span>
            <span class="n">resample</span><span class="o">=</span><span class="n">hparams</span><span class="p">[</span><span class="s1">&#39;interpolation&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;interpolation&#39;</span> <span class="ow">in</span> <span class="n">hparams</span> <span class="k">else</span> <span class="n">_RANDOM_INTERPOLATION</span><span class="p">,</span>
        <span class="p">)</span>

        <span class="c1"># If magnitude_std is &gt; 0, introduce some randomness</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">magnitude_std</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">hparams</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;magnitude_std&#39;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">&lt;</span> <span class="mf">1.0</span> <span class="ow">and</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">img</span>
        <span class="n">magnitude</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">magnitude</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">magnitude_std</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">magnitude_std</span> <span class="o">==</span> <span class="nb">float</span><span class="p">(</span><span class="s1">&#39;inf&#39;</span><span class="p">):</span>
                <span class="n">magnitude</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">magnitude</span><span class="p">)</span>
            <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">magnitude_std</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">magnitude</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">gauss</span><span class="p">(</span><span class="n">magnitude</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">magnitude_std</span><span class="p">)</span>
        <span class="n">magnitude</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">_MAX_MAGNITUDE</span><span class="p">,</span> <span class="nb">max</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">magnitude</span><span class="p">))</span>  <span class="c1"># clip to valid range</span>
        <span class="n">level_args</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">level_fn</span><span class="p">(</span><span class="n">magnitude</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">hparams</span><span class="p">)</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">level_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="nb">tuple</span><span class="p">()</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">aug_fn</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="o">*</span><span class="n">level_args</span><span class="p">,</span> <span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">kwargs</span><span class="p">)</span></div>


<span class="n">_RAND_TRANSFORMS</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s1">&#39;AutoContrast&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Equalize&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Invert&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Rotate&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Posterize&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Solarize&#39;</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeAdd&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Color&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Contrast&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Brightness&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Sharpness&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ShearX&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ShearY&#39;</span><span class="p">,</span>
    <span class="s1">&#39;TranslateXRel&#39;</span><span class="p">,</span>
    <span class="s1">&#39;TranslateYRel&#39;</span><span class="p">,</span>
    <span class="c1"># &#39;Cutout&#39;  # NOTE I&#39;ve implement this as random erasing separately</span>
<span class="p">]</span>


<span class="n">_RAND_INCREASING_TRANSFORMS</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s1">&#39;AutoContrast&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Equalize&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Invert&#39;</span><span class="p">,</span>
    <span class="s1">&#39;Rotate&#39;</span><span class="p">,</span>
    <span class="s1">&#39;PosterizeIncreasing&#39;</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeIncreasing&#39;</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeAdd&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ColorIncreasing&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ContrastIncreasing&#39;</span><span class="p">,</span>
    <span class="s1">&#39;BrightnessIncreasing&#39;</span><span class="p">,</span>
    <span class="s1">&#39;SharpnessIncreasing&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ShearX&#39;</span><span class="p">,</span>
    <span class="s1">&#39;ShearY&#39;</span><span class="p">,</span>
    <span class="s1">&#39;TranslateXRel&#39;</span><span class="p">,</span>
    <span class="s1">&#39;TranslateYRel&#39;</span><span class="p">,</span>
    <span class="c1"># &#39;Cutout&#39;  # NOTE I&#39;ve implement this as random erasing separately</span>
<span class="p">]</span>

<span class="c1"># These experimental weights are based loosely on the relative improvements mentioned in paper.</span>
<span class="c1"># They may not result in increased performance, but could likely be tuned to so.</span>
<span class="n">_RAND_CHOICE_WEIGHTS_0</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s1">&#39;Rotate&#39;</span><span class="p">:</span> <span class="mf">0.3</span><span class="p">,</span>
    <span class="s1">&#39;ShearX&#39;</span><span class="p">:</span> <span class="mf">0.2</span><span class="p">,</span>
    <span class="s1">&#39;ShearY&#39;</span><span class="p">:</span> <span class="mf">0.2</span><span class="p">,</span>
    <span class="s1">&#39;TranslateXRel&#39;</span><span class="p">:</span> <span class="mf">0.1</span><span class="p">,</span>
    <span class="s1">&#39;TranslateYRel&#39;</span><span class="p">:</span> <span class="mf">0.1</span><span class="p">,</span>
    <span class="s1">&#39;Color&#39;</span><span class="p">:</span> <span class="mf">.025</span><span class="p">,</span>
    <span class="s1">&#39;Sharpness&#39;</span><span class="p">:</span> <span class="mf">0.025</span><span class="p">,</span>
    <span class="s1">&#39;AutoContrast&#39;</span><span class="p">:</span> <span class="mf">0.025</span><span class="p">,</span>
    <span class="s1">&#39;Solarize&#39;</span><span class="p">:</span> <span class="mf">.005</span><span class="p">,</span>
    <span class="s1">&#39;SolarizeAdd&#39;</span><span class="p">:</span> <span class="mf">.005</span><span class="p">,</span>
    <span class="s1">&#39;Contrast&#39;</span><span class="p">:</span> <span class="mf">.005</span><span class="p">,</span>
    <span class="s1">&#39;Brightness&#39;</span><span class="p">:</span> <span class="mf">.005</span><span class="p">,</span>
    <span class="s1">&#39;Equalize&#39;</span><span class="p">:</span> <span class="mf">.005</span><span class="p">,</span>
    <span class="s1">&#39;Posterize&#39;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span>
    <span class="s1">&#39;Invert&#39;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span>
<span class="p">}</span>


<span class="k">def</span> <span class="nf">_select_rand_weights</span><span class="p">(</span><span class="n">weight_idx</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">transforms</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="n">transforms</span> <span class="o">=</span> <span class="n">transforms</span> <span class="ow">or</span> <span class="n">_RAND_TRANSFORMS</span>
    <span class="k">assert</span> <span class="n">weight_idx</span> <span class="o">==</span> <span class="mi">0</span>  <span class="c1"># only one set of weights currently</span>
    <span class="n">rand_weights</span> <span class="o">=</span> <span class="n">_RAND_CHOICE_WEIGHTS_0</span>
    <span class="n">probs</span> <span class="o">=</span> <span class="p">[</span><span class="n">rand_weights</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">transforms</span><span class="p">]</span>
    <span class="n">probs</span> <span class="o">/=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">probs</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">probs</span>


<div class="viewcode-block" id="rand_augment_ops"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.rand_augment_ops">[docs]</a><span class="k">def</span> <span class="nf">rand_augment_ops</span><span class="p">(</span><span class="n">magnitude</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">hparams</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">transforms</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="n">hparams</span> <span class="o">=</span> <span class="n">hparams</span> <span class="ow">or</span> <span class="n">_HPARAMS_DEFAULT</span>
    <span class="n">transforms</span> <span class="o">=</span> <span class="n">transforms</span> <span class="ow">or</span> <span class="n">_RAND_TRANSFORMS</span>
    <span class="k">return</span> <span class="p">[</span><span class="n">AugmentOp</span><span class="p">(</span>
        <span class="n">name</span><span class="p">,</span> <span class="n">prob</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span> <span class="n">magnitude</span><span class="o">=</span><span class="n">magnitude</span><span class="p">,</span> <span class="n">hparams</span><span class="o">=</span><span class="n">hparams</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span> <span class="ow">in</span> <span class="n">transforms</span><span class="p">]</span></div>


<div class="viewcode-block" id="RandAugment"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.RandAugment">[docs]</a><span class="k">class</span> <span class="nc">RandAugment</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Random auto augment class, will select auto augment transforms according to probability weights for each op</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ops</span><span class="p">,</span> <span class="n">num_layers</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">choice_weights</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ops</span> <span class="o">=</span> <span class="n">ops</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">num_layers</span> <span class="o">=</span> <span class="n">num_layers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">choice_weights</span> <span class="o">=</span> <span class="n">choice_weights</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">):</span>
        <span class="c1"># no replacement when using weighted choice</span>
        <span class="n">ops</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">ops</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_layers</span><span class="p">,</span> <span class="n">replace</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">choice_weights</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">choice_weights</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">op</span> <span class="ow">in</span> <span class="n">ops</span><span class="p">:</span>
            <span class="n">img</span> <span class="o">=</span> <span class="n">op</span><span class="p">(</span><span class="n">img</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">img</span></div>


<div class="viewcode-block" id="rand_augment_transform"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.auto_augment.rand_augment_transform">[docs]</a><span class="k">def</span> <span class="nf">rand_augment_transform</span><span class="p">(</span><span class="n">config_str</span><span class="p">,</span> <span class="n">hparams</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Create a RandAugment transform</span>

<span class="sd">    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by</span>
<span class="sd">    dashes (&#39;-&#39;). The first section defines the specific variant of rand augment (currently only &#39;rand&#39;). The remaining</span>
<span class="sd">    sections, not order sepecific determine</span>
<span class="sd">        &#39;m&#39; - integer magnitude of rand augment</span>
<span class="sd">        &#39;n&#39; - integer num layers (number of transform ops selected per image)</span>
<span class="sd">        &#39;w&#39; - integer probabiliy weight index (index of a set of weights to influence choice of op)</span>
<span class="sd">        &#39;mstd&#39; -  float std deviation of magnitude noise applied</span>
<span class="sd">        &#39;inc&#39; - integer (bool), use augmentations that increase in severity with magnitude (default: 0)</span>
<span class="sd">    Ex &#39;rand-m9-n3-mstd0.5&#39; results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5</span>
<span class="sd">    &#39;rand-mstd1-w0&#39; results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2</span>

<span class="sd">    :param hparams: Other hparams (kwargs) for the RandAugmentation scheme</span>

<span class="sd">    :return: A PyTorch compatible Transform</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">magnitude</span> <span class="o">=</span> <span class="n">_MAX_MAGNITUDE</span>  <span class="c1"># default to _MAX_MAGNITUDE for magnitude (currently 10)</span>
    <span class="n">num_layers</span> <span class="o">=</span> <span class="mi">2</span>  <span class="c1"># default to 2 ops per image</span>
    <span class="n">weight_idx</span> <span class="o">=</span> <span class="kc">None</span>  <span class="c1"># default to no probability weights for op choice</span>
    <span class="n">transforms</span> <span class="o">=</span> <span class="n">_RAND_TRANSFORMS</span>
    <span class="n">config</span> <span class="o">=</span> <span class="n">config_str</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;-&#39;</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="n">config</span><span class="p">:</span>
        <span class="n">cs</span> <span class="o">=</span> <span class="n">re</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="sa">r</span><span class="s1">&#39;(\d.*)&#39;</span><span class="p">,</span> <span class="n">c</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">cs</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">:</span>
            <span class="k">continue</span>
        <span class="n">key</span><span class="p">,</span> <span class="n">val</span> <span class="o">=</span> <span class="n">cs</span><span class="p">[:</span><span class="mi">2</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">key</span> <span class="o">==</span> <span class="s1">&#39;mstd&#39;</span><span class="p">:</span>
            <span class="c1"># noise param injected via hparams for now</span>
            <span class="n">hparams</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;magnitude_std&#39;</span><span class="p">,</span> <span class="nb">float</span><span class="p">(</span><span class="n">val</span><span class="p">))</span>
        <span class="k">elif</span> <span class="n">key</span> <span class="o">==</span> <span class="s1">&#39;inc&#39;</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">bool</span><span class="p">(</span><span class="n">val</span><span class="p">):</span>
                <span class="n">transforms</span> <span class="o">=</span> <span class="n">_RAND_INCREASING_TRANSFORMS</span>
        <span class="k">elif</span> <span class="n">key</span> <span class="o">==</span> <span class="s1">&#39;m&#39;</span><span class="p">:</span>
            <span class="n">magnitude</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
        <span class="k">elif</span> <span class="n">key</span> <span class="o">==</span> <span class="s1">&#39;n&#39;</span><span class="p">:</span>
            <span class="n">num_layers</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
        <span class="k">elif</span> <span class="n">key</span> <span class="o">==</span> <span class="s1">&#39;w&#39;</span><span class="p">:</span>
            <span class="n">weight_idx</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">assert</span> <span class="kc">False</span><span class="p">,</span> <span class="s1">&#39;Unknown RandAugment config section&#39;</span>
    <span class="n">ra_ops</span> <span class="o">=</span> <span class="n">rand_augment_ops</span><span class="p">(</span><span class="n">magnitude</span><span class="o">=</span><span class="n">magnitude</span><span class="p">,</span> <span class="n">hparams</span><span class="o">=</span><span class="n">hparams</span><span class="p">,</span> <span class="n">transforms</span><span class="o">=</span><span class="n">transforms</span><span class="p">)</span>
    <span class="n">choice_weights</span> <span class="o">=</span> <span class="kc">None</span> <span class="k">if</span> <span class="n">weight_idx</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">_select_rand_weights</span><span class="p">(</span><span class="n">weight_idx</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">RandAugment</span><span class="p">(</span><span class="n">ra_ops</span><span class="p">,</span> <span class="n">num_layers</span><span class="p">,</span> <span class="n">choice_weights</span><span class="o">=</span><span class="n">choice_weights</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.classification_datasets.cifar &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
        <script src="../../../../../_static/jquery.js"></script>
        <script src="../../../../../_static/underscore.js"></script>
        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../../_static/doctools.js"></script>
        <script src="../../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.classification_datasets.cifar</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.classification_datasets.cifar</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">Union</span>

<span class="kn">from</span> <span class="nn">torchvision.transforms</span> <span class="kn">import</span> <span class="n">Compose</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.factories.transforms_factory</span> <span class="kn">import</span> <span class="n">TransformsFactory</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.decorators.factory_decorator</span> <span class="kn">import</span> <span class="n">resolve_param</span>
<span class="kn">from</span> <span class="nn">torchvision.datasets</span> <span class="kn">import</span> <span class="n">CIFAR10</span><span class="p">,</span> <span class="n">CIFAR100</span>


<div class="viewcode-block" id="Cifar10"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.Cifar10">[docs]</a><span class="k">class</span> <span class="nc">Cifar10</span><span class="p">(</span><span class="n">CIFAR10</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    CIFAR10 Dataset</span>

<span class="sd">    :param root:                    Path for the data to be extracted</span>
<span class="sd">    :param train:                   Bool to load training (True) or validation (False) part of the dataset</span>
<span class="sd">    :param transforms:              List of transforms to apply sequentially on sample. Wrapped internally with torchvision.Compose</span>
<span class="sd">    :param target_transform:        Transform to apply to target output</span>
<span class="sd">    :param download:                Download (True) the dataset from source</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s2">&quot;transforms&quot;</span><span class="p">,</span> <span class="n">TransformsFactory</span><span class="p">())</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">root</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
        <span class="n">train</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
        <span class="n">transforms</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span class="p">,</span> <span class="nb">dict</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
        <span class="n">target_transform</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Callable</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
        <span class="n">download</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># TO KEEP BACKWARD COMPATABILITY, WILL BE REMOVED IN THE FUTURE ONCE WE ALLIGN TORCHVISION/NATIVE TRANSFORMS</span>
        <span class="c1"># TREATMENT IN FACTORIES (I.E STATING COMPOSE IN CONFIGS)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">transforms</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">transforms</span> <span class="o">=</span> <span class="n">Compose</span><span class="p">(</span><span class="n">transforms</span><span class="p">)</span>

        <span class="nb">super</span><span class="p">(</span><span class="n">Cifar10</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>
            <span class="n">root</span><span class="o">=</span><span class="n">root</span><span class="p">,</span>
            <span class="n">train</span><span class="o">=</span><span class="n">train</span><span class="p">,</span>
            <span class="n">transform</span><span class="o">=</span><span class="n">transforms</span><span class="p">,</span>
            <span class="n">target_transform</span><span class="o">=</span><span class="n">target_transform</span><span class="p">,</span>
            <span class="n">download</span><span class="o">=</span><span class="n">download</span><span class="p">,</span>
        <span class="p">)</span></div>


<div class="viewcode-block" id="Cifar100"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.Cifar100">[docs]</a><span class="k">class</span> <span class="nc">Cifar100</span><span class="p">(</span><span class="n">CIFAR100</span><span class="p">):</span>
    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s2">&quot;transforms&quot;</span><span class="p">,</span> <span class="n">TransformsFactory</span><span class="p">())</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">root</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
        <span class="n">train</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
        <span class="n">transforms</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span class="p">,</span> <span class="nb">dict</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
        <span class="n">target_transform</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Callable</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
        <span class="n">download</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        CIFAR100 Dataset</span>

<span class="sd">        :param root:                    Path for the data to be extracted</span>
<span class="sd">        :param train:                   Bool to load training (True) or validation (False) part of the dataset</span>
<span class="sd">        :param transforms:              List of transforms to apply sequentially on sample. Wrapped internally with torchvision.Compose</span>
<span class="sd">        :param target_transform:        Transform to apply to target output</span>
<span class="sd">        :param download:                Download (True) the dataset from source</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># TO KEEP BACKWARD COMPATABILITY, WILL BE REMOVED IN THE FUTURE ONCE WE ALLIGN TORCHVISION/NATIVE TRANSFORMS</span>
        <span class="c1"># TREATMENT IN FACTORIES (I.E STATING COMPOSE IN CONFIGS)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">transforms</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">transforms</span> <span class="o">=</span> <span class="n">Compose</span><span class="p">(</span><span class="n">transforms</span><span class="p">)</span>

        <span class="nb">super</span><span class="p">(</span><span class="n">Cifar100</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>
            <span class="n">root</span><span class="o">=</span><span class="n">root</span><span class="p">,</span>
            <span class="n">train</span><span class="o">=</span><span class="n">train</span><span class="p">,</span>
            <span class="n">transform</span><span class="o">=</span><span class="n">transforms</span><span class="p">,</span>
            <span class="n">target_transform</span><span class="o">=</span><span class="n">target_transform</span><span class="p">,</span>
            <span class="n">download</span><span class="o">=</span><span class="n">download</span><span class="p">,</span>
        <span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.classification_datasets.imagenet_dataset &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
        <script src="../../../../../_static/jquery.js"></script>
        <script src="../../../../../_static/underscore.js"></script>
        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../../_static/doctools.js"></script>
        <script src="../../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.classification_datasets.imagenet_dataset</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.classification_datasets.imagenet_dataset</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>

<span class="kn">import</span> <span class="nn">torchvision.datasets</span> <span class="k">as</span> <span class="nn">torch_datasets</span>
<span class="kn">from</span> <span class="nn">torchvision.transforms</span> <span class="kn">import</span> <span class="n">Compose</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.decorators.factory_decorator</span> <span class="kn">import</span> <span class="n">resolve_param</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.factories.transforms_factory</span> <span class="kn">import</span> <span class="n">TransformsFactory</span>


<div class="viewcode-block" id="ImageNetDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.ImageNetDataset">[docs]</a><span class="k">class</span> <span class="nc">ImageNetDataset</span><span class="p">(</span><span class="n">torch_datasets</span><span class="o">.</span><span class="n">ImageFolder</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;ImageNetDataset dataset&quot;&quot;&quot;</span>

    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s2">&quot;transforms&quot;</span><span class="p">,</span> <span class="n">factory</span><span class="o">=</span><span class="n">TransformsFactory</span><span class="p">())</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">root</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">transforms</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span class="p">,</span> <span class="nb">dict</span><span class="p">]</span> <span class="o">=</span> <span class="p">[],</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="c1"># TO KEEP BACKWARD COMPATABILITY, WILL BE REMOVED IN THE FUTURE ONCE WE ALLIGN TORCHVISION/NATIVE TRANSFORMS</span>
        <span class="c1"># TREATMENT IN FACTORIES (I.E STATING COMPOSE IN CONFIGS)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">transforms</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">transforms</span> <span class="o">=</span> <span class="n">Compose</span><span class="p">(</span><span class="n">transforms</span><span class="p">)</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ImageNetDataset</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><span class="n">root</span><span class="p">,</span> <span class="n">transform</span><span class="o">=</span><span class="n">transforms</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.data_augmentation &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.data_augmentation &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -91,9 +93,9 @@
 
                             <span class="kn">from</span> <span class="nn">torchvision.transforms</span> <span class="kn">import</span> <span class="n">RandomErasing</span>
                
 
                            -<div class="viewcode-block" id="DataAugmentation"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.DataAugmentation">[docs]</a><span class="k">class</span> <span class="nc">DataAugmentation</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="DataAugmentation"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.DataAugmentation">[docs]</a><span class="k">class</span> <span class="nc">DataAugmentation</span><span class="p">:</span>
                
 
                            -<div class="viewcode-block" id="DataAugmentation.to_tensor"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.DataAugmentation.to_tensor">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="DataAugmentation.to_tensor"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.DataAugmentation.to_tensor">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">to_tensor</span><span class="p">():</span>
                
 
                                     <span class="k">def</span> <span class="nf">_to_tensor</span><span class="p">(</span><span class="n">image</span><span class="p">):</span>
                
 
                                         <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">image</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
                
@@ -104,7 +106,7 @@
 
                                     <span class="k">return</span> <span class="n">_to_tensor</span></div>
                
 
                            -<div class="viewcode-block" id="DataAugmentation.normalize"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.DataAugmentation.normalize">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="DataAugmentation.normalize"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.DataAugmentation.normalize">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="n">mean</span><span class="p">,</span> <span class="n">std</span><span class="p">):</span>
                
 
                                     <span class="n">mean</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="n">mean</span><span class="p">)</span>
                
 
                                     <span class="n">std</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="n">std</span><span class="p">)</span>
                
@@ -116,7 +118,7 @@
 
                                     <span class="k">return</span> <span class="n">_normalize</span></div>
                
 
                            -<div class="viewcode-block" id="DataAugmentation.cutout"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.DataAugmentation.cutout">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="DataAugmentation.cutout"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.DataAugmentation.cutout">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">cutout</span><span class="p">(</span><span class="n">mask_size</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">cutout_inside</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">mask_color</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="m
                
 
                                     <span class="n">mask_size_half</span> <span class="o">=</span> <span class="n">mask_size</span> <span class="o">//</span> <span class="mi">2</span>
                
 
                                     <span class="n">offset</span> <span class="o">=</span> <span class="mi">1</span> <span class="k">if</span> <span class="n">mask_size</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">==</span> <span class="mi">0</span> <span class="k">else</span> <span class="mi">0</span>
                
@@ -159,7 +161,7 @@
 
                                                         <span class="p">[</span><span class="o">-</span><span class="mf">0.5836</span><span class="p">,</span> <span class="o">-</span><span class="mf">0.6948</span><span class="p">,</span> <span class="mf">0.4203</span><span class="p">]])}</span>
                
 
                            -<div class="viewcode-block" id="Lighting"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.Lighting">[docs]</a><span class="k">class</span> <span class="nc">Lighting</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
                
 
                            +<span class="k">class</span> <span class="nc">Lighting</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Lighting noise(AlexNet - style PCA - based noise)</span>
                
 
                             <span class="sd">    Taken from fastai Imagenet training -</span>
                
@@ -183,10 +185,10 @@
 
                                         <span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="n">alpha</span><span class="o">.</span><span class="n">view</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="o">.</span><span class="n">expand</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">))</span> \
                
 
                                         <span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">eigval</span><span class="o">.</span><span class="n">view</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="o">.</span><span class="n">expand</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span 
                
 
                                         <span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">()</span>
                
 
                            -        <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">rgb</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">expand_as</span><span class="p">(</span><span class="n">img</span><span
                
 
                            +        <span class="k">return</span> <span class="n">img</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">rgb</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">expand_as</span><span class="p">(</span><span class="n">img</span><span
                
 
                            -<div class="viewcode-block" id="RandomErase"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.RandomErase">[docs]</a><span class="k">class</span> <span class="nc">RandomErase</span><span class="p">(</span><span class="n">RandomErasing</span><span class="p">):</span>
                
 
                            +<span class="k">class</span> <span class="nc">RandomErase</span><span class="p">(</span><span class="n">RandomErasing</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    A simple class that translates the parameters supported in SuperGradient&#39;s code base</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
@@ -197,7 +199,7 @@
 
                                         <span class="n">value</span> <span class="o">=</span> <span class="nb">float</span><span class="p">(</span><span class="n">value</span><span class="p">)</span>
                
 
                                     <span class="k">except</span> <span class="ne">ValueError</span><span class="p">:</span>
                
 
                                         <span class="k">pass</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">p</span><span class="o">=</span><span class="n">probability</span><span class="p">,</span> <span class="n">value</span><span class="o">=</span><span class="n">value</span><span class="p">)</span></div>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">p</span><span class="o">=</span><span class="n">probability</span><span class="p">,</span> <span class="n">value</span><span class="o">=</span><span class="n">value</span><span class="p">)</span>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -227,4 +229,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.dataset_interfaces.dataset_interface &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
        <script src="../../../../../_static/jquery.js"></script>
        <script src="../../../../../_static/underscore.js"></script>
        <script src="../../../../../_static/doctools.js"></script>
    <script src="../../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.dataset_interfaces.dataset_interface</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.dataset_interfaces.dataset_interface</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os</span>

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">torchvision</span>
<span class="kn">import</span> <span class="nn">torchvision.datasets</span> <span class="k">as</span> <span class="nn">datasets</span>
<span class="kn">from</span> <span class="nn">torch.utils.data.distributed</span> <span class="kn">import</span> <span class="n">DistributedSampler</span>
<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">ConcatDataset</span><span class="p">,</span> <span class="n">BatchSampler</span><span class="p">,</span> <span class="n">DataLoader</span>
<span class="kn">import</span> <span class="nn">torchvision.transforms</span> <span class="k">as</span> <span class="nn">transforms</span>


<span class="kn">from</span> <span class="nn">super_gradients.common</span> <span class="kn">import</span> <span class="n">DatasetDataInterface</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.environment</span> <span class="kn">import</span> <span class="n">AWS_ENV_NAME</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>

<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.distributed_training_utils</span> <span class="kn">import</span> <span class="n">get_local_rank</span><span class="p">,</span> <span class="n">wait_for_the_master</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">get_param</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">DetectionTargetsFormat</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.datasets</span> <span class="kn">import</span> <span class="n">datasets_utils</span><span class="p">,</span> <span class="n">DataAugmentation</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.datasets_conf</span> <span class="kn">import</span> <span class="n">COCO_DETECTION_CLASSES_LIST</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.data_augmentation</span> <span class="kn">import</span> <span class="n">Lighting</span><span class="p">,</span> <span class="n">RandomErase</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.mixup</span> <span class="kn">import</span> <span class="n">CollateMixup</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.detection_datasets</span> <span class="kn">import</span> <span class="n">COCODetectionDataset</span><span class="p">,</span> <span class="n">PascalVOCDetectionDataset</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.samplers.infinite_sampler</span> <span class="kn">import</span> <span class="n">InfiniteSampler</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets</span> <span class="kn">import</span> <span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">,</span> \
    <span class="n">PascalAUG2012SegmentationDataSet</span><span class="p">,</span> <span class="n">CoCoSegmentationDataSet</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation</span> <span class="kn">import</span> <span class="n">CityscapesDataset</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets.supervisely_persons_segmentation</span> <span class="kn">import</span> \
    <span class="n">SuperviselyPersonsDataset</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.samplers.repeated_augmentation_sampler</span> <span class="kn">import</span> <span class="n">RepeatAugSampler</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.datasets_utils</span> <span class="kn">import</span> <span class="n">RandomResizedCropAndInterpolation</span><span class="p">,</span> <span class="n">worker_init_reset_seed</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.transforms.transforms</span> <span class="kn">import</span> <span class="n">DetectionMosaic</span><span class="p">,</span> <span class="n">DetectionMixup</span><span class="p">,</span> <span class="n">DetectionRandomAffine</span><span class="p">,</span>\
    <span class="n">DetectionTargetsFormatTransform</span><span class="p">,</span> <span class="n">DetectionPaddedRescale</span><span class="p">,</span> <span class="n">DetectionHSV</span><span class="p">,</span> <span class="n">DetectionHorizontalFlip</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.dataset_exceptions</span> <span class="kn">import</span> <span class="n">IllegalDatasetParameterException</span>


<span class="n">default_dataset_params</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;batch_size&quot;</span><span class="p">:</span> <span class="mi">64</span><span class="p">,</span> <span class="s2">&quot;val_batch_size&quot;</span><span class="p">:</span> <span class="mi">200</span><span class="p">,</span> <span class="s2">&quot;test_batch_size&quot;</span><span class="p">:</span> <span class="mi">200</span><span class="p">,</span> <span class="s2">&quot;dataset_dir&quot;</span><span class="p">:</span> <span class="s2">&quot;./data/&quot;</span><span class="p">,</span>
                          <span class="s2">&quot;s3_link&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">}</span>
<span class="n">LIBRARY_DATASETS</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;cifar10&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s1">&#39;class&#39;</span><span class="p">:</span> <span class="n">datasets</span><span class="o">.</span><span class="n">CIFAR10</span><span class="p">,</span> <span class="s1">&#39;mean&#39;</span><span class="p">:</span> <span class="p">(</span><span class="mf">0.4914</span><span class="p">,</span> <span class="mf">0.4822</span><span class="p">,</span> <span class="mf">0.4465</span><span class="p">),</span> <span class="s1">&#39;std&#39;</span><span class="p">:</span> <span class="p">(</span><span class="mf">0.2023</span><span class="p">,</span> <span class="mf">0.1994</span><span class="p">,</span> <span class="mf">0.2010</span><span class="p">)},</span>
    <span class="s2">&quot;cifar100&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s1">&#39;class&#39;</span><span class="p">:</span> <span class="n">datasets</span><span class="o">.</span><span class="n">CIFAR100</span><span class="p">,</span> <span class="s1">&#39;mean&#39;</span><span class="p">:</span> <span class="p">(</span><span class="mf">0.5071</span><span class="p">,</span> <span class="mf">0.4865</span><span class="p">,</span> <span class="mf">0.4409</span><span class="p">),</span> <span class="s1">&#39;std&#39;</span><span class="p">:</span> <span class="p">(</span><span class="mf">0.2673</span><span class="p">,</span> <span class="mf">0.2564</span><span class="p">,</span> <span class="mf">0.2762</span><span class="p">)},</span>
    <span class="s2">&quot;SVHN&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s1">&#39;class&#39;</span><span class="p">:</span> <span class="n">datasets</span><span class="o">.</span><span class="n">SVHN</span><span class="p">,</span> <span class="s1">&#39;mean&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span> <span class="s1">&#39;std&#39;</span><span class="p">:</span> <span class="kc">None</span><span class="p">}</span>
<span class="p">}</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>


<div class="viewcode-block" id="DatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">DatasetInterface</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    DatasetInterface - This class manages all of the &quot;communiation&quot; the Model has with the Data Sets</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">train_loader</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">val_loader</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">test_loader</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">classes</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        @param train_loader: torch.utils.data.Dataloader (optional) dataloader for training.</span>
<span class="sd">        @param test_loader: torch.utils.data.Dataloader (optional) dataloader for testing.</span>
<span class="sd">        @param classes: list of classes.</span>

<span class="sd">        Note: the above parameters will be discarded in case dataset_params is passed.</span>

<span class="sd">        @param dataset_params:</span>

<span class="sd">            - `batch_size` : int (default=64)</span>

<span class="sd">                Number of examples per batch for training. Large batch sizes are recommended.</span>

<span class="sd">            - `val_batch_size` : int (default=200)</span>

<span class="sd">                Number of examples per batch for validation. Large batch sizes are recommended.</span>

<span class="sd">            - `dataset_dir` : str (default=&quot;./data/&quot;)</span>

<span class="sd">                Directory location for the data. Data will be downloaded to this directory when getting it from a</span>
<span class="sd">                remote url.</span>

<span class="sd">            - `s3_link` : str (default=None)</span>

<span class="sd">                remote s3 link to download the data (optional).</span>

<span class="sd">            - `aug_repeat_count` : int (default=0)</span>

<span class="sd">                amount of repetitions (each repetition of an example is augmented differently) for each</span>
<span class="sd">                 example for the trainset.</span>

<span class="sd">        &quot;&quot;&quot;</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">HpmStruct</span><span class="p">(</span><span class="o">**</span><span class="n">default_dataset_params</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">override</span><span class="p">(</span><span class="o">**</span><span class="n">dataset_params</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">testset</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">val_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">test_loader</span> <span class="o">=</span> <span class="n">train_loader</span><span class="p">,</span> <span class="n">val_loader</span><span class="p">,</span> <span class="n">test_loader</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">classes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">s3_link</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">download_from_cloud</span><span class="p">()</span>

<div class="viewcode-block" id="DatasetInterface.download_from_cloud"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DatasetInterface.download_from_cloud">[docs]</a>    <span class="k">def</span> <span class="nf">download_from_cloud</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">s3_link</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">env_name</span> <span class="o">=</span> <span class="n">AWS_ENV_NAME</span>
            <span class="n">downloader</span> <span class="o">=</span> <span class="n">DatasetDataInterface</span><span class="p">(</span><span class="n">env</span><span class="o">=</span><span class="n">env_name</span><span class="p">)</span>
            <span class="n">target_dir</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">dataset_dir</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">target_dir</span><span class="p">):</span>
                <span class="n">os</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">target_dir</span><span class="p">)</span>
            <span class="n">downloader</span><span class="o">.</span><span class="n">load_remote_dataset_file</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">s3_link</span><span class="p">,</span> <span class="n">target_dir</span><span class="p">)</span></div>

<div class="viewcode-block" id="DatasetInterface.build_data_loaders"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DatasetInterface.build_data_loaders">[docs]</a>    <span class="k">def</span> <span class="nf">build_data_loaders</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span> <span class="n">train_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">val_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                           <span class="n">test_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">distributed_sampler</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">        define train, val (and optionally test) loaders. The method deals separately with distributed training and standard</span>
<span class="sd">        (non distributed, or parallel training). In the case of distributed training we need to rely on distributed</span>
<span class="sd">        samplers.</span>
<span class="sd">        :param batch_size_factor: int - factor to multiply the batch size (usually for multi gpu)</span>
<span class="sd">        :param num_workers: int - number of workers (parallel processes) for dataloaders</span>
<span class="sd">        :param train_batch_size: int - batch size for train loader, if None will be taken from dataset_params</span>
<span class="sd">        :param val_batch_size: int - batch size for val loader, if None will be taken from dataset_params</span>
<span class="sd">        :param distributed_sampler: boolean flag for distributed training mode</span>
<span class="sd">        :return: train_loader, val_loader, classes: list of classes</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># CHANGE THE BATCH SIZE ACCORDING TO THE NUMBER OF DEVICES - ONLY IN NON-DISTRIBUTED TRAINING MODE</span>
        <span class="c1"># IN DISTRIBUTED MODE WE NEED DISTRIBUTED SAMPLERS</span>
        <span class="c1"># NO SHUFFLE IN DISTRIBUTED TRAINING</span>

        <span class="n">aug_repeat_count</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;aug_repeat_count&quot;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">aug_repeat_count</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">distributed_sampler</span><span class="p">:</span>
            <span class="k">raise</span> <span class="n">IllegalDatasetParameterException</span><span class="p">(</span><span class="s2">&quot;repeated augmentation is only supported with DDP.&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">distributed_sampler</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="n">train_sampler</span> <span class="o">=</span> <span class="n">RepeatAugSampler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">,</span>
                                             <span class="n">num_repeats</span><span class="o">=</span><span class="n">aug_repeat_count</span><span class="p">)</span> <span class="k">if</span> <span class="n">aug_repeat_count</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">DistributedSampler</span><span class="p">(</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">)</span>
            <span class="n">val_sampler</span> <span class="o">=</span> <span class="n">DistributedSampler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">)</span>
            <span class="n">test_sampler</span> <span class="o">=</span> <span class="n">DistributedSampler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">testset</span><span class="p">)</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">testset</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="kc">None</span>
            <span class="n">train_shuffle</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span> <span class="o">=</span> <span class="n">batch_size_factor</span>
            <span class="n">train_sampler</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">val_sampler</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">test_sampler</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">train_shuffle</span> <span class="o">=</span> <span class="kc">True</span>

        <span class="k">if</span> <span class="n">train_batch_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">train_batch_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">batch_size</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span>
        <span class="k">if</span> <span class="n">val_batch_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">val_batch_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_batch_size</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span>
        <span class="k">if</span> <span class="n">test_batch_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">test_batch_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">test_batch_size</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span>

        <span class="n">train_loader_drop_last</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;train_loader_drop_last&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

        <span class="n">cutmix</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;cutmix&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
        <span class="n">cutmix_params</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;cutmix_params&#39;</span><span class="p">)</span>

        <span class="c1"># WRAPPING collate_fn</span>
        <span class="n">train_collate_fn</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">,</span> <span class="s1">&#39;collate_fn&#39;</span><span class="p">)</span>
        <span class="n">val_collate_fn</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">,</span> <span class="s1">&#39;collate_fn&#39;</span><span class="p">)</span>
        <span class="n">test_collate_fn</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">testset</span><span class="p">,</span> <span class="s1">&#39;collate_fn&#39;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">cutmix</span> <span class="ow">and</span> <span class="n">train_collate_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="n">IllegalDatasetParameterException</span><span class="p">(</span><span class="s2">&quot;cutmix and collate function cannot be used together&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">cutmix</span><span class="p">:</span>
            <span class="c1"># FIXME - cutmix should be available only in classification dataset. once we make sure all classification</span>
            <span class="c1"># datasets inherit from the same super class, we should move cutmix code to that class</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Cutmix/mixup was enabled. This feature is currently supported only &quot;</span>
                           <span class="s2">&quot;for classification datasets.&quot;</span><span class="p">)</span>
            <span class="n">train_collate_fn</span> <span class="o">=</span> <span class="n">CollateMixup</span><span class="p">(</span><span class="o">**</span><span class="n">cutmix_params</span><span class="p">)</span>

        <span class="c1"># FIXME - UNDERSTAND IF THE num_replicas VARIBALE IS NEEDED</span>
        <span class="c1"># train_sampler = DistributedSampler(self.trainset,</span>
        <span class="c1">#                                    num_replicas=distributed_gpus_num) if distributed_sampler else None</span>
        <span class="c1"># val_sampler = DistributedSampler(self.valset,</span>
        <span class="c1">#                                   num_replicas=distributed_gpus_num) if distributed_sampler else None</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">,</span>
                                                        <span class="n">batch_size</span><span class="o">=</span><span class="n">train_batch_size</span><span class="p">,</span>
                                                        <span class="n">shuffle</span><span class="o">=</span><span class="n">train_shuffle</span><span class="p">,</span>
                                                        <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                                        <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                        <span class="n">sampler</span><span class="o">=</span><span class="n">train_sampler</span><span class="p">,</span>
                                                        <span class="n">collate_fn</span><span class="o">=</span><span class="n">train_collate_fn</span><span class="p">,</span>
                                                        <span class="n">drop_last</span><span class="o">=</span><span class="n">train_loader_drop_last</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">val_loader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">,</span>
                                                      <span class="n">batch_size</span><span class="o">=</span><span class="n">val_batch_size</span><span class="p">,</span>
                                                      <span class="n">shuffle</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                                                      <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                                      <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                      <span class="n">sampler</span><span class="o">=</span><span class="n">val_sampler</span><span class="p">,</span>
                                                      <span class="n">collate_fn</span><span class="o">=</span><span class="n">val_collate_fn</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">testset</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">test_loader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">testset</span><span class="p">,</span>
                                                           <span class="n">batch_size</span><span class="o">=</span><span class="n">test_batch_size</span><span class="p">,</span>
                                                           <span class="n">shuffle</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                                                           <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                                           <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                           <span class="n">sampler</span><span class="o">=</span><span class="n">test_sampler</span><span class="p">,</span>
                                                           <span class="n">collate_fn</span><span class="o">=</span><span class="n">test_collate_fn</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span></div>

<div class="viewcode-block" id="DatasetInterface.get_data_loaders"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DatasetInterface.get_data_loaders">[docs]</a>    <span class="k">def</span> <span class="nf">get_data_loaders</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Get self.train_loader, self.val_loader, self.test_loader, self.classes.</span>

<span class="sd">        If the data loaders haven&#39;t been initialized yet, build them first.</span>

<span class="sd">        :param kwargs: kwargs are passed to build_data_loaders.</span>

<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">val_loader</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">build_data_loaders</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">val_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">test_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span></div>

<div class="viewcode-block" id="DatasetInterface.get_val_sample"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DatasetInterface.get_val_sample">[docs]</a>    <span class="k">def</span> <span class="nf">get_val_sample</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_samples</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">num_samples</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Tried to load more samples than val-set size&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">num_samples</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span><span class="n">num_samples</span><span class="p">]</span></div>

<div class="viewcode-block" id="DatasetInterface.get_dataset_params"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DatasetInterface.get_dataset_params">[docs]</a>    <span class="k">def</span> <span class="nf">get_dataset_params</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span></div>

<div class="viewcode-block" id="DatasetInterface.print_dataset_details"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DatasetInterface.print_dataset_details">[docs]</a>    <span class="k">def</span> <span class="nf">print_dataset_details</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{}</span><span class="s2"> training samples, </span><span class="si">{}</span><span class="s2"> val samples, </span><span class="si">{}</span><span class="s2"> classes&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">),</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">),</span>
                                                                             <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span><span class="p">)))</span></div></div>


<div class="viewcode-block" id="ExternalDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.ExternalDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">ExternalDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">train_loader</span><span class="p">,</span> <span class="n">val_loader</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{}):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        ExternalDatasetInterface - A wrapper for external dataset interface that gets dataloaders from keras/TF</span>
<span class="sd">        and converts them to Torch-like dataloaders that return torch.Tensors after</span>
<span class="sd">        optional collate_fn while maintaining the same interface (connect_dataset_interface etc.)</span>
<span class="sd">            :train_loader:       The external train_loader</span>
<span class="sd">            :val_loader:        The external val_loader</span>
<span class="sd">            :num_classes:        The number of classes</span>
<span class="sd">            :dataset_params      The dict that includes the batch_size and/or the collate_fn</span>

<span class="sd">            :return:             DataLoaders that generate torch.Tensors batches after collate_fn</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span> <span class="o">=</span> <span class="n">train_loader</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">val_loader</span> <span class="o">=</span> <span class="n">val_loader</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">num_classes</span>

<div class="viewcode-block" id="ExternalDatasetInterface.get_data_loaders"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.ExternalDatasetInterface.get_data_loaders">[docs]</a>    <span class="k">def</span> <span class="nf">get_data_loaders</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size_factor</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span> <span class="n">num_workers</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">8</span><span class="p">,</span> <span class="n">train_batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                         <span class="n">val_batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">distributed_sampler</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>

        <span class="c1"># CHANGE THE BATCH SIZE ACCORDING TO THE NUMBER OF DEVICES - ONLY IN NON-DISTRIBUED TRAINING MODE</span>
        <span class="c1"># IN DISTRIBUTED MODE WE NEED DISTRIBUTED SAMPLERS</span>
        <span class="c1"># NO SHUFFLE IN DISTRIBUTED TRAINING</span>
        <span class="k">if</span> <span class="n">distributed_sampler</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="n">train_sampler</span> <span class="o">=</span> <span class="n">DistributedSampler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">,</span> <span class="n">shuffle</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
            <span class="n">val_sampler</span> <span class="o">=</span> <span class="n">DistributedSampler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">)</span>
            <span class="n">train_shuffle</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span> <span class="o">=</span> <span class="n">batch_size_factor</span>
            <span class="n">train_sampler</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">val_sampler</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">train_shuffle</span> <span class="o">=</span> <span class="kc">True</span>

        <span class="k">if</span> <span class="n">train_batch_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">train_batch_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">batch_size</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span>
        <span class="k">if</span> <span class="n">val_batch_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">val_batch_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_batch_size</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_size_factor</span>

        <span class="n">train_loader_drop_last</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;train_loader_drop_last&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

        <span class="c1"># WRAPPING collate_fn</span>
        <span class="n">train_collate_fn</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;train_collate_fn&#39;</span><span class="p">)</span>
        <span class="n">val_collate_fn</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;val_collate_fn&#39;</span><span class="p">)</span>

        <span class="c1"># FIXME - UNDERSTAND IF THE num_replicas VARIBALE IS NEEDED</span>
        <span class="c1"># train_sampler = DistributedSampler(self.trainset,</span>
        <span class="c1">#                                    num_replicas=distributed_gpus_num) if distributed_sampler else None</span>
        <span class="c1"># val_sampler = DistributedSampler(self.valset,</span>
        <span class="c1">#                                   num_replicas=distributed_gpus_num) if distributed_sampler else None</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">torch_train_loader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span>
                                                              <span class="n">batch_size</span><span class="o">=</span><span class="n">train_batch_size</span><span class="p">,</span>
                                                              <span class="n">shuffle</span><span class="o">=</span><span class="n">train_shuffle</span><span class="p">,</span>
                                                              <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                                              <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                              <span class="n">sampler</span><span class="o">=</span><span class="n">train_sampler</span><span class="p">,</span>
                                                              <span class="n">collate_fn</span><span class="o">=</span><span class="n">train_collate_fn</span><span class="p">,</span>
                                                              <span class="n">drop_last</span><span class="o">=</span><span class="n">train_loader_drop_last</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">torch_val_loader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">val_loader</span><span class="p">,</span>
                                                            <span class="n">batch_size</span><span class="o">=</span><span class="n">val_batch_size</span><span class="p">,</span>
                                                            <span class="n">shuffle</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                                                            <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                                            <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                            <span class="n">sampler</span><span class="o">=</span><span class="n">val_sampler</span><span class="p">,</span>
                                                            <span class="n">collate_fn</span><span class="o">=</span><span class="n">val_collate_fn</span><span class="p">)</span>

        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">torch_train_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">torch_val_loader</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span></div></div>


<div class="viewcode-block" id="LibraryDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.LibraryDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">LibraryDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;cifar10&quot;</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">to_cutout</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">LibraryDatasetInterface</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><span class="n">dataset_params</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_name</span> <span class="o">=</span> <span class="n">name</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_name</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">LIBRARY_DATASETS</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s1">&#39;dataset not found&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span> <span class="o">=</span> <span class="n">LIBRARY_DATASETS</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_name</span><span class="p">]</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">]</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">trainset</span> <span class="o">=</span> <span class="n">torchvision</span><span class="o">.</span><span class="n">datasets</span><span class="o">.</span><span class="n">SVHN</span><span class="p">(</span><span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">dataset_dir</span><span class="p">,</span> <span class="n">split</span><span class="o">=</span><span class="s1">&#39;train&#39;</span><span class="p">,</span> <span class="n">download</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                 <span class="n">transform</span><span class="o">=</span><span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">())</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;std&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">datasets_utils</span><span class="o">.</span><span class="n">get_mean_and_std</span><span class="p">(</span><span class="n">trainset</span><span class="p">)</span>

        <span class="c1"># OVERWRITE MEAN AND STD IF DEFINED IN DATASET PARAMS</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;img_mean&#39;</span><span class="p">,</span>
                                                               <span class="n">default_val</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">])</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;std&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;img_std&#39;</span><span class="p">,</span>
                                                              <span class="n">default_val</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;std&#39;</span><span class="p">])</span>

        <span class="n">crop_size</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;crop_size&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mi">32</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">to_cutout</span><span class="p">:</span>
            <span class="n">transform_train</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomCrop</span><span class="p">(</span><span class="n">crop_size</span><span class="p">,</span> <span class="n">padding</span><span class="o">=</span><span class="mi">4</span><span class="p">),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomHorizontalFlip</span><span class="p">(),</span>
                <span class="n">DataAugmentation</span><span class="o">.</span><span class="n">normalize</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;std&#39;</span><span class="p">]),</span>
                <span class="n">DataAugmentation</span><span class="o">.</span><span class="n">cutout</span><span class="p">(</span><span class="mi">16</span><span class="p">),</span>
                <span class="n">DataAugmentation</span><span class="o">.</span><span class="n">to_tensor</span><span class="p">()</span>
            <span class="p">])</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">transform_train</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomCrop</span><span class="p">(</span><span class="n">crop_size</span><span class="p">,</span> <span class="n">padding</span><span class="o">=</span><span class="mi">4</span><span class="p">),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomHorizontalFlip</span><span class="p">(),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">Normalize</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;std&#39;</span><span class="p">]),</span>
            <span class="p">])</span>

        <span class="n">transform_val</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">Normalize</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;mean&#39;</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s1">&#39;std&#39;</span><span class="p">]),</span>
        <span class="p">])</span>
        <span class="n">dataset_cls</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">lib_dataset_params</span><span class="p">[</span><span class="s2">&quot;class&quot;</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">dataset_cls</span><span class="p">(</span><span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">dataset_dir</span><span class="p">,</span> <span class="n">train</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">download</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                    <span class="n">transform</span><span class="o">=</span><span class="n">transform_train</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">dataset_cls</span><span class="p">(</span><span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">dataset_dir</span><span class="p">,</span> <span class="n">train</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">download</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                  <span class="n">transform</span><span class="o">=</span><span class="n">transform_val</span><span class="p">)</span></div>


<div class="viewcode-block" id="Cifar10DatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.Cifar10DatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">Cifar10DatasetInterface</span><span class="p">(</span><span class="n">LibraryDatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{}):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">Cifar10DatasetInterface</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><span class="n">name</span><span class="o">=</span><span class="s2">&quot;cifar10&quot;</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="Cifar100DatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.Cifar100DatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">Cifar100DatasetInterface</span><span class="p">(</span><span class="n">LibraryDatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{}):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">Cifar100DatasetInterface</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><span class="n">name</span><span class="o">=</span><span class="s2">&quot;cifar100&quot;</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="TestDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.TestDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">TestDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">trainset</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">classes</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">TestDatasetInterface</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><span class="n">dataset_params</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">trainset</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">testset</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">classes</span>

<div class="viewcode-block" id="TestDatasetInterface.get_data_loaders"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.TestDatasetInterface.get_data_loaders">[docs]</a>    <span class="k">def</span> <span class="nf">get_data_loaders</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span> <span class="n">train_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">val_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                         <span class="n">distributed_sampler</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</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">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span>
        <span class="k">return</span> <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">get_data_loaders</span><span class="p">(</span><span class="n">batch_size_factor</span><span class="o">=</span><span class="n">batch_size_factor</span><span class="p">,</span>
                                        <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                        <span class="n">train_batch_size</span><span class="o">=</span><span class="n">train_batch_size</span><span class="p">,</span>
                                        <span class="n">val_batch_size</span><span class="o">=</span><span class="n">val_batch_size</span><span class="p">,</span>
                                        <span class="n">distributed_sampler</span><span class="o">=</span><span class="n">distributed_sampler</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="ClassificationTestDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.ClassificationTestDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">ClassificationTestDatasetInterface</span><span class="p">(</span><span class="n">TestDatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">image_size</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">5</span><span class="p">,</span> <span class="n">classes</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">trainset</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">TensorDataset</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">))),</span>
                                                  <span class="n">torch</span><span class="o">.</span><span class="n">LongTensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">batch_size</span><span class="p">))))</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ClassificationTestDatasetInterface</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><span class="n">trainset</span><span class="o">=</span><span class="n">trainset</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
                                                                 <span class="n">classes</span><span class="o">=</span><span class="n">classes</span><span class="p">)</span></div>


<div class="viewcode-block" id="SegmentationTestDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.SegmentationTestDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">SegmentationTestDatasetInterface</span><span class="p">(</span><span class="n">TestDatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">image_size</span><span class="o">=</span><span class="mi">512</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">4</span><span class="p">):</span>
        <span class="n">trainset</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">TensorDataset</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">))),</span>
                                                  <span class="n">torch</span><span class="o">.</span><span class="n">LongTensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">))))</span>

        <span class="nb">super</span><span class="p">(</span><span class="n">SegmentationTestDatasetInterface</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><span class="n">trainset</span><span class="o">=</span><span class="n">trainset</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="DetectionTestDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DetectionTestDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">DetectionTestDatasetInterface</span><span class="p">(</span><span class="n">TestDatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">image_size</span><span class="o">=</span><span class="mi">320</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span> <span class="n">classes</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">trainset</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">TensorDataset</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="n">image_size</span><span class="p">))),</span>
                                                  <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="mi">6</span><span class="p">))))</span>

        <span class="nb">super</span><span class="p">(</span><span class="n">DetectionTestDatasetInterface</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><span class="n">trainset</span><span class="o">=</span><span class="n">trainset</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
                                                            <span class="n">classes</span><span class="o">=</span><span class="n">classes</span><span class="p">)</span></div>


<div class="viewcode-block" id="TestYoloDetectionDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.TestYoloDetectionDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">TestYoloDetectionDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    note: the output size is (batch_size, 6) in the test while in real training</span>
<span class="sd">    the size of axis 0 can vary (the number of bounding boxes)</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">input_dims</span><span class="o">=</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">32</span><span class="p">,</span> <span class="mi">32</span><span class="p">),</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">5</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">TensorDataset</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="o">*</span><span class="n">input_dims</span><span class="p">)),</span>
                                                       <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="mi">6</span><span class="p">)))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="p">[</span><span class="mi">0</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">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span></div>


<div class="viewcode-block" id="ImageNetDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.ImageNetDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">ImageNetDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">data_dir</span><span class="o">=</span><span class="s2">&quot;/data/Imagenet&quot;</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ImageNetDatasetInterface</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><span class="n">dataset_params</span><span class="p">)</span>

        <span class="n">data_dir</span> <span class="o">=</span> <span class="n">dataset_params</span><span class="p">[</span><span class="s1">&#39;dataset_dir&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;dataset_dir&#39;</span> <span class="ow">in</span> <span class="n">dataset_params</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> <span class="k">else</span> <span class="n">data_dir</span>
        <span class="n">traindir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">data_dir</span><span class="p">),</span> <span class="s1">&#39;train&#39;</span><span class="p">)</span>
        <span class="n">valdir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">data_dir</span><span class="p">,</span> <span class="s1">&#39;val&#39;</span><span class="p">)</span>
        <span class="n">img_mean</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;img_mean&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="p">[</span><span class="mf">0.485</span><span class="p">,</span> <span class="mf">0.456</span><span class="p">,</span> <span class="mf">0.406</span><span class="p">])</span>
        <span class="n">img_std</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;img_std&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="p">[</span><span class="mf">0.229</span><span class="p">,</span> <span class="mf">0.224</span><span class="p">,</span> <span class="mf">0.225</span><span class="p">])</span>
        <span class="n">normalize</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Normalize</span><span class="p">(</span><span class="n">mean</span><span class="o">=</span><span class="n">img_mean</span><span class="p">,</span>
                                         <span class="n">std</span><span class="o">=</span><span class="n">img_std</span><span class="p">)</span>

        <span class="n">crop_size</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;crop_size&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mi">224</span><span class="p">)</span>
        <span class="n">resize_size</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;resize_size&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mi">256</span><span class="p">)</span>
        <span class="n">color_jitter</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;color_jitter&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mf">0.0</span><span class="p">)</span>
        <span class="n">imagenet_pca_aug</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;imagenet_pca_aug&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mf">0.0</span><span class="p">)</span>
        <span class="n">train_interpolation</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;train_interpolation&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="s1">&#39;default&#39;</span><span class="p">)</span>
        <span class="n">rand_augment_config_string</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;rand_augment_config_string&#39;</span><span class="p">,</span>
                                                          <span class="n">default_val</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>

        <span class="n">color_jitter</span> <span class="o">=</span> <span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">color_jitter</span><span class="p">),)</span> <span class="o">*</span> <span class="mi">3</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">color_jitter</span><span class="p">,</span> <span class="nb">float</span><span class="p">)</span> <span class="k">else</span> <span class="n">color_jitter</span>
        <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">color_jitter</span><span class="p">)</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">),</span> <span class="s2">&quot;color_jitter must be a scalar or tuple of len 3 or 4&quot;</span>

        <span class="n">color_augmentation</span> <span class="o">=</span> <span class="n">datasets_utils</span><span class="o">.</span><span class="n">get_color_augmentation</span><span class="p">(</span><span class="n">rand_augment_config_string</span><span class="p">,</span> <span class="n">color_jitter</span><span class="p">,</span>
                                                                   <span class="n">crop_size</span><span class="o">=</span><span class="n">crop_size</span><span class="p">,</span> <span class="n">img_mean</span><span class="o">=</span><span class="n">img_mean</span><span class="p">)</span>

        <span class="n">train_transformation_list</span> <span class="o">=</span> <span class="p">[</span>
            <span class="n">RandomResizedCropAndInterpolation</span><span class="p">(</span><span class="n">crop_size</span><span class="p">,</span> <span class="n">interpolation</span><span class="o">=</span><span class="n">train_interpolation</span><span class="p">),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">RandomHorizontalFlip</span><span class="p">(),</span>
            <span class="n">color_augmentation</span><span class="p">,</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
            <span class="n">Lighting</span><span class="p">(</span><span class="n">imagenet_pca_aug</span><span class="p">),</span>
            <span class="n">normalize</span><span class="p">]</span>

        <span class="n">rndm_erase_prob</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;random_erase_prob&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mf">0.</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">rndm_erase_prob</span><span class="p">:</span>
            <span class="n">train_transformation_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">RandomErase</span><span class="p">(</span><span class="n">rndm_erase_prob</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">random_erase_value</span><span class="p">))</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">datasets</span><span class="o">.</span><span class="n">ImageFolder</span><span class="p">(</span><span class="n">traindir</span><span class="p">,</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">(</span><span class="n">train_transformation_list</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">datasets</span><span class="o">.</span><span class="n">ImageFolder</span><span class="p">(</span><span class="n">valdir</span><span class="p">,</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">Resize</span><span class="p">(</span><span class="n">resize_size</span><span class="p">),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">CenterCrop</span><span class="p">(</span><span class="n">crop_size</span><span class="p">),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
            <span class="n">normalize</span><span class="p">,</span>
        <span class="p">]))</span></div>


<div class="viewcode-block" id="TinyImageNetDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.TinyImageNetDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">TinyImageNetDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">data_dir</span><span class="o">=</span><span class="s2">&quot;/data/TinyImagenet&quot;</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">TinyImageNetDatasetInterface</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><span class="n">dataset_params</span><span class="p">)</span>

        <span class="n">data_dir</span> <span class="o">=</span> <span class="n">dataset_params</span><span class="p">[</span><span class="s1">&#39;dataset_dir&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;dataset_dir&#39;</span> <span class="ow">in</span> <span class="n">dataset_params</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> <span class="k">else</span> <span class="n">data_dir</span>
        <span class="n">traindir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">data_dir</span><span class="p">),</span> <span class="s1">&#39;train&#39;</span><span class="p">)</span>
        <span class="n">valdir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">data_dir</span><span class="p">,</span> <span class="s1">&#39;val&#39;</span><span class="p">)</span>

        <span class="n">img_mean</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;img_mean&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="p">[</span><span class="mf">0.4802</span><span class="p">,</span> <span class="mf">0.4481</span><span class="p">,</span> <span class="mf">0.3975</span><span class="p">])</span>
        <span class="n">img_std</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;img_std&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="p">[</span><span class="mf">0.2770</span><span class="p">,</span> <span class="mf">0.2691</span><span class="p">,</span> <span class="mf">0.2821</span><span class="p">])</span>
        <span class="n">normalize</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Normalize</span><span class="p">(</span><span class="n">mean</span><span class="o">=</span><span class="n">img_mean</span><span class="p">,</span>
                                         <span class="n">std</span><span class="o">=</span><span class="n">img_std</span><span class="p">)</span>

        <span class="n">crop_size</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;crop_size&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mi">56</span><span class="p">)</span>
        <span class="n">resize_size</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;resize_size&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="mi">64</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">datasets</span><span class="o">.</span><span class="n">ImageFolder</span><span class="p">(</span>
            <span class="n">traindir</span><span class="p">,</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomResizedCrop</span><span class="p">(</span><span class="n">crop_size</span><span class="p">),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomHorizontalFlip</span><span class="p">(),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
                <span class="n">normalize</span><span class="p">,</span>
            <span class="p">]))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">datasets</span><span class="o">.</span><span class="n">ImageFolder</span><span class="p">(</span><span class="n">valdir</span><span class="p">,</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">Resize</span><span class="p">(</span><span class="n">resize_size</span><span class="p">),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">CenterCrop</span><span class="p">(</span><span class="n">crop_size</span><span class="p">),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
            <span class="n">normalize</span><span class="p">,</span>
        <span class="p">]))</span></div>


<div class="viewcode-block" id="ClassificationDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.ClassificationDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">ClassificationDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">normalization_mean</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">normalization_std</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">resolution</span><span class="o">=</span><span class="mi">64</span><span class="p">,</span>
                 <span class="n">dataset_params</span><span class="o">=</span><span class="p">{}):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ClassificationDatasetInterface</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><span class="n">dataset_params</span><span class="p">)</span>
        <span class="n">data_dir</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">dataset_dir</span>

        <span class="n">traindir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">data_dir</span><span class="p">),</span> <span class="s1">&#39;train&#39;</span><span class="p">)</span>
        <span class="n">valdir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">data_dir</span><span class="p">,</span> <span class="s1">&#39;val&#39;</span><span class="p">)</span>
        <span class="n">normalize</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Normalize</span><span class="p">(</span><span class="n">mean</span><span class="o">=</span><span class="n">normalization_mean</span><span class="p">,</span>
                                         <span class="n">std</span><span class="o">=</span><span class="n">normalization_std</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">datasets</span><span class="o">.</span><span class="n">ImageFolder</span><span class="p">(</span>
            <span class="n">traindir</span><span class="p">,</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomResizedCrop</span><span class="p">(</span><span class="n">resolution</span><span class="p">),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">RandomHorizontalFlip</span><span class="p">(),</span>
                <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
                <span class="n">normalize</span><span class="p">,</span>
            <span class="p">]))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">datasets</span><span class="o">.</span><span class="n">ImageFolder</span><span class="p">(</span><span class="n">valdir</span><span class="p">,</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">Resize</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">resolution</span> <span class="o">*</span> <span class="mf">1.15</span><span class="p">)),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">CenterCrop</span><span class="p">(</span><span class="n">resolution</span><span class="p">),</span>
            <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">(),</span>
            <span class="n">normalize</span><span class="p">,</span>
        <span class="p">]))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span> <span class="o">=</span> <span class="n">data_dir</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">normalization_mean</span> <span class="o">=</span> <span class="n">normalization_mean</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">normalization_std</span> <span class="o">=</span> <span class="n">normalization_std</span></div>


<div class="viewcode-block" id="PascalVOC2012SegmentationDataSetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.PascalVOC2012SegmentationDataSetInterface">[docs]</a><span class="k">class</span> <span class="nc">PascalVOC2012SegmentationDataSetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache_labels</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">cache_images</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">dataset_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">dataset_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span> <span class="o">=</span> <span class="n">dataset_params</span><span class="p">[</span><span class="s1">&#39;dataset_dir&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;dataset_dir&#39;</span> <span class="ow">in</span> <span class="n">dataset_params</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> \
            <span class="k">else</span> <span class="s1">&#39;/data/pascal_voc_2012/VOCdevkit/VOC2012/&#39;</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">(</span><span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span>
                                                         <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;ImageSets/Segmentation/train.txt&#39;</span><span class="p">,</span>
                                                         <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s1">&#39;JPEGImages&#39;</span><span class="p">,</span>
                                                         <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s1">&#39;SegmentationClass&#39;</span><span class="p">,</span> <span class="n">augment</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                         <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
                                                         <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">(</span><span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span>
                                                       <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;ImageSets/Segmentation/val.txt&#39;</span><span class="p">,</span>
                                                       <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s1">&#39;JPEGImages&#39;</span><span class="p">,</span>
                                                       <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s1">&#39;SegmentationClass&#39;</span><span class="p">,</span> <span class="n">augment</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                       <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
                                                       <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span></div>


<div class="viewcode-block" id="PascalAUG2012SegmentationDataSetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.PascalAUG2012SegmentationDataSetInterface">[docs]</a><span class="k">class</span> <span class="nc">PascalAUG2012SegmentationDataSetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache_labels</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">cache_images</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">dataset_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">dataset_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span> <span class="o">=</span> <span class="n">dataset_params</span><span class="p">[</span><span class="s1">&#39;dataset_dir&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;dataset_dir&#39;</span> <span class="ow">in</span> <span class="n">dataset_params</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> \
            <span class="k">else</span> <span class="s1">&#39;/data/pascal_voc_2012/VOCaug/dataset/&#39;</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">PascalAUG2012SegmentationDataSet</span><span class="p">(</span>
            <span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;trainval.txt&#39;</span><span class="p">,</span>
            <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s1">&#39;img&#39;</span><span class="p">,</span>
            <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s1">&#39;cls&#39;</span><span class="p">,</span> <span class="n">augment</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">PascalAUG2012SegmentationDataSet</span><span class="p">(</span>
            <span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;val.txt&#39;</span><span class="p">,</span>
            <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s1">&#39;img&#39;</span><span class="p">,</span>
            <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s1">&#39;cls&#39;</span><span class="p">,</span> <span class="n">augment</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span> <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span></div>


<div class="viewcode-block" id="CoCoDataSetInterfaceBase"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.CoCoDataSetInterfaceBase">[docs]</a><span class="k">class</span> <span class="nc">CoCoDataSetInterfaceBase</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">dataset_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">dataset_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span> <span class="o">=</span> <span class="n">dataset_params</span><span class="p">[</span><span class="s1">&#39;dataset_dir&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;dataset_dir&#39;</span> <span class="ow">in</span> <span class="n">dataset_params</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> <span class="k">else</span> <span class="s1">&#39;/data/coco/&#39;</span></div>


<div class="viewcode-block" id="CoCoSegmentationDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.CoCoSegmentationDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">CoCoSegmentationDatasetInterface</span><span class="p">(</span><span class="n">CoCoDataSetInterfaceBase</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache_labels</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">cache_images</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                 <span class="n">dataset_classes_inclusion_tuples_list</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">CoCoSegmentationDataSet</span><span class="p">(</span>
            <span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;instances_train2017.json&#39;</span><span class="p">,</span>
            <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s1">&#39;images/train2017&#39;</span><span class="p">,</span>
            <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s1">&#39;annotations&#39;</span><span class="p">,</span> <span class="n">augment</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
            <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">,</span>
            <span class="n">dataset_classes_inclusion_tuples_list</span><span class="o">=</span><span class="n">dataset_classes_inclusion_tuples_list</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">CoCoSegmentationDataSet</span><span class="p">(</span>
            <span class="n">root</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;instances_val2017.json&#39;</span><span class="p">,</span>
            <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s1">&#39;images/val2017&#39;</span><span class="p">,</span>
            <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s1">&#39;annotations&#39;</span><span class="p">,</span> <span class="n">augment</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
            <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">,</span>
            <span class="n">dataset_classes_inclusion_tuples_list</span><span class="o">=</span><span class="n">dataset_classes_inclusion_tuples_list</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">coco_classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span></div>


<div class="viewcode-block" id="CityscapesDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.CityscapesDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">CityscapesDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache_labels</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">cache_images</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>
        <span class="n">root_dir</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;dataset_dir&quot;</span><span class="p">,</span> <span class="s2">&quot;/data/cityscapes&quot;</span><span class="p">)</span>
        <span class="n">img_size</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;img_size&quot;</span><span class="p">,</span> <span class="mi">1024</span><span class="p">)</span>
        <span class="n">crop_size</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;crop_size&quot;</span><span class="p">,</span> <span class="mi">512</span><span class="p">)</span>
        <span class="n">image_mask_transforms</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;image_mask_transforms&quot;</span><span class="p">)</span>
        <span class="n">image_mask_transforms_aug</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;image_mask_transforms_aug&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">CityscapesDataset</span><span class="p">(</span>
            <span class="n">root_dir</span><span class="o">=</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;lists/train.lst&#39;</span><span class="p">,</span>
            <span class="n">labels_csv_path</span><span class="o">=</span><span class="s2">&quot;lists/labels.csv&quot;</span><span class="p">,</span>
            <span class="n">img_size</span><span class="o">=</span><span class="n">img_size</span><span class="p">,</span>
            <span class="n">crop_size</span><span class="o">=</span><span class="n">crop_size</span><span class="p">,</span>
            <span class="n">augment</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
            <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">,</span>
            <span class="n">image_mask_transforms</span><span class="o">=</span><span class="n">image_mask_transforms</span><span class="p">,</span>
            <span class="n">image_mask_transforms_aug</span><span class="o">=</span><span class="n">image_mask_transforms_aug</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">CityscapesDataset</span><span class="p">(</span>
            <span class="n">root_dir</span><span class="o">=</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;lists/val.lst&#39;</span><span class="p">,</span>
            <span class="n">labels_csv_path</span><span class="o">=</span><span class="s2">&quot;lists/labels.csv&quot;</span><span class="p">,</span>
            <span class="n">img_size</span><span class="o">=</span><span class="n">img_size</span><span class="p">,</span>
            <span class="n">crop_size</span><span class="o">=</span><span class="n">crop_size</span><span class="p">,</span>
            <span class="n">augment</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
            <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">,</span>
            <span class="n">image_mask_transforms</span><span class="o">=</span><span class="n">image_mask_transforms</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span></div>


<div class="viewcode-block" id="SuperviselyPersonsDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.SuperviselyPersonsDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">SuperviselyPersonsDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cache_labels</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">cache_images</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>
        <span class="n">root_dir</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;dataset_dir&quot;</span><span class="p">,</span> <span class="s2">&quot;/data/supervisely-persons&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">SuperviselyPersonsDataset</span><span class="p">(</span>
            <span class="n">root_dir</span><span class="o">=</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;train.csv&#39;</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
            <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">,</span>
            <span class="n">image_mask_transforms_aug</span><span class="o">=</span><span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;image_mask_transforms_aug&quot;</span><span class="p">,</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([])),</span>
            <span class="n">augment</span><span class="o">=</span><span class="kc">True</span>
        <span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">SuperviselyPersonsDataset</span><span class="p">(</span>
            <span class="n">root_dir</span><span class="o">=</span><span class="n">root_dir</span><span class="p">,</span>
            <span class="n">list_file</span><span class="o">=</span><span class="s1">&#39;val.csv&#39;</span><span class="p">,</span>
            <span class="n">dataset_hyper_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">,</span>
            <span class="n">cache_labels</span><span class="o">=</span><span class="n">cache_labels</span><span class="p">,</span>
            <span class="n">cache_images</span><span class="o">=</span><span class="n">cache_images</span><span class="p">,</span>
            <span class="n">image_mask_transforms</span><span class="o">=</span><span class="n">get_param</span><span class="p">(</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;image_mask_transforms&quot;</span><span class="p">,</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([])),</span>
            <span class="n">augment</span><span class="o">=</span><span class="kc">False</span>
        <span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span></div>


<div class="viewcode-block" id="DetectionDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DetectionDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">DetectionDatasetInterface</span><span class="p">(</span><span class="n">DatasetInterface</span><span class="p">):</span>
<div class="viewcode-block" id="DetectionDatasetInterface.build_data_loaders"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.DetectionDatasetInterface.build_data_loaders">[docs]</a>    <span class="k">def</span> <span class="nf">build_data_loaders</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size_factor</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span> <span class="n">train_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">val_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                           <span class="n">test_batch_size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">distributed_sampler</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>

        <span class="n">train_sampler</span> <span class="o">=</span> <span class="n">InfiniteSampler</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">),</span> <span class="n">seed</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>

        <span class="n">train_batch_sampler</span> <span class="o">=</span> <span class="n">BatchSampler</span><span class="p">(</span>
            <span class="n">sampler</span><span class="o">=</span><span class="n">train_sampler</span><span class="p">,</span>
            <span class="n">batch_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">batch_size</span><span class="p">,</span>
            <span class="n">drop_last</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span> <span class="o">=</span> <span class="n">DataLoader</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="p">,</span>
                                       <span class="n">batch_sampler</span><span class="o">=</span><span class="n">train_batch_sampler</span><span class="p">,</span>
                                       <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                       <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                       <span class="n">worker_init_fn</span><span class="o">=</span><span class="n">worker_init_reset_seed</span><span class="p">,</span>
                                       <span class="n">collate_fn</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_collate_fn</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">distributed_sampler</span><span class="p">:</span>
            <span class="n">sampler</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">distributed</span><span class="o">.</span><span class="n">DistributedSampler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">,</span> <span class="n">shuffle</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">sampler</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">SequentialSampler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">)</span>

        <span class="n">val_loader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valset</span><span class="p">,</span>
                                                 <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span>
                                                 <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                                                 <span class="n">sampler</span><span class="o">=</span><span class="n">sampler</span><span class="p">,</span>
                                                 <span class="n">batch_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_batch_size</span><span class="p">,</span>
                                                 <span class="n">collate_fn</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_collate_fn</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">val_loader</span> <span class="o">=</span> <span class="n">val_loader</span></div></div>


<div class="viewcode-block" id="PascalVOCUnifiedDetectionDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.PascalVOCUnifiedDetectionDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">PascalVOCUnifiedDetectionDatasetInterface</span><span class="p">(</span><span class="n">DetectionDatasetInterface</span><span class="p">):</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">dataset_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">dataset_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">data_dir</span>
        <span class="n">train_input_dim</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_image_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_image_size</span><span class="p">)</span>
        <span class="n">val_input_dim</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_image_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_image_size</span><span class="p">)</span>
        <span class="n">train_max_num_samples</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;train_max_num_samples&quot;</span><span class="p">)</span>
        <span class="n">val_max_num_samples</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;val_max_num_samples&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">download</span><span class="p">:</span>
            <span class="n">PascalVOCDetectionDataset</span><span class="o">.</span><span class="n">download</span><span class="p">(</span><span class="n">data_dir</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="p">)</span>

        <span class="n">train_dataset_names</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;train2007&quot;</span><span class="p">,</span> <span class="s2">&quot;val2007&quot;</span><span class="p">,</span> <span class="s2">&quot;train2012&quot;</span><span class="p">,</span> <span class="s2">&quot;val2012&quot;</span><span class="p">]</span>
        <span class="c1"># We divide train_max_num_samples between the datasets</span>
        <span class="k">if</span> <span class="n">train_max_num_samples</span><span class="p">:</span>
            <span class="n">max_num_samples_per_train_dataset</span> <span class="o">=</span> <span class="p">[</span><span class="nb">len</span><span class="p">(</span><span class="n">segment</span><span class="p">)</span> <span class="k">for</span> <span class="n">segment</span> <span class="ow">in</span> <span class="n">np</span><span class="o">.</span><span class="n">array_split</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">train_max_num_samples</span><span class="p">),</span> <span class="nb">len</span><span class="p">(</span><span class="n">train_dataset_names</span><span class="p">))]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">max_num_samples_per_train_dataset</span> <span class="o">=</span> <span class="p">[</span><span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">train_dataset_names</span><span class="p">)</span>
        <span class="n">train_sets</span> <span class="o">=</span> <span class="p">[</span><span class="n">PascalVOCDetectionDataset</span><span class="p">(</span><span class="n">data_dir</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="p">,</span>
                                                <span class="n">input_dim</span><span class="o">=</span><span class="n">train_input_dim</span><span class="p">,</span>
                                                <span class="n">cache</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_train_images</span><span class="p">,</span>
                                                <span class="n">cache_path</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_dir</span> <span class="o">+</span> <span class="s2">&quot;cache_train&quot;</span><span class="p">,</span>
                                                <span class="n">transforms</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_transforms</span><span class="p">,</span>
                                                <span class="n">images_sub_directory</span><span class="o">=</span><span class="s1">&#39;images/&#39;</span> <span class="o">+</span> <span class="n">trainset_name</span> <span class="o">+</span> <span class="s1">&#39;/&#39;</span><span class="p">,</span>
                                                <span class="n">class_inclusion_list</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">class_inclusion_list</span><span class="p">,</span>
                                                <span class="n">max_num_samples</span><span class="o">=</span><span class="n">max_num_samples_per_train_dataset</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
                      <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">trainset_name</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">train_dataset_names</span><span class="p">)]</span>

        <span class="n">testset2007</span> <span class="o">=</span> <span class="n">PascalVOCDetectionDataset</span><span class="p">(</span><span class="n">data_dir</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="p">,</span>
                                                <span class="n">input_dim</span><span class="o">=</span><span class="n">val_input_dim</span><span class="p">,</span>
                                                <span class="n">cache</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_val_images</span><span class="p">,</span>
                                                <span class="n">cache_path</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_dir</span> <span class="o">+</span> <span class="s2">&quot;cache_valid&quot;</span><span class="p">,</span>
                                                <span class="n">transforms</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_transforms</span><span class="p">,</span>
                                                <span class="n">images_sub_directory</span><span class="o">=</span><span class="s1">&#39;images/test2007/&#39;</span><span class="p">,</span>
                                                <span class="n">class_inclusion_list</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">class_inclusion_list</span><span class="p">,</span>
                                                <span class="n">max_num_samples</span><span class="o">=</span><span class="n">val_max_num_samples</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">train_sets</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">classes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">ConcatDataset</span><span class="p">(</span><span class="n">train_sets</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">testset2007</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">collate_fn</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_collate_fn</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">img_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_image_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span><span class="o">.</span><span class="n">cache_labels</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_train_images</span></div>


<div class="viewcode-block" id="CoCoDetectionDatasetInterface"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.dataset_interfaces.html#super_gradients.training.CoCoDetectionDatasetInterface">[docs]</a><span class="k">class</span> <span class="nc">CoCoDetectionDatasetInterface</span><span class="p">(</span><span class="n">DetectionDatasetInterface</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="p">{}):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">CoCoDetectionDatasetInterface</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><span class="n">dataset_params</span><span class="o">=</span><span class="n">dataset_params</span><span class="p">)</span>

        <span class="n">train_input_dim</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_image_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_image_size</span><span class="p">)</span>
        <span class="n">targets_format</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s2">&quot;targets_format&quot;</span><span class="p">,</span> <span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">LABEL_CXCYWH</span><span class="p">)</span>

        <span class="n">train_transforms</span> <span class="o">=</span> <span class="p">[</span><span class="n">DetectionMosaic</span><span class="p">(</span><span class="n">input_dim</span><span class="o">=</span><span class="n">train_input_dim</span><span class="p">,</span>
                                            <span class="n">prob</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">mosaic_prob</span><span class="p">),</span>
                            <span class="n">DetectionRandomAffine</span><span class="p">(</span><span class="n">degrees</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">degrees</span><span class="p">,</span>
                                                  <span class="n">translate</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">translate</span><span class="p">,</span>
                                                  <span class="n">scales</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">mosaic_scale</span><span class="p">,</span>
                                                  <span class="n">shear</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">shear</span><span class="p">,</span>
                                                  <span class="n">target_size</span><span class="o">=</span><span class="n">train_input_dim</span><span class="p">,</span>
                                                  <span class="n">filter_box_candidates</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">filter_box_candidates</span><span class="p">,</span>
                                                  <span class="n">wh_thr</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">wh_thr</span><span class="p">,</span>
                                                  <span class="n">area_thr</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">area_thr</span><span class="p">,</span>
                                                  <span class="n">ar_thr</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">ar_thr</span>
                                                  <span class="p">),</span>
                            <span class="n">DetectionMixup</span><span class="p">(</span><span class="n">input_dim</span><span class="o">=</span><span class="n">train_input_dim</span><span class="p">,</span>
                                           <span class="n">mixup_scale</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">mixup_scale</span><span class="p">,</span>
                                           <span class="n">prob</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">mixup_prob</span><span class="p">,</span>
                                           <span class="n">flip_prob</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">flip_prob</span><span class="p">),</span>
                            <span class="n">DetectionHSV</span><span class="p">(</span><span class="n">prob</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">hsv_prob</span><span class="p">,</span>
                                         <span class="n">hgain</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">hgain</span><span class="p">,</span>
                                         <span class="n">sgain</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">sgain</span><span class="p">,</span>
                                         <span class="n">vgain</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">vgain</span>
                                         <span class="p">),</span>
                            <span class="n">DetectionHorizontalFlip</span><span class="p">(</span><span class="n">prob</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">flip_prob</span><span class="p">),</span>
                            <span class="n">DetectionPaddedRescale</span><span class="p">(</span><span class="n">input_dim</span><span class="o">=</span><span class="n">train_input_dim</span><span class="p">,</span> <span class="n">max_targets</span><span class="o">=</span><span class="mi">120</span><span class="p">),</span>
                            <span class="n">DetectionTargetsFormatTransform</span><span class="p">(</span><span class="n">output_format</span><span class="o">=</span><span class="n">targets_format</span><span class="p">)</span>
                            <span class="p">]</span>

        <span class="c1"># IF CACHE- CREATING THE CACHE FILE WILL HAPPEN ONLY FOR RANK 0, THEN ALL THE OTHER RANKS SIMPLY READ FROM IT.</span>
        <span class="n">local_rank</span> <span class="o">=</span> <span class="n">get_local_rank</span><span class="p">()</span>
        <span class="k">with</span> <span class="n">wait_for_the_master</span><span class="p">(</span><span class="n">local_rank</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">trainset</span> <span class="o">=</span> <span class="n">COCODetectionDataset</span><span class="p">(</span><span class="n">data_dir</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">data_dir</span><span class="p">,</span>
                                                 <span class="n">name</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_subdir</span><span class="p">,</span>
                                                 <span class="n">json_file</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_json_file</span><span class="p">,</span>
                                                 <span class="n">img_size</span><span class="o">=</span><span class="n">train_input_dim</span><span class="p">,</span>
                                                 <span class="n">cache</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_train_images</span><span class="p">,</span>
                                                 <span class="n">cache_dir_path</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_dir_path</span><span class="p">,</span>
                                                 <span class="n">transforms</span><span class="o">=</span><span class="n">train_transforms</span><span class="p">,</span>
                                                 <span class="n">with_crowd</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

        <span class="n">val_input_dim</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_image_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_image_size</span><span class="p">)</span>
        <span class="n">with_crowd</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="s1">&#39;with_crowd&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

        <span class="c1"># IF CACHE- CREATING THE CACHE FILE WILL HAPPEN ONLY FOR RANK 0, THEN ALL THE OTHER RANKS SIMPLY READ FROM IT.</span>
        <span class="k">with</span> <span class="n">wait_for_the_master</span><span class="p">(</span><span class="n">local_rank</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">valset</span> <span class="o">=</span> <span class="n">COCODetectionDataset</span><span class="p">(</span>
                <span class="n">data_dir</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">data_dir</span><span class="p">,</span>
                <span class="n">json_file</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_json_file</span><span class="p">,</span>
                <span class="n">name</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_subdir</span><span class="p">,</span>
                <span class="n">img_size</span><span class="o">=</span><span class="n">val_input_dim</span><span class="p">,</span>
                <span class="n">transforms</span><span class="o">=</span><span class="p">[</span><span class="n">DetectionPaddedRescale</span><span class="p">(</span><span class="n">input_dim</span><span class="o">=</span><span class="n">val_input_dim</span><span class="p">),</span>
                            <span class="n">DetectionTargetsFormatTransform</span><span class="p">(</span><span class="n">max_targets</span><span class="o">=</span><span class="mi">50</span><span class="p">,</span> <span class="n">output_format</span><span class="o">=</span><span class="n">targets_format</span><span class="p">)],</span>
                <span class="n">cache</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_val_images</span><span class="p">,</span>
                <span class="n">cache_dir_path</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">cache_dir_path</span><span class="p">,</span>
                <span class="n">with_crowd</span><span class="o">=</span><span class="n">with_crowd</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">COCO_DETECTION_CLASSES_LIST</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.datasets_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.datasets_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.datasets_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">copy</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">from</span> <span class="nn">abc</span> <span class="kn">import</span> <span class="n">ABC</span><span class="p">,</span> <span class="n">abstractmethod</span>
<span class="kn">from</span> <span class="nn">multiprocessing</span> <span class="kn">import</span> <span class="n">Value</span><span class="p">,</span> <span class="n">Lock</span>
<span class="kn">import</span> <span class="nn">random</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">torch.nn.functional</span> <span class="k">as</span> <span class="nn">F</span>
<span class="kn">import</span> <span class="nn">torchvision</span>
<span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">torch.distributed</span> <span class="k">as</span> <span class="nn">dist</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.sg_loggers.abstract_sg_logger</span> <span class="kn">import</span> <span class="n">AbstractSGLogger</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
<span class="kn">from</span> <span class="nn">deprecated</span> <span class="kn">import</span> <span class="n">deprecated</span>
<span class="kn">from</span> <span class="nn">matplotlib.patches</span> <span class="kn">import</span> <span class="n">Rectangle</span>
<span class="kn">from</span> <span class="nn">torchvision.datasets</span> <span class="kn">import</span> <span class="n">ImageFolder</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.auto_augment</span> <span class="kn">import</span> <span class="n">rand_augment_transform</span>
<span class="kn">from</span> <span class="nn">torchvision.transforms</span> <span class="kn">import</span> <span class="n">transforms</span><span class="p">,</span> <span class="n">InterpolationMode</span><span class="p">,</span> <span class="n">RandomResizedCrop</span>
<span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">AverageMeter</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">DetectionVisualization</span><span class="p">,</span> <span class="n">Anchors</span>
<span class="kn">import</span> <span class="nn">uuid</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.distributed_training_utils</span> <span class="kn">import</span> <span class="n">get_local_rank</span><span class="p">,</span> <span class="n">get_world_size</span>

<span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="nn">plt</span>


<div class="viewcode-block" id="get_mean_and_std_torch"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.get_mean_and_std_torch">[docs]</a><span class="k">def</span> <span class="nf">get_mean_and_std_torch</span><span class="p">(</span><span class="n">data_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dataloader</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span> <span class="n">RandomResizeSize</span><span class="o">=</span><span class="mi">224</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A function for getting the mean and std of large datasets using pytorch dataloader and gpu functionality.</span>

<span class="sd">    :param data_dir: String, path to none-library dataset folder. For example &quot;/data/Imagenette&quot; or &quot;/data/TinyImagenet&quot;</span>
<span class="sd">    :param dataloader: a torch DataLoader, as it would feed the data into the trainer (including transforms etc).</span>
<span class="sd">    :param RandomResizeSize: Int, the size of the RandomResizeCrop as it appears in the DataInterface (for example, for Imagenet,</span>
<span class="sd">    this value should be 224).</span>
<span class="sd">    :return: 2 lists,mean and std, each one of len 3 (1 for each channel)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">assert</span> <span class="n">data_dir</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">dataloader</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">,</span> <span class="s1">&#39;Please provide either path to data folder or DataLoader, not both.&#39;</span>

    <span class="k">if</span> <span class="n">dataloader</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">traindir</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">abspath</span><span class="p">(</span><span class="n">data_dir</span><span class="p">),</span> <span class="s1">&#39;train&#39;</span><span class="p">)</span>
        <span class="n">trainset</span> <span class="o">=</span> <span class="n">ImageFolder</span><span class="p">(</span><span class="n">traindir</span><span class="p">,</span> <span class="n">transforms</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span><span class="n">transforms</span><span class="o">.</span><span class="n">RandomResizedCrop</span><span class="p">(</span><span class="n">RandomResizeSize</span><span class="p">),</span>
                                                             <span class="n">transforms</span><span class="o">.</span><span class="n">RandomHorizontalFlip</span><span class="p">(),</span>
                                                             <span class="n">transforms</span><span class="o">.</span><span class="n">ToTensor</span><span class="p">()]))</span>
        <span class="n">dataloader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="n">trainset</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">)</span>

    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Calculating on </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">dataloader</span><span class="o">.</span><span class="n">dataset</span><span class="o">.</span><span class="n">targets</span><span class="p">)</span><span class="si">}</span><span class="s1"> Training Samples&#39;</span><span class="p">)</span>

    <span class="n">device</span> <span class="o">=</span> <span class="s1">&#39;cuda:0&#39;</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="s1">&#39;cpu&#39;</span>
    <span class="n">h</span><span class="p">,</span> <span class="n">w</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span>
    <span class="k">for</span> <span class="n">batch_idx</span><span class="p">,</span> <span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">dataloader</span><span class="p">):</span>
        <span class="n">inputs</span> <span class="o">=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">batch_idx</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">h</span><span class="p">,</span> <span class="n">w</span> <span class="o">=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">2</span><span class="p">),</span> <span class="n">inputs</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
            <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Min: </span><span class="si">{</span><span class="n">inputs</span><span class="o">.</span><span class="n">min</span><span class="p">()</span><span class="si">}</span><span class="s1">, Max: </span><span class="si">{</span><span class="n">inputs</span><span class="o">.</span><span class="n">max</span><span class="p">()</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
            <span class="n">chsum</span> <span class="o">=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">chsum</span> <span class="o">+=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="n">mean</span> <span class="o">=</span> <span class="n">chsum</span> <span class="o">/</span> <span class="nb">len</span><span class="p">(</span><span class="n">trainset</span><span class="p">)</span> <span class="o">/</span> <span class="n">h</span> <span class="o">/</span> <span class="n">w</span>
    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;mean: </span><span class="si">{</span><span class="n">mean</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>

    <span class="n">chsum</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">for</span> <span class="n">batch_idx</span><span class="p">,</span> <span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">dataloader</span><span class="p">):</span>
        <span class="n">inputs</span> <span class="o">=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">batch_idx</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">chsum</span> <span class="o">=</span> <span class="p">(</span><span class="n">inputs</span> <span class="o">-</span> <span class="n">mean</span><span class="p">)</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">chsum</span> <span class="o">+=</span> <span class="p">(</span><span class="n">inputs</span> <span class="o">-</span> <span class="n">mean</span><span class="p">)</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="n">std</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">chsum</span> <span class="o">/</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">trainset</span><span class="p">)</span> <span class="o">*</span> <span class="n">h</span> <span class="o">*</span> <span class="n">w</span> <span class="o">-</span> <span class="mi">1</span><span class="p">))</span>
    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;std: </span><span class="si">{</span><span class="n">std</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">mean</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</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><span class="o">.</span><span class="n">tolist</span><span class="p">(),</span> <span class="n">std</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</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><span class="o">.</span><span class="n">tolist</span><span class="p">()</span></div>


<div class="viewcode-block" id="get_mean_and_std"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.get_mean_and_std">[docs]</a><span class="nd">@deprecated</span><span class="p">(</span><span class="n">reason</span><span class="o">=</span><span class="s1">&#39;Use get_mean_and_std_torch() instead. It is faster and more accurate&#39;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">get_mean_and_std</span><span class="p">(</span><span class="n">dataset</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;Compute the mean and std value of dataset.&#39;&#39;&#39;</span>
    <span class="n">dataloader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">shuffle</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">mean</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
    <span class="n">std</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
    <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;==&gt; Computing mean and std..&#39;</span><span class="p">)</span>
    <span class="n">j</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">for</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span> <span class="ow">in</span> <span class="n">dataloader</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">j</span> <span class="o">%</span> <span class="mi">10</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="n">j</span><span class="p">)</span>
        <span class="n">j</span> <span class="o">+=</span> <span class="mi">1</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">3</span><span class="p">):</span>
            <span class="n">mean</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">inputs</span><span class="p">[:,</span> <span class="n">i</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:]</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>
            <span class="n">std</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">inputs</span><span class="p">[:,</span> <span class="n">i</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:]</span><span class="o">.</span><span class="n">std</span><span class="p">()</span>
    <span class="n">mean</span><span class="o">.</span><span class="n">div_</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">dataset</span><span class="p">))</span>
    <span class="n">std</span><span class="o">.</span><span class="n">div_</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">dataset</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">mean</span><span class="p">,</span> <span class="n">std</span></div>


<div class="viewcode-block" id="AbstractCollateFunction"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.AbstractCollateFunction">[docs]</a><span class="k">class</span> <span class="nc">AbstractCollateFunction</span><span class="p">(</span><span class="n">ABC</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A collate function (for torch DataLoader)</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@abstractmethod</span>
    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">):</span>
        <span class="k">pass</span></div>


<div class="viewcode-block" id="ComposedCollateFunction"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.ComposedCollateFunction">[docs]</a><span class="k">class</span> <span class="nc">ComposedCollateFunction</span><span class="p">(</span><span class="n">AbstractCollateFunction</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A function (for torch DataLoader) which executes a sequence of sub collate functions</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">functions</span><span class="p">:</span> <span class="nb">list</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">functions</span> <span class="o">=</span> <span class="n">functions</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">f</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">functions</span><span class="p">:</span>
            <span class="n">batch</span> <span class="o">=</span> <span class="n">f</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">batch</span></div>


<div class="viewcode-block" id="AtomicInteger"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.AtomicInteger">[docs]</a><span class="k">class</span> <span class="nc">AtomicInteger</span><span class="p">:</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_value</span> <span class="o">=</span> <span class="n">Value</span><span class="p">(</span><span class="s1">&#39;i&#39;</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__set__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">instance</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_value</span><span class="o">.</span><span class="n">value</span> <span class="o">=</span> <span class="n">value</span>

    <span class="k">def</span> <span class="fm">__get__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">instance</span><span class="p">,</span> <span class="n">owner</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_value</span><span class="o">.</span><span class="n">value</span></div>


<div class="viewcode-block" id="MultiScaleCollateFunction"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.MultiScaleCollateFunction">[docs]</a><span class="k">class</span> <span class="nc">MultiScaleCollateFunction</span><span class="p">(</span><span class="n">AbstractCollateFunction</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    a collate function to implement multi-scale data augmentation</span>
<span class="sd">    according to https://arxiv.org/pdf/1612.08242.pdf</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_counter</span> <span class="o">=</span> <span class="n">AtomicInteger</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
    <span class="n">_current_size</span> <span class="o">=</span> <span class="n">AtomicInteger</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
    <span class="n">_lock</span> <span class="o">=</span> <span class="n">Lock</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">target_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">min_image_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">max_image_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                 <span class="n">image_size_steps</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">32</span><span class="p">,</span>
                 <span class="n">change_frequency</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">10</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        set parameters for the multi-scale collate function</span>
<span class="sd">        the possible image sizes are in range [min_image_size, max_image_size] in steps of image_size_steps</span>
<span class="sd">        a new size will be randomly selected every change_frequency calls to the collate_fn()</span>
<span class="sd">            :param target_size: scales will be [0.66 * target_size, 1.5 * target_size]</span>
<span class="sd">            :param min_image_size: the minimum size to scale down to (in pixels)</span>
<span class="sd">            :param max_image_size: the maximum size to scale up to (in pixels)</span>
<span class="sd">            :param image_size_steps: typically, the stride of the net, which defines the possible image</span>
<span class="sd">                    size multiplications</span>
<span class="sd">            :param change_frequency:</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">assert</span> <span class="n">target_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="p">(</span><span class="n">max_image_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">min_image_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">),</span> \
            <span class="s1">&#39;either target_size or min_image_size and max_image_size has to be set&#39;</span>
        <span class="k">assert</span> <span class="n">target_size</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">max_image_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">,</span> <span class="s1">&#39;target_size and max_image_size cannot be both defined&#39;</span>

        <span class="k">if</span> <span class="n">target_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">min_image_size</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="mf">0.66</span> <span class="o">*</span> <span class="n">target_size</span> <span class="o">-</span> <span class="p">((</span><span class="mf">0.66</span> <span class="o">*</span> <span class="n">target_size</span><span class="p">)</span> <span class="o">%</span> <span class="n">image_size_steps</span><span class="p">)</span> <span class="o">+</span> <span class="n">image_size_steps</span><span class="p">)</span>
            <span class="n">max_image_size</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="mf">1.5</span> <span class="o">*</span> <span class="n">target_size</span> <span class="o">-</span> <span class="p">((</span><span class="mf">1.5</span> <span class="o">*</span> <span class="n">target_size</span><span class="p">)</span> <span class="o">%</span> <span class="n">image_size_steps</span><span class="p">))</span>

        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Using multi-scale </span><span class="si">%g</span><span class="s1"> - </span><span class="si">%g</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="n">min_image_size</span><span class="p">,</span> <span class="n">max_image_size</span><span class="p">))</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">sizes</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="n">min_image_size</span><span class="p">,</span> <span class="n">max_image_size</span> <span class="o">+</span> <span class="n">image_size_steps</span><span class="p">,</span> <span class="n">image_size_steps</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span> <span class="o">=</span> <span class="n">image_size_steps</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">frequency</span> <span class="o">=</span> <span class="n">change_frequency</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_current_size</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">sizes</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">):</span>

        <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">_lock</span><span class="p">:</span>

            <span class="c1"># Important: this implementation was tailored for a specific input. it assumes the batch is a tuple where</span>
            <span class="c1"># the images are the first item</span>
            <span class="k">assert</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">batch</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">),</span> <span class="s1">&#39;this collate function expects the input to be a tuple (images, labels)&#39;</span>
            <span class="n">images</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_counter</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">frequency</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">_current_size</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">sizes</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_counter</span> <span class="o">+=</span> <span class="mi">1</span>

            <span class="k">assert</span> <span class="n">images</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span> <span class="o">==</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">images</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> \
                <span class="s1">&#39;images sized not divisible by </span><span class="si">%d</span><span class="s1">. (resize images before calling multi_scale)&#39;</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span>

            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_current_size</span> <span class="o">!=</span> <span class="nb">max</span><span class="p">(</span><span class="n">images</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">:]):</span>
                <span class="n">ratio</span> <span class="o">=</span> <span class="nb">float</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_current_size</span><span class="p">)</span> <span class="o">/</span> <span class="nb">max</span><span class="p">(</span><span class="n">images</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">:])</span>
                <span class="n">new_size</span> <span class="o">=</span> <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">images</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="n">ratio</span><span class="p">)),</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">images</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">*</span> <span class="n">ratio</span><span class="p">)))</span>
                <span class="n">images</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">interpolate</span><span class="p">(</span><span class="n">images</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="n">new_size</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s1">&#39;bilinear&#39;</span><span class="p">,</span> <span class="n">align_corners</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

            <span class="k">return</span> <span class="n">images</span><span class="p">,</span> <span class="n">batch</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span></div>


<div class="viewcode-block" id="AbstractPrePredictionCallback"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.AbstractPrePredictionCallback">[docs]</a><span class="k">class</span> <span class="nc">AbstractPrePredictionCallback</span><span class="p">(</span><span class="n">ABC</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Abstract class for forward pass preprocessing function, to be used by passing its inheritors through training_params</span>
<span class="sd">     pre_prediction_callback keyword arg.</span>

<span class="sd">    Should implement __call__ and return images, targets after applying the desired preprocessing.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="nd">@abstractmethod</span>
    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">):</span>
        <span class="k">pass</span></div>


<div class="viewcode-block" id="MultiscalePrePredictionCallback"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.MultiscalePrePredictionCallback">[docs]</a><span class="k">class</span> <span class="nc">MultiscalePrePredictionCallback</span><span class="p">(</span><span class="n">AbstractPrePredictionCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Mutiscale pre-prediction callback pass function.</span>

<span class="sd">    When passed through train_params images, targets will be applied by the below transform to support multi scaling</span>
<span class="sd">    on the fly.</span>

<span class="sd">    After each self.frequency forward passes, change size randomly from</span>
<span class="sd">     (input_size-self.multiscale_range*self.image_size_steps, input_size-(self.multiscale_range-1)*self.image_size_steps,</span>
<span class="sd">     ...input_size+self.multiscale_range*self.image_size_steps)</span>


<span class="sd">    Attributes:</span>
<span class="sd">        multiscale_range: (int) Range of values for resize sizes as discussed above (default=5)</span>
<span class="sd">        image_size_steps: (int) Image step sizes as discussed abov (default=32)</span>
<span class="sd">        change_frequency: (int) The frequency to apply change in input size.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">multiscale_range</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">5</span><span class="p">,</span>
                 <span class="n">image_size_steps</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">32</span><span class="p">,</span>
                 <span class="n">change_frequency</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">10</span><span class="p">):</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">multiscale_range</span> <span class="o">=</span> <span class="n">multiscale_range</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span> <span class="o">=</span> <span class="n">image_size_steps</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">frequency</span> <span class="o">=</span> <span class="n">change_frequency</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">is_distributed</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sampled_imres_once</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">new_input_size</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="o">=</span> <span class="n">get_local_rank</span><span class="p">()</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_distributed</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">is_distributed</span> <span class="o">=</span> <span class="n">get_world_size</span><span class="p">()</span> <span class="o">&gt;</span> <span class="mi">1</span>

        <span class="c1"># GENERATE A NEW SIZE AND BROADCAST IT TO THE THE OTHER RANKS SO THEY HAVE THE SAME SCALE</span>
        <span class="n">input_size</span> <span class="o">=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">:]</span>
        <span class="k">if</span> <span class="n">batch_idx</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">frequency</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">tensor</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">LongTensor</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">inputs</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>

            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">size_factor</span> <span class="o">=</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="mf">1.0</span> <span class="o">/</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">min_size</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span><span class="p">)</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">multiscale_range</span>
                <span class="n">max_size</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">multiscale_range</span>
                <span class="n">random_size</span> <span class="o">=</span> <span class="p">(</span><span class="n">min_size</span><span class="p">,</span> <span class="n">max_size</span><span class="p">)</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">sampled_imres_once</span><span class="p">:</span>
                    <span class="n">size</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="o">*</span><span class="n">random_size</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="c1"># sample the biggest resolution first to make sure the run fits into the GPU memory</span>
                    <span class="n">size</span> <span class="o">=</span> <span class="n">max_size</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">sampled_imres_once</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="n">size</span> <span class="o">=</span> <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span> <span class="o">*</span> <span class="n">size</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_size_steps</span> <span class="o">*</span> <span class="nb">int</span><span class="p">(</span><span class="n">size</span> <span class="o">*</span> <span class="n">size_factor</span><span class="p">))</span>
                <span class="n">tensor</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">tensor</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>

            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_distributed</span><span class="p">:</span>
                <span class="n">dist</span><span class="o">.</span><span class="n">barrier</span><span class="p">()</span>
                <span class="n">dist</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="n">tensor</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">new_input_size</span> <span class="o">=</span> <span class="p">(</span><span class="n">tensor</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">item</span><span class="p">(),</span> <span class="n">tensor</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">item</span><span class="p">())</span>

        <span class="n">scale_y</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">new_input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">scale_x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">new_input_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">scale_x</span> <span class="o">!=</span> <span class="mi">1</span> <span class="ow">or</span> <span class="n">scale_y</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">inputs</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">functional</span><span class="o">.</span><span class="n">interpolate</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">new_input_size</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;bilinear&quot;</span><span class="p">,</span> <span class="n">align_corners</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span></div>


<div class="viewcode-block" id="DetectionMultiscalePrePredictionCallback"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.DetectionMultiscalePrePredictionCallback">[docs]</a><span class="k">class</span> <span class="nc">DetectionMultiscalePrePredictionCallback</span><span class="p">(</span><span class="n">MultiscalePrePredictionCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Mutiscalepre-prediction callback for object detection.</span>

<span class="sd">    When passed through train_params images, targets will be applied by the below transform to support multi scaling</span>
<span class="sd">    on the fly.</span>

<span class="sd">    After each self.frequency forward passes, change size randomly from</span>
<span class="sd">     (input_size-self.multiscale_range*self.image_size_steps, input_size-(self.multiscale_range-1)*self.image_size_steps,</span>
<span class="sd">     ...input_size+self.multiscale_range*self.image_size_steps) and apply the same rescaling to the box coordinates.</span>



<span class="sd">    Attributes:</span>
<span class="sd">        multiscale_range: (int) Range of values for resize sizes as discussed above (default=5)</span>
<span class="sd">        image_size_steps: (int) Image step sizes as discussed abov (default=32)</span>
<span class="sd">        change_frequency: (int) The frequency to apply change in input size.</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">):</span>
        <span class="c1"># RESCALE THE IMAGE FIRST WITH SUPER(), AND IF RESCALING HAS ACTUALLY BEEN DONE APPLY TO BOXES AS WELL</span>
        <span class="n">input_size</span> <span class="o">=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">:]</span>
        <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span> <span class="o">=</span> <span class="nb">super</span><span class="p">(</span><span class="n">DetectionMultiscalePrePredictionCallback</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__call__</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">)</span>
        <span class="n">new_input_size</span> <span class="o">=</span> <span class="n">inputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">:]</span>
        <span class="n">scale_y</span> <span class="o">=</span> <span class="n">new_input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">scale_x</span> <span class="o">=</span> <span class="n">new_input_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">scale_x</span> <span class="o">!=</span> <span class="mi">1</span> <span class="ow">or</span> <span class="n">scale_y</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">targets</span><span class="p">[</span><span class="o">...</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="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="o">...</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="o">*</span> <span class="n">scale_x</span>
            <span class="n">targets</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">3</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">3</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="n">scale_y</span>
        <span class="k">return</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span></div>


<span class="n">_pil_interpolation_to_str</span> <span class="o">=</span> <span class="p">{</span>
    <span class="n">Image</span><span class="o">.</span><span class="n">NEAREST</span><span class="p">:</span> <span class="s1">&#39;PIL.Image.NEAREST&#39;</span><span class="p">,</span>
    <span class="n">Image</span><span class="o">.</span><span class="n">BILINEAR</span><span class="p">:</span> <span class="s1">&#39;PIL.Image.BILINEAR&#39;</span><span class="p">,</span>
    <span class="n">Image</span><span class="o">.</span><span class="n">BICUBIC</span><span class="p">:</span> <span class="s1">&#39;PIL.Image.BICUBIC&#39;</span><span class="p">,</span>
    <span class="n">Image</span><span class="o">.</span><span class="n">LANCZOS</span><span class="p">:</span> <span class="s1">&#39;PIL.Image.LANCZOS&#39;</span><span class="p">,</span>
    <span class="n">Image</span><span class="o">.</span><span class="n">HAMMING</span><span class="p">:</span> <span class="s1">&#39;PIL.Image.HAMMING&#39;</span><span class="p">,</span>
    <span class="n">Image</span><span class="o">.</span><span class="n">BOX</span><span class="p">:</span> <span class="s1">&#39;PIL.Image.BOX&#39;</span><span class="p">,</span>
<span class="p">}</span>


<span class="k">def</span> <span class="nf">_pil_interp</span><span class="p">(</span><span class="n">method</span><span class="p">):</span>
    <span class="k">if</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;bicubic&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">BICUBIC</span>
    <span class="k">elif</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;lanczos&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">LANCZOS</span>
    <span class="k">elif</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;hamming&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">HAMMING</span>
    <span class="k">elif</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;nearest&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">NEAREST</span>
    <span class="k">elif</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;bilinear&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">BILINEAR</span>
    <span class="k">elif</span> <span class="n">method</span> <span class="o">==</span> <span class="s1">&#39;box&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">BOX</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;interpolation type must be one of [&#39;bilinear&#39;, &#39;bicubic&#39;, &#39;lanczos&#39;, &#39;hamming&#39;, &quot;</span>
                         <span class="s2">&quot;&#39;nearest&#39;, &#39;box&#39;] for explicit interpolation type, or &#39;random&#39; for random&quot;</span><span class="p">)</span>


<span class="n">_RANDOM_INTERPOLATION</span> <span class="o">=</span> <span class="p">(</span><span class="n">InterpolationMode</span><span class="o">.</span><span class="n">BILINEAR</span><span class="p">,</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">BICUBIC</span><span class="p">)</span>


<div class="viewcode-block" id="RandomResizedCropAndInterpolation"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.RandomResizedCropAndInterpolation">[docs]</a><span class="k">class</span> <span class="nc">RandomResizedCropAndInterpolation</span><span class="p">(</span><span class="n">RandomResizedCrop</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Crop the given PIL Image to random size and aspect ratio with explicitly chosen or random interpolation.</span>

<span class="sd">    A crop of random size (default: of 0.08 to 1.0) of the original size and a random</span>
<span class="sd">    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop</span>
<span class="sd">    is finally resized to given size.</span>
<span class="sd">    This is popularly used to train the Inception networks.</span>

<span class="sd">    Args:</span>
<span class="sd">        size: expected output size of each edge</span>
<span class="sd">        scale: range of size of the origin size cropped</span>
<span class="sd">        ratio: range of aspect ratio of the origin aspect ratio cropped</span>
<span class="sd">        interpolation: Default: PIL.Image.BILINEAR</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">scale</span><span class="o">=</span><span class="p">(</span><span class="mf">0.08</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">),</span> <span class="n">ratio</span><span class="o">=</span><span class="p">(</span><span class="mf">3.</span> <span class="o">/</span> <span class="mf">4.</span><span class="p">,</span> <span class="mf">4.</span> <span class="o">/</span> <span class="mf">3.</span><span class="p">),</span>
                 <span class="n">interpolation</span><span class="o">=</span><span class="s1">&#39;default&#39;</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">RandomResizedCropAndInterpolation</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><span class="n">size</span><span class="o">=</span><span class="n">size</span><span class="p">,</span> <span class="n">scale</span><span class="o">=</span><span class="n">scale</span><span class="p">,</span> <span class="n">ratio</span><span class="o">=</span><span class="n">ratio</span><span class="p">,</span> <span class="n">interpolation</span><span class="o">=</span><span class="n">interpolation</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">interpolation</span> <span class="o">==</span> <span class="s1">&#39;random&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span> <span class="o">=</span> <span class="n">_RANDOM_INTERPOLATION</span>
        <span class="k">elif</span> <span class="n">interpolation</span> <span class="o">==</span> <span class="s1">&#39;default&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span> <span class="o">=</span> <span class="n">InterpolationMode</span><span class="o">.</span><span class="n">BILINEAR</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span> <span class="o">=</span> <span class="n">_pil_interp</span><span class="p">(</span><span class="n">interpolation</span><span class="p">)</span>

<div class="viewcode-block" id="RandomResizedCropAndInterpolation.forward"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.RandomResizedCropAndInterpolation.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Args:</span>
<span class="sd">            img (PIL Image): Image to be cropped and resized.</span>

<span class="sd">        Returns:</span>
<span class="sd">            PIL Image: Randomly cropped and resized image.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_params</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">ratio</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
            <span class="n">interpolation</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">choice</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">interpolation</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span>
        <span class="k">return</span> <span class="n">torchvision</span><span class="o">.</span><span class="n">transforms</span><span class="o">.</span><span class="n">functional</span><span class="o">.</span><span class="n">resized_crop</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">size</span><span class="p">,</span> <span class="n">interpolation</span><span class="p">)</span></div>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
            <span class="n">interpolate_str</span> <span class="o">=</span> <span class="s1">&#39; &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">_pil_interpolation_to_str</span><span class="p">[</span><span class="n">x</span><span class="p">]</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span><span class="p">])</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">interpolate_str</span> <span class="o">=</span> <span class="n">_pil_interpolation_to_str</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">interpolation</span><span class="p">]</span>
        <span class="n">format_string</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span> <span class="o">+</span> <span class="s1">&#39;(size=</span><span class="si">{0}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">size</span><span class="p">)</span>
        <span class="n">format_string</span> <span class="o">+=</span> <span class="s1">&#39;, scale=</span><span class="si">{0}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="mi">4</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale</span><span class="p">))</span>
        <span class="n">format_string</span> <span class="o">+=</span> <span class="s1">&#39;, ratio=</span><span class="si">{0}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">r</span><span class="p">,</span> <span class="mi">4</span><span class="p">)</span> <span class="k">for</span> <span class="n">r</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">ratio</span><span class="p">))</span>
        <span class="n">format_string</span> <span class="o">+=</span> <span class="s1">&#39;, interpolation=</span><span class="si">{0}</span><span class="s1">)&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">interpolate_str</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">format_string</span></div>


<span class="n">STAT_LOGGER_FONT_SIZE</span> <span class="o">=</span> <span class="mi">15</span>


<div class="viewcode-block" id="DatasetStatisticsTensorboardLogger"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.DatasetStatisticsTensorboardLogger">[docs]</a><span class="k">class</span> <span class="nc">DatasetStatisticsTensorboardLogger</span><span class="p">:</span>

    <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
    <span class="n">DEFAULT_SUMMARY_PARAMS</span> <span class="o">=</span> <span class="p">{</span>
        <span class="s1">&#39;sample_images&#39;</span><span class="p">:</span> <span class="mi">32</span><span class="p">,</span>  <span class="c1"># by default, 32 images will be sampled from each dataset</span>
        <span class="s1">&#39;plot_class_distribution&#39;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
        <span class="s1">&#39;plot_box_size_distribution&#39;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
        <span class="s1">&#39;plot_anchors_coverage&#39;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
        <span class="s1">&#39;max_batches&#39;</span><span class="p">:</span> <span class="mi">30</span>
    <span class="p">}</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sg_logger</span><span class="p">:</span> <span class="n">AbstractSGLogger</span><span class="p">,</span> <span class="n">summary_params</span><span class="p">:</span> <span class="nb">dict</span> <span class="o">=</span> <span class="n">DEFAULT_SUMMARY_PARAMS</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span> <span class="o">=</span> <span class="n">sg_logger</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">summary_params</span> <span class="o">=</span> <span class="p">{</span><span class="o">**</span><span class="n">DatasetStatisticsTensorboardLogger</span><span class="o">.</span><span class="n">DEFAULT_SUMMARY_PARAMS</span><span class="p">,</span> <span class="o">**</span><span class="n">summary_params</span><span class="p">}</span>

<div class="viewcode-block" id="DatasetStatisticsTensorboardLogger.analyze"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.DatasetStatisticsTensorboardLogger.analyze">[docs]</a>    <span class="k">def</span> <span class="nf">analyze</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data_loader</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">,</span> <span class="n">title</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                <span class="n">all_classes</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="n">anchors</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param data_loader: the dataset data loader</span>
<span class="sd">        :param dataset_params: the dataset parameters</span>
<span class="sd">        :param title: the title for this dataset (i.e. Coco 2017 test set)</span>
<span class="sd">        :param anchors: the list of anchors used by the model. applicable only for detection datasets</span>
<span class="sd">        :param all_classes: the list of all classes names</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># FIXME: UNCOMMENT AND APPLY TO NEW DetectionDataSet ONCE ITS MERGED</span>
        <span class="c1"># if isinstance(data_loader.dataset, DetectionDataSet):</span>
        <span class="c1">#     self._analyze_detection(data_loader=data_loader, title=title,</span>
        <span class="c1">#                             all_classes=all_classes, anchors=anchors)</span>
        <span class="c1"># else:</span>
        <span class="c1">#     DatasetStatisticsTensorboardLogger.logger.warning(&#39;only DetectionDataSet are currently supported&#39;)</span>
        <span class="n">DatasetStatisticsTensorboardLogger</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s1">&#39;only DetectionDataSet are currently supported&#39;</span><span class="p">)</span></div>

    <span class="k">def</span> <span class="nf">_analyze_detection</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data_loader</span><span class="p">,</span> <span class="n">title</span><span class="p">,</span> <span class="n">all_classes</span><span class="p">,</span> <span class="n">anchors</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Analyze a detection dataset</span>

<span class="sd">        :param data_loader: the dataset data loader</span>
<span class="sd">        :param dataset_params: the dataset parameters</span>
<span class="sd">        :param all_classes: the list of all classes names</span>
<span class="sd">        :param title: the title for this dataset (i.e. Coco 2017 test set)</span>
<span class="sd">        :param anchors: the list of anchors used by the model. if not provided, anchors coverage will not be analyzed</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">color_mean</span> <span class="o">=</span> <span class="n">AverageMeter</span><span class="p">()</span>
            <span class="n">color_std</span> <span class="o">=</span> <span class="n">AverageMeter</span><span class="p">()</span>
            <span class="n">all_labels</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">image_size</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="p">(</span><span class="n">images</span><span class="p">,</span> <span class="n">labels</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">tqdm</span><span class="p">(</span><span class="n">data_loader</span><span class="p">)):</span>

                <span class="k">if</span> <span class="n">i</span> <span class="o">&gt;=</span> <span class="bp">self</span><span class="o">.</span><span class="n">summary_params</span><span class="p">[</span><span class="s1">&#39;max_batches&#39;</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="k">break</span>

                <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">image_size</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">images</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">images</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">])</span>
                    <span class="k">if</span> <span class="n">images</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">summary_params</span><span class="p">[</span><span class="s1">&#39;sample_images&#39;</span><span class="p">]:</span>
                        <span class="n">samples</span> <span class="o">=</span> <span class="n">images</span><span class="p">[:</span><span class="bp">self</span><span class="o">.</span><span class="n">summary_params</span><span class="p">[</span><span class="s1">&#39;sample_images&#39;</span><span class="p">]]</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="n">samples</span> <span class="o">=</span> <span class="n">images</span>

                    <span class="n">pred</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">6</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">samples</span><span class="p">))]</span>
                    <span class="k">try</span><span class="p">:</span>
                        <span class="n">result_images</span> <span class="o">=</span> <span class="n">DetectionVisualization</span><span class="o">.</span><span class="n">visualize_batch</span><span class="p">(</span><span class="n">image_tensor</span><span class="o">=</span><span class="n">samples</span><span class="p">,</span> <span class="n">pred_boxes</span><span class="o">=</span><span class="n">pred</span><span class="p">,</span>
                                                                               <span class="n">target_boxes</span><span class="o">=</span><span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">labels</span><span class="p">),</span>
                                                                               <span class="n">batch_name</span><span class="o">=</span><span class="n">title</span><span class="p">,</span>
                                                                               <span class="n">class_names</span><span class="o">=</span><span class="n">all_classes</span><span class="p">,</span>
                                                                               <span class="n">box_thickness</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                                                                               <span class="n">gt_alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">)</span>

                        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_images</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s1"> sample images&#39;</span><span class="p">,</span> <span class="n">images</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">result_images</span><span class="p">)</span>
                                                  <span class="o">.</span><span class="n">transpose</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">3</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="p">::</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:])</span>
                    <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
                        <span class="n">DatasetStatisticsTensorboardLogger</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span>
                            <span class="sa">f</span><span class="s1">&#39;Dataset Statistics failed at adding an example batch:</span><span class="se">\n</span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
                        <span class="k">return</span>

                <span class="n">all_labels</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">labels</span><span class="p">)</span>
                <span class="n">color_mean</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">images</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">]),</span> <span class="mi">1</span><span class="p">)</span>
                <span class="n">color_std</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">std</span><span class="p">(</span><span class="n">images</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">]),</span> <span class="mi">1</span><span class="p">)</span>

            <span class="n">all_labels</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">all_labels</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">:]</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>

            <span class="k">try</span><span class="p">:</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">summary_params</span><span class="p">[</span><span class="s1">&#39;plot_class_distribution&#39;</span><span class="p">]:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">_analyze_class_distribution</span><span class="p">(</span><span class="n">labels</span><span class="o">=</span><span class="n">all_labels</span><span class="p">,</span> <span class="n">num_classes</span><span class="o">=</span><span class="nb">len</span><span class="p">(</span><span class="n">all_classes</span><span class="p">),</span> <span class="n">title</span><span class="o">=</span><span class="n">title</span><span class="p">)</span>
            <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
                <span class="n">DatasetStatisticsTensorboardLogger</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Dataset Statistics failed at analyzing class distributions.</span><span class="se">\n</span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="k">return</span>

            <span class="k">try</span><span class="p">:</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">summary_params</span><span class="p">[</span><span class="s1">&#39;plot_box_size_distribution&#39;</span><span class="p">]:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">_analyze_object_size_distribution</span><span class="p">(</span><span class="n">labels</span><span class="o">=</span><span class="n">all_labels</span><span class="p">,</span> <span class="n">title</span><span class="o">=</span><span class="n">title</span><span class="p">)</span>
            <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
                <span class="n">DatasetStatisticsTensorboardLogger</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Dataset Statistics failed at analyzing object size &#39;</span>
                                                                <span class="sa">f</span><span class="s1">&#39;distributions.</span><span class="se">\n</span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="k">return</span>

            <span class="n">summary</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span>
            <span class="n">summary</span> <span class="o">+=</span> <span class="sa">f</span><span class="s1">&#39;dataset size: </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">data_loader</span><span class="p">)</span><span class="si">}</span><span class="s1">  </span><span class="se">\n</span><span class="s1">&#39;</span>
            <span class="n">summary</span> <span class="o">+=</span> <span class="sa">f</span><span class="s1">&#39;color mean: </span><span class="si">{</span><span class="n">color_mean</span><span class="o">.</span><span class="n">average</span><span class="si">}</span><span class="s1">  </span><span class="se">\n</span><span class="s1">&#39;</span>
            <span class="n">summary</span> <span class="o">+=</span> <span class="sa">f</span><span class="s1">&#39;color std: </span><span class="si">{</span><span class="n">color_std</span><span class="o">.</span><span class="n">average</span><span class="si">}</span><span class="s1">  </span><span class="se">\n</span><span class="s1">&#39;</span>

            <span class="k">try</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">anchors</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">image_size</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">coverage</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_analyze_anchors_coverage</span><span class="p">(</span><span class="n">anchors</span><span class="o">=</span><span class="n">anchors</span><span class="p">,</span> <span class="n">image_size</span><span class="o">=</span><span class="n">image_size</span><span class="p">,</span>
                                                              <span class="n">title</span><span class="o">=</span><span class="n">title</span><span class="p">,</span> <span class="n">labels</span><span class="o">=</span><span class="n">all_labels</span><span class="p">)</span>
                    <span class="n">summary</span> <span class="o">+=</span> <span class="sa">f</span><span class="s1">&#39;anchors: </span><span class="si">{</span><span class="n">anchors</span><span class="si">}</span><span class="s1">  </span><span class="se">\n</span><span class="s1">&#39;</span>
                    <span class="n">summary</span> <span class="o">+=</span> <span class="sa">f</span><span class="s1">&#39;anchors coverage: </span><span class="si">{</span><span class="n">coverage</span><span class="si">}</span><span class="s1">  </span><span class="se">\n</span><span class="s1">&#39;</span>
            <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
                <span class="n">DatasetStatisticsTensorboardLogger</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Dataset Statistics failed at analyzing anchors &#39;</span>
                                                                <span class="sa">f</span><span class="s1">&#39;coverage.</span><span class="se">\n</span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
                <span class="k">return</span>

            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_text</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s1"> Statistics&#39;</span><span class="p">,</span> <span class="n">text_string</span><span class="o">=</span><span class="n">summary</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>

        <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
            <span class="n">DatasetStatisticsTensorboardLogger</span><span class="o">.</span><span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;dataset analysis failed!</span><span class="se">\n</span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_analyze_class_distribution</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">labels</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">title</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">hist</span><span class="p">,</span> <span class="n">edges</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">labels</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">num_classes</span><span class="p">)</span>

        <span class="n">f</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="p">[</span><span class="mi">10</span><span class="p">,</span> <span class="mi">8</span><span class="p">])</span>

        <span class="n">plt</span><span class="o">.</span><span class="n">bar</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">num_classes</span><span class="p">),</span> <span class="n">hist</span><span class="p">,</span> <span class="n">width</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s1">&#39;#0504aa&#39;</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.7</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">xlim</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;y&#39;</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.75</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">xlabel</span><span class="p">(</span><span class="s1">&#39;Value&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">ylabel</span><span class="p">(</span><span class="s1">&#39;Frequency&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">xticks</span><span class="p">(</span><span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">yticks</span><span class="p">(</span><span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">title</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s1"> class distribution&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_figure</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s2"> class distribution&quot;</span><span class="p">,</span> <span class="n">figure</span><span class="o">=</span><span class="n">f</span><span class="p">)</span>
        <span class="n">text_dist</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">val</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">hist</span><span class="p">):</span>
            <span class="n">text_dist</span> <span class="o">+=</span> <span class="sa">f</span><span class="s1">&#39;[</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s1">]: </span><span class="si">{</span><span class="n">val</span><span class="si">}</span><span class="s1">, &#39;</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_text</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s2"> class distribution&quot;</span><span class="p">,</span> <span class="n">text_string</span><span class="o">=</span><span class="n">text_dist</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_analyze_object_size_distribution</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">labels</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">title</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This function will add two plots to the tensorboard.</span>
<span class="sd">        one is a 2D histogram and the other is a scatter plot. in both cases the X axis is the object width and Y axis</span>
<span class="sd">        is the object width (both normalized by image size)</span>
<span class="sd">        :param labels: all the labels of the dataset of the shape [class_label, x_center, y_center, w, h]</span>
<span class="sd">        :param title: the dataset title</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="c1"># histogram plot</span>
        <span class="n">hist</span><span class="p">,</span> <span class="n">xedges</span><span class="p">,</span> <span class="n">yedges</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram2d</span><span class="p">(</span><span class="n">labels</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">],</span> <span class="n">labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">],</span> <span class="mi">50</span><span class="p">)</span>  <span class="c1"># x and y are deliberately switched</span>

        <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><span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">10</span><span class="p">,</span> <span class="mi">6</span><span class="p">))</span>
        <span class="n">fig</span><span class="o">.</span><span class="n">suptitle</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s1"> boxes w/h distribution&#39;</span><span class="p">)</span>
        <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">121</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s1">&#39;W&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s1">&#39;H&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">imshow</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">hist</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span> <span class="n">interpolation</span><span class="o">=</span><span class="s1">&#39;nearest&#39;</span><span class="p">,</span> <span class="n">origin</span><span class="o">=</span><span class="s1">&#39;lower&#39;</span><span class="p">,</span>
                   <span class="n">extent</span><span class="o">=</span><span class="p">[</span><span class="n">xedges</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">xedges</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">yedges</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">yedges</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]])</span>

        <span class="c1"># scatter plot</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">labels</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">10000</span><span class="p">:</span>
            <span class="c1"># we randomly sample just 10000 objects so that the scatter plot will not get too dense</span>
            <span class="n">labels</span> <span class="o">=</span> <span class="n">labels</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">labels</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">10000</span><span class="p">)]</span>
        <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">122</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s1">&#39;W&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s1">&#39;H&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>

        <span class="n">plt</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span><span class="n">labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">],</span> <span class="n">labels</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">],</span> <span class="n">marker</span><span class="o">=</span><span class="s1">&#39;.&#39;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_figure</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s1"> boxes w/h distribution&#39;</span><span class="p">,</span> <span class="n">figure</span><span class="o">=</span><span class="n">fig</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_get_rect</span><span class="p">(</span><span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">):</span>
        <span class="n">min_w</span> <span class="o">=</span> <span class="n">w</span> <span class="o">/</span> <span class="mf">4.0</span>
        <span class="n">min_h</span> <span class="o">=</span> <span class="n">h</span> <span class="o">/</span> <span class="mf">4.0</span>
        <span class="k">return</span> <span class="n">Rectangle</span><span class="p">((</span><span class="n">min_w</span><span class="p">,</span> <span class="n">min_h</span><span class="p">),</span> <span class="n">w</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">-</span> <span class="n">min_w</span><span class="p">,</span> <span class="n">h</span> <span class="o">*</span> <span class="mi">4</span> <span class="o">-</span> <span class="n">min_h</span><span class="p">,</span> <span class="n">linewidth</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">edgecolor</span><span class="o">=</span><span class="s1">&#39;b&#39;</span><span class="p">,</span> <span class="n">facecolor</span><span class="o">=</span><span class="s1">&#39;none&#39;</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_get_score</span><span class="p">(</span><span class="n">anchors</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">points</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">image_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Calculate the ratio (and 1/ratio) between each anchor width and height and each point (representing a possible</span>
<span class="sd">        object width and height).</span>
<span class="sd">        i.e. for an anchor with w=10,h=20 the point w=11,h=25 will have the ratios 11/10=1.1 and 25/20=1.25</span>
<span class="sd">        or 10/11=0.91 and 20/25=0.8 respectively</span>

<span class="sd">        :param anchors: array of anchors of the shape [2,N]</span>
<span class="sd">        :param points: array of points of the shape [2,M]</span>
<span class="sd">        :param image_size the size of the input image</span>

<span class="sd">        :returns: an array of size [image_size - 1, image_size - 1] where each cell i,j represent the minimum ratio</span>
<span class="sd">        for that cell (point) from all anchors</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="n">ratio</span> <span class="o">=</span> <span class="n">anchors</span><span class="p">[:,</span> <span class="p">:,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">/</span> <span class="n">points</span><span class="p">[:,</span> <span class="p">]</span>
        <span class="n">inv_ratio</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">/</span> <span class="n">ratio</span>
        <span class="n">min_ratio</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">np</span><span class="o">.</span><span class="n">minimum</span><span class="p">(</span><span class="n">ratio</span><span class="p">,</span> <span class="n">inv_ratio</span><span class="p">)</span>
        <span class="n">min_ratio</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">min_ratio</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">to_closest_anchor</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">min_ratio</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
        <span class="n">to_closest_anchor</span><span class="p">[</span><span class="n">to_closest_anchor</span> <span class="o">&gt;</span> <span class="mf">0.75</span><span class="p">]</span> <span class="o">=</span> <span class="mi">2</span>
        <span class="k">return</span> <span class="n">to_closest_anchor</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">image_size</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_analyze_anchors_coverage</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">anchors</span><span class="p">:</span> <span class="n">Anchors</span><span class="p">,</span> <span class="n">image_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">labels</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">title</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This function will add anchors coverage plots to the tensorboard.</span>
<span class="sd">        :param anchors: a list of anchors</span>
<span class="sd">        :param image_size: the input image size for this training</span>
<span class="sd">        :param labels: all the labels of the dataset of the shape [class_label, x_center, y_center, w, h]</span>
<span class="sd">        :param title: the dataset title</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <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><span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">12</span><span class="p">,</span> <span class="mi">5</span><span class="p">))</span>
        <span class="n">fig</span><span class="o">.</span><span class="n">suptitle</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s1"> anchors coverage&#39;</span><span class="p">)</span>

        <span class="c1"># box style plot</span>
        <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">121</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s1">&#39;W&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s1">&#39;H&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_xlim</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="n">image_size</span><span class="p">])</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_ylim</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="n">image_size</span><span class="p">])</span>

        <span class="n">anchors_boxes</span> <span class="o">=</span> <span class="n">anchors</span><span class="o">.</span><span class="n">anchors</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>
        <span class="n">anchors_len</span> <span class="o">=</span> <span class="n">anchors</span><span class="o">.</span><span class="n">num_anchors</span>

        <span class="n">anchors_boxes</span> <span class="o">=</span> <span class="n">anchors_boxes</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>

        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">anchors_len</span><span class="p">):</span>
            <span class="n">rect</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_rect</span><span class="p">(</span><span class="n">anchors_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span> <span class="n">anchors_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">1</span><span class="p">])</span>
            <span class="n">rect</span><span class="o">.</span><span class="n">set_alpha</span><span class="p">(</span><span class="mf">0.3</span><span class="p">)</span>
            <span class="n">rect</span><span class="o">.</span><span class="n">set_facecolor</span><span class="p">([</span><span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">(),</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">(),</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">(),</span> <span class="mf">0.3</span><span class="p">])</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">add_patch</span><span class="p">(</span><span class="n">rect</span><span class="p">)</span>

        <span class="c1"># distance from anchor plot</span>
        <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">122</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s1">&#39;W&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s1">&#39;H&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>

        <span class="n">x</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">image_size</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="n">y</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">image_size</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="n">xx</span><span class="p">,</span> <span class="n">yy</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> <span class="n">sparse</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
        <span class="n">points</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">xx</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">yy</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)])</span>

        <span class="n">color</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_score</span><span class="p">(</span><span class="n">anchors_boxes</span><span class="p">,</span> <span class="n">points</span><span class="p">,</span> <span class="n">image_size</span><span class="p">)</span>

        <span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s1">&#39;W&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s1">&#39;H&#39;</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">STAT_LOGGER_FONT_SIZE</span><span class="p">)</span>
        <span class="n">plt</span><span class="o">.</span><span class="n">imshow</span><span class="p">(</span><span class="n">color</span><span class="p">,</span> <span class="n">interpolation</span><span class="o">=</span><span class="s1">&#39;nearest&#39;</span><span class="p">,</span> <span class="n">origin</span><span class="o">=</span><span class="s1">&#39;lower&#39;</span><span class="p">,</span>
                   <span class="n">extent</span><span class="o">=</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="n">image_size</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">image_size</span><span class="p">])</span>

        <span class="c1"># calculate the coverage for the dataset labels</span>
        <span class="n">cover_masks</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">anchors_len</span><span class="p">):</span>
            <span class="n">w_max</span> <span class="o">=</span> <span class="p">(</span><span class="n">anchors_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">image_size</span><span class="p">)</span> <span class="o">*</span> <span class="mi">4</span>
            <span class="n">w_min</span> <span class="o">=</span> <span class="p">(</span><span class="n">anchors_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">image_size</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.25</span>
            <span class="n">h_max</span> <span class="o">=</span> <span class="p">(</span><span class="n">anchors_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">image_size</span><span class="p">)</span> <span class="o">*</span> <span class="mi">4</span>
            <span class="n">h_min</span> <span class="o">=</span> <span class="p">(</span><span class="n">anchors_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">image_size</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.25</span>
            <span class="n">cover_masks</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span>
                <span class="n">np</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="n">labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">&lt;</span> <span class="n">w_max</span><span class="p">,</span> <span class="n">labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">&gt;</span> <span class="n">w_min</span><span class="p">),</span> <span class="n">labels</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">&lt;</span> <span class="n">h_max</span><span class="p">),</span>
                <span class="n">labels</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">&gt;</span> <span class="n">h_min</span><span class="p">))</span>
        <span class="n">cover_masks</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">cover_masks</span><span class="p">)</span>
        <span class="n">coverage</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">cover_masks</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">))</span> <span class="o">/</span> <span class="nb">len</span><span class="p">(</span><span class="n">labels</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_figure</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">title</span><span class="si">}</span><span class="s1"> anchors coverage&#39;</span><span class="p">,</span> <span class="n">figure</span><span class="o">=</span><span class="n">fig</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">coverage</span></div>


<div class="viewcode-block" id="get_color_augmentation"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.get_color_augmentation">[docs]</a><span class="k">def</span> <span class="nf">get_color_augmentation</span><span class="p">(</span><span class="n">rand_augment_config_string</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">color_jitter</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span> <span class="n">crop_size</span><span class="o">=</span><span class="mi">224</span><span class="p">,</span> <span class="n">img_mean</span><span class="o">=</span><span class="p">[</span><span class="mf">0.485</span><span class="p">,</span> <span class="mf">0.456</span><span class="p">,</span> <span class="mf">0.406</span><span class="p">]):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns color augmentation class. As these augmentation cannot work on top one another, only one is returned</span>
<span class="sd">    according to rand_augment_config_string</span>

<span class="sd">    :param rand_augment_config_string: string which defines the auto augment configurations.</span>
<span class="sd">                                       If none, color jitter will be returned. For possibile values see auto_augment.py</span>
<span class="sd">    :param color_jitter: tuple for color jitter value.</span>
<span class="sd">    :param crop_size: relevant only for auto augment</span>
<span class="sd">    :param img_mean: relevant only for auto augment</span>
<span class="sd">    :return: RandAugment transform or ColorJitter</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">rand_augment_config_string</span><span class="p">:</span>
        <span class="n">auto_augment_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">translate_const</span><span class="o">=</span><span class="nb">int</span><span class="p">(</span><span class="n">crop_size</span> <span class="o">*</span> <span class="mf">0.45</span><span class="p">),</span>
                                   <span class="n">img_mean</span><span class="o">=</span><span class="nb">tuple</span><span class="p">([</span><span class="nb">min</span><span class="p">(</span><span class="mi">255</span><span class="p">,</span> <span class="nb">round</span><span class="p">(</span><span class="mi">255</span> <span class="o">*</span> <span class="n">x</span><span class="p">))</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">img_mean</span><span class="p">]))</span>

        <span class="n">color_augmentation</span> <span class="o">=</span> <span class="n">rand_augment_transform</span><span class="p">(</span><span class="n">rand_augment_config_string</span><span class="p">,</span> <span class="n">auto_augment_params</span><span class="p">)</span>

    <span class="k">else</span><span class="p">:</span>  <span class="c1"># RandAugment includes colorjitter like augmentations, both cannot be applied together.</span>
        <span class="n">color_augmentation</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">ColorJitter</span><span class="p">(</span><span class="o">*</span><span class="n">color_jitter</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">color_augmentation</span></div>


<div class="viewcode-block" id="worker_init_reset_seed"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.datasets_utils.worker_init_reset_seed">[docs]</a><span class="k">def</span> <span class="nf">worker_init_reset_seed</span><span class="p">(</span><span class="n">worker_id</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Make sure each process has different random seed, especially for &#39;fork&#39; method.</span>
<span class="sd">    Check https://github.com/pytorch/pytorch/issues/63311 for more details.</span>

<span class="sd">    :param worker_id: placeholder (needs to be passed to DataLoader init).</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">seed</span> <span class="o">=</span> <span class="n">uuid</span><span class="o">.</span><span class="n">uuid4</span><span class="p">()</span><span class="o">.</span><span class="n">int</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">**</span> <span class="mi">32</span>
    <span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">set_rng_state</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">)</span><span class="o">.</span><span class="n">get_state</span><span class="p">())</span>
    <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="n">seed</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.detection_datasets.coco_detection &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.detection_datasets.coco_detection &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -87,106 +89,97 @@
 
                               <h1>Source code for super_gradients.training.datasets.detection_datasets.coco_detection</h1><div class="highlight"><pre>
                
 
                             <span></span><span class="kn">import</span> <span class="nn">os</span>
                
 
                            -<span class="kn">from</span> <span class="nn">pycocotools.coco</span> <span class="kn">import</span> <span class="n">COCO</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                            -<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">Dataset</span>
                
 
                            -<span class="kn">import</span> <span class="nn">random</span>
                
 
                             <span class="kn">import</span> <span class="nn">cv2</span>
                
 
                            +
                
 
                             <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
                
 
                            -<span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                            -<span class="kn">from</span> <span class="nn">multiprocessing.pool</span> <span class="kn">import</span> <span class="n">ThreadPool</span>
                
 
                            +<span class="kn">from</span> <span class="nn">pycocotools.coco</span> <span class="kn">import</span> <span class="n">COCO</span>
                
 
                            +
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.detection_datasets.detection_dataset</span> <span class="kn">import</span> <span class="n">DetectionDataset</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">DetectionTargetsFormat</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.datasets_conf</span> <span class="kn">import</span> <span class="n">COCO_DETECTION_CLASSES_LIST</span>
                
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="COCODetectionDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.COCODetectionDataset">[docs]</a><span class="k">class</span> <span class="nc">COCODetectionDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
                
 
                            -    <span class="sd">&quot;&quot;&quot;Detection dataset COCO implementation&quot;&quot;&quot;</span>
                
 
                            +<div class="viewcode-block" id="COCODetectionDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.COCODetectionDataset">[docs]</a><span class="k">class</span> <span class="nc">COCODetectionDataset</span><span class="p">(</span><span class="n">DetectionDataset</span><span class="p">):</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Dataset for COCO object detection.&quot;&quot;&quot;</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            -            <span class="bp">self</span><span class="p">,</span> <span class="n">img_size</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span>
                
 
                            -            <span class="n">data_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -            <span class="n">json_file</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;instances_train2017.json&quot;</span><span class="p">,</span>
                
 
                            -            <span class="n">name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;images/train2017&quot;</span><span class="p">,</span>
                
 
                            -            <span class="n">cache</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -            <span class="n">cache_dir_path</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -            <span class="n">tight_box_rotation</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -            <span class="n">transforms</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="p">[],</span>
                
 
                            -            <span class="n">with_crowd</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">json_file</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;instances_train2017.json&quot;</span><span class="p">,</span>
                
 
                            +        <span class="n">subdir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;images/train2017&quot;</span><span class="p">,</span>
                
 
                            +        <span class="n">tight_box_rotation</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">with_crowd</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +        <span class="o">*</span><span class="n">args</span><span class="p">,</span>
                
 
                            +        <span class="o">**</span><span class="n">kwargs</span><span class="p">,</span>
                
 
                                 <span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        :param img_size: tuple, Image size (when loaded, before transforms)</span>
                
 
                            -<span class="sd">        :param data_dir: str, root path to coco data.</span>
                
 
                            -<span class="sd">        :param json_file: str, path to coco json file, that resides in data_dir/annotations/json_file.</span>
                
 
                            -<span class="sd">        :param name: str, sub directory of data_dir containing the data.</span>
                
 
                            -<span class="sd">        :param cache: bool, whether to cache images</span>
                
 
                            -<span class="sd">        :param cache_dir_path: str, path to a directory that will be used for caching (with memmap).</span>
                
 
                            -<span class="sd">        :param tight_box_rotation: bool, whether to use of segmentation maps convex hull</span>
                
 
                            -<span class="sd">         as target_seg (see load_sample docs).</span>
                
 
                            -<span class="sd">        :param transforms: list of transforms to apply sequentially on sample in __getitem__</span>
                
 
                            +<span class="sd">        :param json_file:           Name of the coco json file, that resides in data_dir/annotations/json_file.</span>
                
 
                            +<span class="sd">        :param subdir:              Sub directory of data_dir containing the data.</span>
                
 
                            +<span class="sd">        :param tight_box_rotation:  bool, whether to use of segmentation maps convex hull as target_seg</span>
                
 
                            +<span class="sd">                                    (check get_sample docs).</span>
                
 
                             <span class="sd">        :param with_crowd: Add the crowd groundtruths to __getitem__</span>
                
 
                            +
                
 
                            +<span class="sd">        kwargs:</span>
                
 
                            +<span class="sd">            all_classes_list: all classes list, default is COCO_DETECTION_CLASSES_LIST.</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span> <span class="o">=</span> <span class="n">data_dir</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span> <span class="o">=</span> <span class="n">img_size</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span> <span class="o">=</span> <span class="n">data_dir</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">subdir</span> <span class="o">=</span> <span class="n">subdir</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">json_file</span> <span class="o">=</span> <span class="n">json_file</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">=</span> <span class="n">name</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">tight_box_rotation</span> <span class="o">=</span> <span class="n">tight_box_rotation</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">with_crowd</span> <span class="o">=</span> <span class="n">with_crowd</span>
                
 
                            +        <span class="n">target_fields</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">,</span> <span class="s2">&quot;crowd_target&quot;</span><span class="p">]</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">with_crowd</span> <span class="k">else</span> <span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span>
                
 
                            +        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;target_fields&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">target_fields</span>
                
 
                            +        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;output_fields&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">,</span> <span class="o">*</span><span class="n">target_fields</span><span class="p">]</span>
                
 
                            +        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;original_target_format&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">XYXY_LABEL</span>
                
 
                            +        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;all_classes_list&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;all_classes_list&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="n">COCO_DETECTION_CLASSES_LIST</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_setup_data_source</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;Initialize img_and_target_path_list and warn if label file is missing</span>
                
 
                            +
                
 
                            +<span class="sd">        :return: List of tuples made of (img_path,target_path)</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">coco</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_init_coco</span><span class="p">()</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">class_ids</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">getCatIds</span><span class="p">())</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="nb">list</span><span class="p">([</span><span class="n">category</span><span class="p">[</span><span class="s2">&quot;name&quot;</span><span class="p">]</span> <span class="k">for</span> <span class="n">category</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">loadC
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">sample_id_to_coco_id</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">getImgIds</span><span class="p">()</span>
                
 
                            +        <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">sample_id_to_coco_id</span><span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_init_coco</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">COCO</span><span class="p">:</span>
                
 
                                     <span class="n">annotation_file_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="p">,</span> <span class="s2">&quot;annotations&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">json_file</span><span c
                
 
                                     <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">annotation_file_path</span><span class="p">):</span>
                
 
                                         <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Could not find annotation file under &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">annotation_file_path</span><span class="p">))</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">coco</span> <span class="o">=</span> <span class="n">COCO</span><span class="p">(</span><span class="n">annotation_file_path</span><span class="p">)</span>  <span class="c1"># duplicate</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">tight_box_rotation</span> <span class="o">=</span> <span class="n">tight_box_rotation</span>
                
 
                            -        <span class="n">remove_useless_info</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">tight_box_rotation</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">getImgIds</span><span class="p">()</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">class_ids</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">getCatIds</span><span class="p">())</span>
                
 
                            -        <span class="n">cats</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">loadCats</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">getCatIds</span><span class="p">())</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">_classes</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="n">c</span><span class="p">[</span><span class="s2">&quot;name&quot;</span><span class="p">]</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="n">cats</span><span class="p">])</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_coco_annotations</span><span class="p">()</span>
                
 
                            -        <span class="k">if</span> <span class="n">cache</span><span class="p">:</span>  <span class="c1"># cache after merged</span>
                
 
                            -            <span class="k">if</span> <span class="n">cache_dir_path</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">cache_dir_path</span><span class="p">):</span>
                
 
                            -                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Must pass valid path through cache_dir_path when caching. Got &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">cache_dir_path</span><span class="p">))</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">cache_dir_path</span> <span class="o">=</span> <span class="n">cache_dir_path</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">_cache_images</span><span class="p">()</span>
                
 
                            -
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span> <span class="o">=</span> <span class="n">transforms</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="fm">__getitem__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">idx</span><span class="p">):</span>
                
 
                            -        <span class="n">sample</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">load_sample</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>
                
 
                            -        <span class="n">sample</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">apply_transforms</span><span class="p">(</span><span class="n">sample</span><span class="p">)</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">with_crowd</span><span class="p">:</span>
                
 
                            -            <span class="k">return</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;crowd_target&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;info&quot;</s
                
 
                            -        <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="k">return</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;info&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;id&quot;</span><span 
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="fm">__len__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            -        <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ids</span><span class="p">)</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_load_coco_annotations</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            -        <span class="k">return</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">_load_anno_from_ids</span><span class="p">(</span><span class="n">_ids</span><span class="p">)</span> <span class="k">for</span> <span class="n">_ids</span> <span class="ow">in</span> <span class="n">tqdm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ids</span><span class="p">,</span> <span class="n">desc</span><spa
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_load_anno_from_ids</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">id_</span><span class="p">):</span>
                
 
                            +
                
 
                            +        <span class="n">coco</span> <span class="o">=</span> <span class="n">COCO</span><span class="p">(</span><span class="n">annotation_file_path</span><span class="p">)</span>
                
 
                            +        <span class="n">remove_useless_info</span><span class="p">(</span><span class="n">coco</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">tight_box_rotation</span><span class="p">)</span>
                
 
                            +        <span class="k">return</span> <span class="n">coco</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_load_annotation</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample_id</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        Load relevant information of a specific image</span>
                
 
                            -
                
 
                            -<span class="sd">        :param id_: image id</span>
                
 
                            -<span class="sd">        :return res:            Target Bboxes (detection)</span>
                
 
                            -<span class="sd">        :return res_crowd:      Crowd target Bboxes (detection)</span>
                
 
                            -<span class="sd">        :return res_seg:        Segmentation</span>
                
 
                            -<span class="sd">        :return img_info:       Image (height, width)</span>
                
 
                            -<span class="sd">        :return resized_info:   Resides image (height, width)</span>
                
 
                            -<span class="sd">        :return img_path:       Path to the associated image</span>
                
 
                            +<span class="sd">        Load relevant information of a specific image.</span>
                
 
                            +
                
 
                            +<span class="sd">        :param sample_id:               Sample_id in the dataset</span>
                
 
                            +<span class="sd">        :return target:                 Target Bboxes (detection) in XYXY_LABEL format</span>
                
 
                            +<span class="sd">        :return crowd_target:           Crowd target Bboxes (detection) in XYXY_LABEL format</span>
                
 
                            +<span class="sd">        :return target_segmentation:    Segmentation</span>
                
 
                            +<span class="sd">        :return initial_img_shape:      Image (height, width)</span>
                
 
                            +<span class="sd">        :return resized_img_shape:      Resides image (height, width)</span>
                
 
                            +<span class="sd">        :return img_path:               Path to the associated image</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">im_ann</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">loadImgs</span><span class="p">(</span><span class="n">id_</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
                
 
                            -        <span class="n">width</span> <span class="o">=</span> <span class="n">im_ann</span><span class="p">[</span><span class="s2">&quot;width&quot;</span><span class="p">]</span>
                
 
                            -        <span class="n">height</span> <span class="o">=</span> <span class="n">im_ann</span><span class="p">[</span><span class="s2">&quot;height&quot;</span><span class="p">]</span>
                
 
                            -        <span class="n">anno_ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">getAnnIds</span><span class="p">(</span><span class="n">imgIds</span><span class="o">=</span><span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">id_</span><span class="p">)])</span>
                
 
                            -        <span class="n">annotations</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">loadAnns</span><span class="p">(</span><span class="n">anno_ids</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">img_id</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">sample_id_to_coco_id</span><span class="p">[</span><span class="n">sample_id</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="n">img_metadata</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">loadImgs</span><span class="p">(</span><span class="n">img_id</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
                
 
                            +        <span class="n">width</span> <span class="o">=</span> <span class="n">img_metadata</span><span class="p">[</span><span class="s2">&quot;width&quot;</span><span class="p">]</span>
                
 
                            +        <span class="n">height</span> <span class="o">=</span> <span class="n">img_metadata</span><span class="p">[</span><span class="s2">&quot;height&quot;</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="n">img_annotation_ids</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">getAnnIds</span><span class="p">(</span><span class="n">imgIds</span><span class="o">=</span><span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">img_id</span><span class="p">)])</span>
                
 
                            +        <span class="n">img_annotations</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">coco</span><span class="o">.</span><span class="n">loadAnns</span><span class="p">(</span><span class="n">img_annotation_ids</span><span class="p">)</span>
                
 
                                     <span class="n">cleaned_annotations</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                            -        <span class="k">for</span> <span class="n">annotation</span> <span class="ow">in</span> <span class="n">annotations</span><span class="p">:</span>
                
 
                            +        <span class="k">for</span> <span class="n">annotation</span> <span class="ow">in</span> <span class="n">img_annotations</span><span class="p">:</span>
                
 
                                         <span class="n">x1</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;bbox&quot;</span><span class="p">][</span><span class="mi">0</span><span class="p">]))</span>
                
 
                                         <span class="n">y1</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;bbox&quot;</span><span class="p">][</span><span class="mi">1</span><span class="p">]))</span>
                
 
                                         <span class="n">x2</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">min</span><span class="p">((</span><span class="n">width</span><span class="p">,</span> <span class="n">x1</span> <span class="o">+</span> <span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;bb
                
@@ -197,14 +190,14 @@
 
                                     <span class="n">non_crowd_annotations</span> <span class="o">=</span> <span class="p">[</span><span class="n">annotation</span> <span class="k">for</span> <span class="n">annotation</span> <span class="ow">in</span> <span class="n">cleaned_annotations</span> <span class="k">if</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;iscrowd&quot;</span><span class="p">]</span> <span class="o">==</span> <span class="mi">0</span><span class="p">]</span>
                
 
                            -        <span class="n">res</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="nb">len</span><span class="p">(</span><span class="n">non_crowd_annotations</span><span class="p">),</span> <span class="mi">5</span><span class="p">))</span>
                
 
                            +        <span class="n">target</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="nb">len</span><span class="p">(</span><span class="n">non_crowd_annotations</span><span class="p">),</span> <span class="mi">5</span><span class="p">))</span>
                
 
                                     <span class="n">num_seg_values</span> <span class="o">=</span> <span class="mi">98</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">tight_box_rotation</span> <span class="k">else</span> <span class="mi">0</span>
                
 
                            -        <span class="n">res_seg</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="nb">len</span><span class="p">(</span><span class="n">non_crowd_annotations</span><span class="p">),</span> <span class="n">num_seg_values</span><span class="p">))</span>
                
 
                            -        <span class="n">res_seg</span><span class="o">.</span><span class="n">fill</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">nan</span><span class="p">)</span>
                
 
                            +        <span class="n">target_segmentation</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="nb">len</span><span class="p">(</span><span class="n">non_crowd_annotations</span><span class="p">),</span> <span class="n">num_seg_values</span><span class="p">))</span>
                
 
                            +        <span class="n">target_segmentation</span><span class="o">.</span><span class="n">fill</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">nan</span><span class="p">)</span>
                
 
                                     <span class="k">for</span> <span class="n">ix</span><span class="p">,</span> <span class="n">annotation</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">non_crowd_annotations</span><span class="p">):</span>
                
 
                                         <span class="bp">cls</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">class_ids</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;category_id&quot;</span><span class="p">])</span>
                
 
                            -            <span class="n">res</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">0</span><span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;clean_bbox&quot;</span><span class="p">]</span>
                
 
                            -            <span class="n">res</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="bp">cls</span>
                
 
                            +            <span class="n">target</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">0</span><span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;clean_bbox&quot;</span><span class="p">]</span>
                
 
                            +            <span class="n">target</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="bp">cls</span>
                
 
                                         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">tight_box_rotation</span><span class="p">:</span>
                
 
                                             <span class="n">seg_points</span> <span class="o">=</span> <span class="p">[</span><span class="n">j</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">annotation</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;segmentation&quot;</span><span class="p">,</span> <span class="p">[])</span> <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span clas
                
 
                                             <span class="k">if</span> <span class="n">seg_points</span><span class="p">:</span>
                
@@ -212,181 +205,41 @@
 
                                                 <span class="n">seg_points_convex</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">convexHull</span><span class="p">(</span><span class="n">seg_points_c</span><span class="p">)</span><span class="o">.</span><span class="n">ravel</span><span class="p">()</span>
                
 
                                             <span class="k">else</span><span class="p">:</span>
                
 
                                                 <span class="n">seg_points_convex</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                            -                <span class="n">res_seg</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="p">:</span><span class="nb">len</span><span class="p">(</span><span class="n">seg_points_convex</span><span class="p">)]</span> <span class="o">=</span> <span class="n">seg_points_convex</span>
                
 
                            +                <span class="n">target_segmentation</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="p">:</span> <span class="nb">len</span><span class="p">(</span><span class="n">seg_points_convex</span><span class="p">)]</span> <span class="o">=</span> <span class="n">seg_points_convex</span>
                
 
                                     <span class="n">crowd_annotations</span> <span class="o">=</span> <span class="p">[</span><span class="n">annotation</span> <span class="k">for</span> <span class="n">annotation</span> <span class="ow">in</span> <span class="n">cleaned_annotations</span> <span class="k">if</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;iscrowd&quot;</span><span class="p">]</span> <span class="o">==</span> <span class="mi">1</span><span class="p">]</span>
                
 
                            -        <span class="n">res_crowd</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="nb">len</span><span class="p">(</span><span class="n">crowd_annotations</span><span class="p">),</span> <span class="mi">5</span><span class="p">))</span>
                
 
                            +        <span class="n">crowd_target</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="nb">len</span><span class="p">(</span><span class="n">crowd_annotations</span><span class="p">),</span> <span class="mi">5</span><span class="p">))</span>
                
 
                                     <span class="k">for</span> <span class="n">ix</span><span class="p">,</span> <span class="n">annotation</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">crowd_annotations</span><span class="p">):</span>
                
 
                                         <span class="bp">cls</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">class_ids</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;category_id&quot;</span><span class="p">])</span>
                
 
                            -            <span class="n">res_crowd</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">0</span><span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;clean_bbox&quot;</span><span class="p">]</span>
                
 
                            -            <span class="n">res_crowd</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="bp">cls</span>
                
 
                            +            <span class="n">crowd_target</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">0</span><span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;clean_bbox&quot;</span><span class="p">]</span>
                
 
                            +            <span class="n">crowd_target</span><span class="p">[</span><span class="n">ix</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="bp">cls</span>
                
 
                                     <span class="n">r</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">height</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">]</s
                
 
                            -        <span class="n">res</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">r</span>
                
 
                            -        <span class="n">res_crowd</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">r</span>
                
 
                            -        <span class="n">res_seg</span> <span class="o">*=</span> <span class="n">r</span>
                
 
                            -
                
 
                            -        <span class="n">img_info</span> <span class="o">=</span> <span class="p">(</span><span class="n">height</span><span class="p">,</span> <span class="n">width</span><span class="p">)</span>
                
 
                            -        <span class="n">resized_info</span> <span class="o">=</span> <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">height</span> <span class="o">*</span> <span class="n">r</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">width</span> <span class="o">*</span> <span class="n">r</span><span class="p">))</span>
                
 
                            -
                
 
                            -        <span class="n">file_name</span> <span class="o">=</span> <span class="p">(</span>
                
 
                            -            <span class="n">im_ann</span><span class="p">[</span><span class="s2">&quot;file_name&quot;</span><span class="p">]</span>
                
 
                            -            <span class="k">if</span> <span class="s2">&quot;file_name&quot;</span> <span class="ow">in</span> <span class="n">im_ann</span>
                
 
                            -            <span class="k">else</span> <span class="s2">&quot;</span><span class="si">{:012}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">id_</span><span class="p">)</span> <span class="o">+</span> <span class="s2">&quot;.jpg&quot;</span>
                
 
                            -        <span class="p">)</span>
                
 
                            -        <span class="n">img_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="n">file_name</span><span class="p">)</span>
                
 
                            -        <span class="k">return</span> <span class="n">res</span><span class="p">,</span> <span class="n">res_crowd</span><span class="p">,</span> <span class="n">res_seg</span><span class="p">,</span> <span class="n">img_info</span><span class="p">,</span> <span class="n">resized_info</span><span class="p">,</span> <span class="n">img_path</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_cache_images</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            -        <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
                
 
                            -            <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">********************************************************************************</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -            <span class="s2">&quot;You are using cached images in RAM to accelerate training.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -            <span class="s2">&quot;This requires large system RAM.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -            <span class="s2">&quot;Make sure you have 200G+ RAM and 136G available disk space for training COCO.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -            <span class="s2">&quot;********************************************************************************</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -        <span class="p">)</span>
                
 
                            -        <span class="n">max_h</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                
 
                            -        <span class="n">max_w</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                
 
                            -        <span class="n">cache_file</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cache_dir_path</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;img_resized_cache_</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">n
                
 
                            -        <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">cache_file</span><span class="p">):</span>
                
 
                            -            <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span>
                
 
                            -                <span class="s2">&quot;Caching images for the first time. This might take about 20 minutes for COCO&quot;</span>
                
 
                            -            <span class="p">)</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">memmap</span><span class="p">(</span>
                
 
                            -                <span class="n">cache_file</span><span class="p">,</span>
                
 
                            -                <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ids</span><span class="p">),</span> <span class="n">max_h</span><span class="p">,</span> <span class="n">max_w</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span>
                
 
                            -                <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">,</span>
                
 
                            -                <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;w+&quot;</span><span class="p">,</span>
                
 
                            -            <span class="p">)</span>
                
 
                            -
                
 
                            -            <span class="n">NUM_THREADs</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="mi">8</span><span class="p">,</span> <span class="n">os</span><span class="o">.</span><span class="n">cpu_count</span><span class="p">())</span>
                
 
                            -            <span class="n">loaded_images</span> <span class="o">=</span> <span class="n">ThreadPool</span><span class="p">(</span><span class="n">NUM_THREADs</span><span class="p">)</span><span class="o">.</span><span class="n">imap</span><span class="p">(</span>
                
 
                            -                <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">load_resized_img</span><span class="p">(</span><span class="n">x</span><span class="p">),</span>
                
 
                            -                <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">)),</span>
                
 
                            -            <span class="p">)</span>
                
 
                            -            <span class="n">pbar</span> <span class="o">=</span> <span class="n">tqdm</span><span class="p">(</span><span class="nb">enumerate</span><span class="p">(</span><span class="n">loaded_images</span><span class="p">),</span> <span class="n">total</span><span class="o">=</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">))</span>
                
 
                            -            <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">out</span> <span class="ow">in</span> <span class="n">pbar</span><span class="p">:</span>
                
 
                            -                <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span><span class="p">[</span><span class="n">k</span><span class="p">][:</span> <span class="n">out</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="p">:</span> <span class="n">out</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span>
                
 
                            -            <span class="n">pbar</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
                
 
                            -        <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
                
 
                            -                <span class="s2">&quot;You are using cached imgs! Make sure your dataset is not changed!!</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -                <span class="s2">&quot;Everytime the self.input_size is changed in your exp file, you need to delete</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -                <span class="s2">&quot;the cached data and re-generate them.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -            <span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Loading cached imgs...&quot;</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">memmap</span><span class="p">(</span>
                
 
                            -            <span class="n">cache_file</span><span class="p">,</span>
                
 
                            -            <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ids</span><span class="p">),</span> <span class="n">max_h</span><span class="p">,</span> <span class="n">max_w</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span>
                
 
                            -            <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">,</span>
                
 
                            -            <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;r+&quot;</span><span class="p">,</span>
                
 
                            -        <span class="p">)</span>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="COCODetectionDataset.load_resized_img"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.COCODetectionDataset.load_resized_img">[docs]</a>    <span class="k">def</span> <span class="nf">load_resized_img</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">):</span>
                
 
                            -        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        Loads image at index, and resizes it to self.input_dim</span>
                
 
                            -
                
 
                            -<span class="sd">        :param index: index to load the image from</span>
                
 
                            -<span class="sd">        :return: resized_img</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">load_image</span><span class="p">(</span><span class="n">index</span><span class="p">)</span>
                
 
                            -        <span class="n">r</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><s
                
 
                            -        <span class="n">resized_img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span>
                
 
                            -            <span class="n">img</span><span class="p">,</span>
                
 
                            -            <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">r</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span c
                
 
                            -            <span class="n">interpolation</span><span class="o">=</span><span class="n">cv2</span><span class="o">.</span><span class="n">INTER_LINEAR</span><span class="p">,</span>
                
 
                            -        <span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span>
                
 
                            -        <span class="k">return</span> <span class="n">resized_img</span></div>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="COCODetectionDataset.load_sample"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.COCODetectionDataset.load_sample">[docs]</a>    <span class="k">def</span> <span class="nf">load_sample</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">):</span>
                
 
                            -        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        Loads sample at self.ids[index] as dictionary that holds:</span>
                
 
                            -<span class="sd">            &quot;image&quot;: Image resized to self.input_dim</span>
                
 
                            -<span class="sd">            &quot;target&quot;: Detection ground truth, np.array shaped (num_targets, 5), format is [class,x1,y1,x2,y2] with</span>
                
 
                            -<span class="sd">                image coordinates.</span>
                
 
                            -<span class="sd">            &quot;target_seg&quot;: Segmentation map convex hull derived detection target.</span>
                
 
                            -<span class="sd">            &quot;info&quot;: Original shape (height,width).</span>
                
 
                            -<span class="sd">            &quot;id&quot;: COCO image id</span>
                
 
                            -
                
 
                            -<span class="sd">        :param index: Sample index</span>
                
 
                            -<span class="sd">        :return: sample as described above</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">id_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ids</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>
                
 
                            -        <span class="n">res</span><span class="p">,</span> <span class="n">res_crowd</span><span class="p">,</span> <span class="n">res_seg</span><span class="p">,</span> <span class="n">img_info</span><span class="p">,</span> <span class="n">resized_info</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="n">pad_img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>
                
 
                            -            <span class="n">img</span> <span class="o">=</span> <span class="n">pad_img</span><span class="p">[:</span> <span class="n">resized_info</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="p">:</span> <span class="n">resized_info</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="p">:]</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                
 
                            -        <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">load_resized_img</span><span class="p">(</span><span class="n">index</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="n">sample</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;image&quot;</span><span class="p">:</span> <span class="n">img</span><span class="p">,</span> <span class="s2">&quot;target&quot;</span><span class="p">:</span> <span class="n">res</span><span class="o">.</span><span class="n">copy</span><span class="p">(),</span> <span class="s2">&quot;target_seg&quot;</span><span class="p">:</span> <span class="n">res_seg</span><span class="p">,</span
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">with_crowd</span><span class="p">:</span>
                
 
                            -            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;crowd_target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">res_crowd</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                
 
                            -        <span class="k">return</span> <span class="n">sample</span></div>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="COCODetectionDataset.load_image"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.COCODetectionDataset.load_image">[docs]</a>    <span class="k">def</span> <span class="nf">load_image</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">):</span>
                
 
                            -        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        Loads image at index with its original resolution</span>
                
 
                            -<span class="sd">        :param index: index in self.annotations</span>
                
 
                            -<span class="sd">        :return: image (np.array)</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">file_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">5</span><span class="p">]</span>
                
 
                            -
                
 
                            -        <span class="n">img_file</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">file_name</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="n">img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">imread</span><span class="p">(</span><span class="n">img_file</span><span class="p">)</span>
                
 
                            -        <span class="k">assert</span> <span class="n">img</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                
 
                            -
                
 
                            -        <span class="k">return</span> <span class="n">img</span></div>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="fm">__del__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            -        <span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">imgs</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_load_anno</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">):</span>
                
 
                            -        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_get_random_non_empty_target_idx</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            -        <span class="n">target</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                            -        <span class="k">while</span> <span class="nb">len</span><span class="p">(</span><span class="n">target</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                
 
                            -            <span class="n">idx</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ids</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
                
 
                            -            <span class="n">target</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_anno</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>
                
 
                            -        <span class="k">return</span> <span class="n">idx</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_load_random_samples</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="p">,</span> <span class="n">non_empty_targets_only</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
                
 
                            -        <span class="n">inds</span> <span class="o">=</span> <span class="p">[</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">_get_random_non_empty_target_idx</span><span class="p">()</span> <span class="k">if</span> <span class="n">non_empty_targets_only</span> <span class="k">else</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.<
                
 
                            -            <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="n">count</span><span class="p">)]</span>
                
 
                            -        <span class="k">return</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">load_sample</span><span class="p">(</span><span class="n">ind</span><span class="p">)</span> <span class="k">for</span> <span class="n">ind</span> <span class="ow">in</span> <span class="n">inds</span><span class="p">]</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_load_additional_inputs_for_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">,</span> <span class="n">transform</span><span class="p">):</span>
                
 
                            -        <span class="n">additional_samples_count</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">additional_samples_count</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span>
                
 
                            -                                                                                 <span class="s2">&quot;additional_samples_count&quot;</span><span class="p">)</span> <span class="k">else</span> <span class="mi">0</span>
                
 
                            -        <span class="n">non_empty_targets</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">non_empty_targets</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span> <span class="s2">&quot;non_empty_targets&quot;</span><span class="p">)</span> <span class="k">else</span> <span class="kc">False</span>
                
 
                            -        <span class="n">additional_samples</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_random_samples</span><span class="p">(</span><span class="n">additional_samples_count</span><span class="p">,</span> <span class="n">non_empty_targets</span><span class="p">)</span>
                
 
                            -        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;additional_samples&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">additional_samples</span>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="COCODetectionDataset.apply_transforms"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.COCODetectionDataset.apply_transforms">[docs]</a>    <span class="k">def</span> <span class="nf">apply_transforms</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class
                
 
                            -        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        Applies self.transforms sequentially to sample</span>
                
 
                            -
                
 
                            -<span class="sd">        If a transforms has the attribute &#39;additional_samples_count&#39;, additional samples will be loaded and stored in</span>
                
 
                            -<span class="sd">         sample[&quot;additional_samples&quot;] prior to applying it. Combining with the attribute &quot;non_empty_targets&quot; will load</span>
                
 
                            -<span class="sd">         only additional samples with objects in them.</span>
                
 
                            -
                
 
                            -<span class="sd">        :param sample: Sample to apply the transforms on to (loaded with self.load_sample)</span>
                
 
                            -<span class="sd">        :return: Transformed sample</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">for</span> <span class="n">transform</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span><span class="p">:</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">_load_additional_inputs_for_transform</span><span class="p">(</span><span class="n">sample</span><span class="p">,</span> <span class="n">transform</span><span class="p">)</span>
                
 
                            -            <span class="n">sample</span> <span class="o">=</span> <span class="n">transform</span><span class="p">(</span><span class="n">sample</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="k">return</span> <span class="n">sample</span></div></div>
                
 
                            -
                
 
                            -
                
 
                            -<div class="viewcode-block" id="remove_useless_info"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.remove_useless_info">[docs]</a><span class="k">def</span> <span class="nf">remove_useless_info</span><span class="p">(</span><span class="n">coco</span><span class="p">,</span> <span class="n">use_seg_info</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
                
 
                            +        <span class="n">target</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">r</span>
                
 
                            +        <span class="n">crowd_target</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">r</span>
                
 
                            +        <span class="n">target_segmentation</span> <span class="o">*=</span> <span class="n">r</span>
                
 
                            +
                
 
                            +        <span class="n">initial_img_shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">height</span><span class="p">,</span> <span class="n">width</span><span class="p">)</span>
                
 
                            +        <span class="n">resized_img_shape</span> <span class="o">=</span> <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">height</span> <span class="o">*</span> <span class="n">r</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">width</span> <span class="o">*</span> <span class="n">r</span><span class="p">))</span>
                
 
                            +
                
 
                            +        <span class="n">file_name</span> <span class="o">=</span> <span class="n">img_metadata</span><span class="p">[</span><span class="s2">&quot;file_name&quot;</span><span class="p">]</span> <span class="k">if</span> <span class="s2">&quot;file_name&quot;</span> <span class="ow">in</span> <span class="n">img_metadata</span> <span class="k">else</span> <span class="s2">&quot;</span><span class="si">{:012}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</spa
                
 
                            +        <span class="n">img_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">subdir</span><span class="p">,</span> <span class="n">file_name</span><span class="p">)</span>
                
 
                            +        <span class="n">img_id</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">sample_id_to_coco_id</span><span class="p">[</span><span class="n">sample_id</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="n">annotation</span> <span class="o">=</span> <span class="p">{</span>
                
 
                            +            <span class="s2">&quot;target&quot;</span><span class="p">:</span> <span class="n">target</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;crowd_target&quot;</span><span class="p">:</span> <span class="n">crowd_target</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;target_segmentation&quot;</span><span class="p">:</span> <span class="n">target_segmentation</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;initial_img_shape&quot;</span><span class="p">:</span> <span class="n">initial_img_shape</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;resized_img_shape&quot;</span><span class="p">:</span> <span class="n">resized_img_shape</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;img_path&quot;</span><span class="p">:</span> <span class="n">img_path</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;id&quot;</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">img_id</span><span class="p">]),</span>
                
 
                            +        <span class="p">}</span>
                
 
                            +        <span class="k">return</span> <span class="n">annotation</span></div>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">remove_useless_info</span><span class="p">(</span><span class="n">coco</span><span class="p">,</span> <span class="n">use_seg_info</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Remove useless info in coco dataset. COCO object is modified inplace.</span>
                
 
                             <span class="sd">    This function is mainly used for saving memory (save about 30% mem).</span>
                
@@ -402,7 +255,7 @@
 
                                         <span class="n">img</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;flickr_url&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="s2">&quot;annotations&quot;</span> <span class="ow">in</span> <span class="n">coco</span><span class="o">.</span><span class="n">dataset</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">use_seg_info</span><span class="p">:</span>
                
 
                                         <span class="k">for</span> <span class="n">anno</span> <span class="ow">in</span> <span class="n">coco</span><span class="o">.</span><span class="n">dataset</span><span class="p">[</span><span class="s2">&quot;annotations&quot;</span><span class="p">]:</span>
                
 
                            -                <span class="n">anno</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;segmentation&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span></div>
                
 
                            +                <span class="n">anno</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;segmentation&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -432,4 +285,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.detection_datasets.detection_dataset &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.detection_datasets.detection_dataset &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -93,21 +95,25 @@
 
                             <span class="kn">import</span> <span class="nn">cv2</span>
                
 
                             <span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="nn">plt</span>
                
 
                             <span class="kn">from</span> <span class="nn">pathlib</span> <span class="kn">import</span> <span class="n">Path</span>
                
 
                            +<span class="kn">from</span> <span class="nn">copy</span> <span class="kn">import</span> <span class="n">deepcopy</span>
                
 
                             <span class="kn">import</span> <span class="nn">hashlib</span>
                
 
                             <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
                
 
                             <span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                             <span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">Dataset</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.decorators.factory_decorator</span> <span class="kn">import</span> <span class="n">resolve_param</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">get_cls_posx_in_target</span><span class="p">,</span> <span class="n">DetectionTargetsFormat</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.transforms.transforms</span> <span class="kn">import</span> <span class="n">DetectionTransform</span><span class="p">,</span> <span class="n">DetectionTargetsFormatTransform</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.dataset_exceptions</span> <span class="kn">import</span> <span class="n">EmptyDatasetException</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.factories.list_factory</span> <span class="kn">import</span> <span class="n">ListFactory</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.factories.transforms_factory</span> <span class="kn">import</span> <span class="n">TransformsFactory</span>
                
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="DetectionDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset">[docs]</a><span class="k">class</span> <span class="nc">DetectionDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="DetectionDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset">[docs]</a><span class="k">class</span> <span class="nc">DetectionDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Detection dataset.</span>
                
 
                             <span class="sd">    This is a boilerplate class to facilitate the implementation of datasets.</span>
                
@@ -146,20 +152,21 @@
 
                             <span class="sd">                                &gt; Therefore, we also have len(self) = 100</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s2">&quot;transforms&quot;</span><span class="p">,</span> <span class="n">ListFactory</span><span class="p">(</span><span class="n">TransformsFactory</span><span class="p">()))</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            -            <span class="bp">self</span><span class="p">,</span>
                
 
                            -            <span class="n">data_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                
 
                            -            <span class="n">input_dim</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span>
                
 
                            -            <span class="n">original_target_format</span><span class="p">:</span> <span class="n">DetectionTargetsFormat</span><span class="p">,</span>
                
 
                            -            <span class="n">max_num_samples</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -            <span class="n">cache</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -            <span class="n">cache_path</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -            <span class="n">transforms</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">DetectionTransform</span><span class="p">]</span> <span class="o">=</span> <span class="p">[],</span>
                
 
                            -            <span class="n">all_classes_list</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -            <span class="n">class_inclusion_list</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -            <span class="n">ignore_empty_annotations</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            -            <span class="n">target_fields</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -            <span class="n">output_fields</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">data_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                
 
                            +        <span class="n">input_dim</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span>
                
 
                            +        <span class="n">original_target_format</span><span class="p">:</span> <span class="n">DetectionTargetsFormat</span><span class="p">,</span>
                
 
                            +        <span class="n">max_num_samples</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">cache</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">cache_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">transforms</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">DetectionTransform</span><span class="p">]</span> <span class="o">=</span> <span class="p">[],</span>
                
 
                            +        <span class="n">all_classes_list</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">class_inclusion_list</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">ignore_empty_annotations</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +        <span class="n">target_fields</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">output_fields</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                                 <span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Detection dataset.</span>
                
@@ -169,7 +176,7 @@
 
                             <span class="sd">                                        differ based on transforms.</span>
                
 
                             <span class="sd">        :param max_num_samples:         If not None, set the maximum size of the dataset by only indexing the first n annotations/images.</span>
                
 
                             <span class="sd">        :param cache:                   Whether to cache images or not.</span>
                
 
                            -<span class="sd">        :param cache_path:              Path to the directory where cached images will be stored in an optimized format.</span>
                
 
                            +<span class="sd">        :param cache_dir:              Path to the directory where cached images will be stored in an optimized format.</span>
                
 
                             <span class="sd">        :param transforms:              List of transforms to apply sequentially on sample.</span>
                
 
                             <span class="sd">        :param all_classes_list:        All the class names.</span>
                
 
                             <span class="sd">        :param class_inclusion_list:    If not None,every class not included will be ignored.</span>
                
@@ -208,11 +215,12 @@
 
                                     <span class="k">if</span> <span class="s2">&quot;target&quot;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">target_fields</span><span class="p">:</span>
                
 
                                         <span class="k">raise</span> <span class="ne">KeyError</span><span class="p">(</span><span class="s1">&#39;&quot;target&quot; is expected to be in the fields to subclass but it was not included&#39;</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_required_annotation_fields</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;target&quot;</span><span class="p">,</span> <span class="s2">&quot;img_path&quot;</span><span class="p">,</span> <span class="s2">&quot;resized_img_shape&quot;</span><span class="p">}</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_cache_annotations</span><span class="p">()</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">cache</span> <span class="o">=</span> <span class="n">cache</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">cache_path</span> <span class="o">=</span> <span class="n">cache_path</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">cached_imgs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_cache_images</span><span class="p">()</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache</span> <span class="k">else</span> <span class="kc">None</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">cache_dir</span> <span class="o">=</span> <span class="n">cache_dir</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">cached_imgs_padded</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_cache_images</span><span class="p">()</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache</span> <span class="k">else</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span> <span class="o">=</span> <span class="n">transforms</span>
                
@@ -231,7 +239,7 @@
 
                             <span class="sd">        Please note that the targets should be resized according to self.input_dim!</span>
                
 
                             <span class="sd">        :param sample_id:   Id of the sample to load annotations from.</span>
                
 
                            -<span class="sd">        :return:            Annotation, a dict with any field but has to include at least &quot;target&quot; and &quot;img_path&quot;.</span>
                
 
                            +<span class="sd">        :return:            Annotation, a dict with any field but has to include at least the fields specified in self._required_annotation_fields.</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="k">raise</span> <span class="ne">NotImplementedError</span>
                
@@ -246,8 +254,10 @@
 
                                             <span class="k">break</span>
                
 
                                         <span class="n">img_annotation</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_annotation</span><span class="p">(</span><span class="n">img_id</span><span class="p">)</span>
                
 
                            -            <span class="k">if</span> <span class="s2">&quot;target&quot;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">img_annotation</span> <span class="ow">or</span> <span class="s2">&quot;img_path&quot;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">img_annotation</span><span class="p">:</span>
                
 
                            -                <span class="k">raise</span> <span class="ne">KeyError</span><span class="p">(</span><span class="s1">&#39;_load_annotation is expected to return at least the field &quot;target&quot; and &quot;img_path&quot;&#39;</span><span class="p">)</span>
                
 
                            +            <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">_required_annotation_fields</span><span class="o">.</span><span class="n">issubset</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">img_annotation</span><span class="o">.</span><span class="n">keys</span><span class="p">())):</span>
                
 
                            +                <span class="k">raise</span> <span class="ne">KeyError</span><span class="p">(</span>
                
 
                            +                    <span class="sa">f</span><span class="s2">&quot;_load_annotation is expected to return at least the fields </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_required_annotation_fields</span><span class="si">}</span><span class="s2"> &quot;</span> <span class="sa">f</span><span class="s2">&quot;but got </span><span class="si">{</span><span class="nb">set</span><span class="p">(</span><span class="n">img_annotation</span><span
                
 
                            +                <span class="p">)</span>
                
 
                                         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">class_inclusion_list</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                             <span class="n">img_annotation</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_sub_class_annotation</span><span class="p">(</span><span class="n">img_annotation</span><span class="p">)</span>
                
@@ -258,8 +268,9 @@
 
                                         <span class="n">annotations</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">img_annotation</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">annotations</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="n">EmptyDatasetException</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Out of </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">n_available_samples</span><span class="si">}</span><span class="s2"> images, not a single one was found with&quot;</span>
                
 
                            -                                        <span class="sa">f</span><span class="s2">&quot;any of these classes: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">class_inclusion_list</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
                
 
                            +            <span class="k">raise</span> <span class="n">EmptyDatasetException</span><span class="p">(</span>
                
 
                            +                <span class="sa">f</span><span class="s2">&quot;Out of </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">n_available_samples</span><span class="si">}</span><span class="s2"> images, not a single one was found with&quot;</span> <span class="sa">f</span><span class="s2">&quot;any of these classes: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">class_inclusion_list</span><span cla
                
 
                            +            <span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="n">annotations</span>
                
 
                                 <span class="k">def</span> <span class="nf">_sub_class_annotation</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">annotation</span><span class="p">:</span> <span class="nb">dict</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Union</span><span class="p">[</span><span class="nb">dict</span><span class="p">,</span> <span class="kc">None</span><span class="p">]:</span>
                
@@ -296,16 +307,17 @@
 
                                     <span class="sd">&quot;&quot;&quot;Cache the images. The cached image are stored in a file to be loaded faster mext time.</span>
                
 
                             <span class="sd">        :return: Cached images</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">cache_path</span> <span class="o">=</span> <span class="n">Path</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cache_path</span><span class="p">)</span>
                
 
                            -        <span class="k">if</span> <span class="n">cache_path</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;You must specify a cache_path if you want to cache your images.&quot;</span>
                
 
                            -                             <span class="s2">&quot;If you did not mean to use cache, please set cache=False &quot;</span><span class="p">)</span>
                
 
                            -        <span class="n">cache_path</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">parents</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;</span><span class="se">\n</span><span class="s2">********************************************************************************</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -                       <span class="s2">&quot;You are using cached images in RAM to accelerate training.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -                       <span class="s2">&quot;This requires large system RAM.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            -                       <span class="s2">&quot;********************************************************************************&quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">cache_dir</span> <span class="o">=</span> <span class="n">Path</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cache_dir</span><span class="p">)</span>
                
 
                            +        <span class="k">if</span> <span class="n">cache_dir</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;You must specify a cache_dir if you want to cache your images.&quot;</span> <span class="s2">&quot;If you did not mean to use cache, please set cache=False &quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">cache_dir</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">parents</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
                
 
                            +            <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">********************************************************************************</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            +            <span class="s2">&quot;You are using cached images in RAM to accelerate training.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            +            <span class="s2">&quot;This requires large system RAM.</span><span class="se">\n</span><span class="s2">&quot;</span>
                
 
                            +            <span class="s2">&quot;********************************************************************************&quot;</span>
                
 
                            +        <span class="p">)</span>
                
 
                                     <span class="n">max_h</span><span class="p">,</span> <span class="n">max_w</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</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">input_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                
@@ -314,10 +326,10 @@
 
                                     <span class="k">for</span> <span class="n">annotation</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">:</span>
                
 
                                         <span class="n">values_to_hash</span> <span class="o">=</span> <span class="p">[</span><span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;resized_img_shape&quot;</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span> <span class="n">annotation</span><span class="p">[</span><span class="s2">&quot;resized_img_shape&quot;</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span> <span class="n">Path</span><
                
 
                                         <span class="k">for</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">values_to_hash</span><span class="p">:</span>
                
 
                            -                <span class="nb">hash</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s1">&#39;utf-8&#39;</span><span class="p">))</span>
                
 
                            +                <span class="nb">hash</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">value</span><span class="p">)</span><span class="o">.</span><span class="n">encode</span><span class="p">(</span><span class="s2">&quot;utf-8&quot;</span><span class="p">))</span>
                
 
                                     <span class="n">cache_hash</span> <span class="o">=</span> <span class="nb">hash</span><span class="o">.</span><span class="n">hexdigest</span><span class="p">()</span>
                
 
                            -        <span class="n">img_resized_cache_path</span> <span class="o">=</span> <span class="n">cache_path</span> <span class="o">/</span> <span class="sa">f</span><span class="s2">&quot;img_resized_cache_</span><span class="si">{</span><span class="n">cache_hash</span><span class="si">}</span><span class="s2">.array&quot;</span>
                
 
                            +        <span class="n">img_resized_cache_path</span> <span class="o">=</span> <span class="n">cache_dir</span> <span class="o">/</span> <span class="sa">f</span><span class="s2">&quot;img_resized_cache_</span><span class="si">{</span><span class="n">cache_hash</span><span class="si">}</span><span class="s2">.array&quot;</span>
                
 
                                     <span class="k">if</span> <span class="ow">not</span> <span class="n">img_resized_cache_path</span><span class="o">.</span><span class="n">exists</span><span class="p">():</span>
                
 
                                         <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Caching images for the first time. Be aware that this will stay in the disk until you delete it yourself.&quot;</span><span class="p">)</span>
                
@@ -325,8 +337,7 @@
 
                                         <span class="n">loaded_images</span> <span class="o">=</span> <span class="n">ThreadPool</span><span class="p">(</span><span class="n">NUM_THREADs</span><span class="p">)</span><span class="o">.</span><span class="n">imap</span><span class="p">(</span><span class="n">func</span><span class="o">=</span><span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_resized_img</span><span class="
                
 
                                         <span class="c1"># Initialize placeholder for images</span>
                
 
                            -            <span class="n">cached_imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">memmap</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">img_resized_cache_path</span><span class="p">),</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">),</span> <span class="n">max_h<
                
 
                            -                                    <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;w+&quot;</span><span class="p">)</span>
                
 
                            +            <span class="n">cached_imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">memmap</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">img_resized_cache_path</span><span class="p">),</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">),</span> <span class="n">max_h<
                
 
                                         <span class="c1"># Store images in the placeholder</span>
                
 
                                         <span class="n">loaded_images_pbar</span> <span class="o">=</span> <span class="n">tqdm</span><span class="p">(</span><span class="nb">enumerate</span><span class="p">(</span><span class="n">loaded_images</span><span class="p">),</span> <span class="n">total</span><span class="o">=</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span>
                
@@ -338,8 +349,7 @@
 
                                         <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;You are using cached imgs!&quot;</span><span class="p">)</span>
                
 
                                     <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Loading cached imgs...&quot;</span><span class="p">)</span>
                
 
                            -        <span class="n">cached_imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">memmap</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">img_resized_cache_path</span><span class="p">),</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">),</span> <span class="n">max_h</spa
                
 
                            -                                <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;r+&quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">cached_imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">memmap</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">img_resized_cache_path</span><span class="p">),</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">),</span> <span class="n">max_h</spa
                
 
                                     <span class="k">return</span> <span class="n">cached_imgs</span>
                
 
                                 <span class="k">def</span> <span class="nf">_load_resized_img</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">:</span>
                
@@ -366,14 +376,13 @@
 
                                     <span class="n">img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">imread</span><span class="p">(</span><span class="n">img_file</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="n">img</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">FileNotFoundError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">img_file</span><span class="si">}</span><span class="s2"> was no found. Please make sure that the dataset was&quot;</span>
                
 
                            -                                    <span class="sa">f</span><span class="s2">&quot;downloaded and that the path is correct&quot;</span><span class="p">)</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">FileNotFoundError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">img_file</span><span class="si">}</span><span class="s2"> was no found. Please make sure that the dataset was&quot;</span> <span class="sa">f</span><span class="s2">&quot;downloaded and that the path is correct&quot;</span><span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="n">img</span>
                
 
                                 <span class="k">def</span> <span class="fm">__del__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Clear the cached images&quot;&quot;&quot;</span>
                
 
                            -        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;cached_imgs&quot;</span><span class="p">):</span>
                
 
                            -            <span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">cached_imgs</span>
                
 
                            +        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;cached_imgs_padded&quot;</span><span class="p">):</span>
                
 
                            +            <span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">cached_imgs_padded</span>
                
 
                                 <span class="k">def</span> <span class="fm">__len__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Get the length of the dataset.&quot;&quot;&quot;</span>
                
@@ -385,34 +394,37 @@
 
                                     <span class="n">sample</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">apply_transforms</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">get_sample</span><span class="p">(</span><span class="n">index</span><span class="p">))</span>
                
 
                                     <span class="k">for</span> <span class="n">field</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_fields</span><span class="p">:</span>
                
 
                                         <span class="k">if</span> <span class="n">field</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">sample</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                            -                <span class="k">raise</span> <span class="ne">KeyError</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;The field </span><span class="si">{</span><span class="n">field</span><span class="si">}</span><span class="s1"> must be present in the sample but was not found.&#39;</span>
                
 
                            -                               <span class="s1">&#39;Please check the output fields of your transforms.&#39;</span><span class="p">)</span>
                
 
                            +                <span class="k">raise</span> <span class="ne">KeyError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The field </span><span class="si">{</span><span class="n">field</span><span class="si">}</span><span class="s2"> must be present in the sample but was not found.&quot;</span> <span class="s2">&quot;Please check the output fields of your transforms.&quot;</span><span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">sample</span><span class="p">[</span><span class="n">field</span><span class="p">]</span> <span class="k">for</span> <span class="n">field</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_fields</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="DetectionDataset.get_random_item"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset.get_random_item">[docs]</a>    <span class="k">def</span> <span class="nf">get_random_item</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="DetectionDataset.get_random_item"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset.get_random_item">[docs]</a>    <span class="k">def</span> <span class="nf">get_random_item</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="k">return</span> <span class="bp">self</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">_random_index</span><span class="p">()]</span></div>
                
 
                            -<div class="viewcode-block" id="DetectionDataset.get_sample"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset.get_sample">[docs]</a>    <span class="k">def</span> <span class="nf">get_sample</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o
                
 
                            +<div class="viewcode-block" id="DetectionDataset.get_sample"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset.get_sample">[docs]</a>    <span class="k">def</span> <span class="nf">get_sample</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="
                
 
                                     <span class="sd">&quot;&quot;&quot;Get raw sample, before any transform (beside subclassing).</span>
                
 
                             <span class="sd">        :param index:   Image index</span>
                
 
                             <span class="sd">        :return:        Sample, i.e. a dictionary including at least &quot;image&quot; and &quot;target&quot;</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="n">img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_resized_image</span><span class="p">(</span><span class="n">index</span><span class="p">)</span>
                
 
                            -        <span class="n">annotation</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>
                
 
                            +        <span class="n">annotation</span> <span class="o">=</span> <span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">[</span><span class="n">index</span><span class="p">])</span>
                
 
                                     <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;image&quot;</span><span class="p">:</span> <span class="n">img</span><span class="p">,</span> <span class="o">**</span><span class="n">annotation</span><span class="p">}</span></div>
                
 
                            -<div class="viewcode-block" id="DetectionDataset.get_resized_image"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset.get_resized_image">[docs]</a>    <span class="k">def</span> <span class="nf">get_resized_image</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">:</span> <span class="nb">int</span><span class="p">)<
                
 
                            +<div class="viewcode-block" id="DetectionDataset.get_resized_image"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset.get_resized_image">[docs]</a>    <span class="k">def</span> <span class="nf">get_resized_image</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        Get the resized image at a specific sample_id, either from cache or by loading from disk, based on self.cached_imgs</span>
                
 
                            +<span class="sd">        Get the resized image (i.e. either width or height reaches its input_dim) at a specific sample_id,</span>
                
 
                            +<span class="sd">        either from cache or by loading from disk, based on self.cached_imgs_padded</span>
                
 
                             <span class="sd">        :param index:  Image index</span>
                
 
                             <span class="sd">        :return:       Resized image</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cache</span><span class="p">:</span>
                
 
                            -            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">cached_imgs</span><span class="p">[</span><span class="n">index</span><span class="p">]</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                
 
                            +            <span class="n">padded_image</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cached_imgs_padded</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>
                
 
                            +            <span class="n">resized_height</span><span class="p">,</span> <span class="n">resized_width</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">annotations</span><span class="p">[</span><span class="n">index</span><span class="p">][</span><span class="s2">&quot;resized_img_shape&quot;</span><span class="p">]</span>
                
 
                            +            <span class="n">resized_image</span> <span class="o">=</span> <span class="n">padded_image</span><span class="p">[:</span><span class="n">resized_height</span><span class="p">,</span> <span class="p">:</span><span class="n">resized_width</span><span class="p">,</span> <span class="p">:]</span>
                
 
                            +            <span class="k">return</span> <span class="n">resized_image</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_resized_img</span><span class="p">(</span><span class="n">index</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="DetectionDataset.apply_transforms"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset.apply_transforms">[docs]</a>    <span class="k">def</span> <span class="nf">apply_transforms</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</s
                
 
                            +<div class="viewcode-block" id="DetectionDataset.apply_transforms"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset.apply_transforms">[docs]</a>    <span class="k">def</span> <span class="nf">apply_transforms</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</sp
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Applies self.transforms sequentially to sample</span>
                
@@ -429,17 +441,14 @@
 
                                         <span class="n">sample</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;additional_samples&quot;</span><span class="p">)</span>  <span class="c1"># additional_samples is not useful after the transform</span>
                
 
                                     <span class="k">return</span> <span class="n">sample</span></div>
                
 
                            -    <span class="k">def</span> <span class="nf">_add_additional_inputs_for_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Union</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">An
                
 
                            -                                             <span class="n">transform</span><span class="p">:</span> <span class="n">DetectionTransform</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="nf">_add_additional_inputs_for_transform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Union</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">An
                
 
                                     <span class="sd">&quot;&quot;&quot;Add additional inputs required by a transform to the sample&quot;&quot;&quot;</span>
                
 
                            -        <span class="n">additional_samples_count</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">additional_samples_count</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span>
                
 
                            -                                                                                 <span class="s2">&quot;additional_samples_count&quot;</span><span class="p">)</span> <span class="k">else</span> <span class="mi">0</span>
                
 
                            +        <span class="n">additional_samples_count</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">additional_samples_count</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span> <span class="s2">&quot;additional_samples_count&quot;</span><span class="p">)</span> <span class="k">else</span> <span class="mi">0</span>
                
 
                                     <span class="n">non_empty_annotations</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">non_empty_annotations</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span> <span class="s2">&quot;non_empty_annotations&quot;</span><span class="p">)</span> <span class="k">else</span> <span class="kc">False</span>
                
 
                                     <span class="n">additional_samples</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_random_samples</span><span class="p">(</span><span class="n">additional_samples_count</span><span class="p">,</span> <span class="n">non_empty_annotations</span><span class="p">)</span>
                
 
                                     <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;additional_samples&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">additional_samples</span>
                
 
                            -<div class="viewcode-block" id="DetectionDataset.get_random_samples"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset.get_random_samples">[docs]</a>    <span class="k">def</span> <span class="nf">get_random_samples</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="p">:</span> <span class="nb">int</span><span class="p"
                
 
                            -                           <span class="n">non_empty_annotations_only</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Union</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span 
                
 
                            +<div class="viewcode-block" id="DetectionDataset.get_random_samples"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset.get_random_samples">[docs]</a>    <span class="k">def</span> <span class="nf">get_random_samples</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">count</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">no
                
 
                                     <span class="sd">&quot;&quot;&quot;Load random samples.</span>
                
 
                             <span class="sd">        :param count: The number of samples wanted</span>
                
@@ -448,7 +457,7 @@
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="k">return</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">get_random_sample</span><span class="p">(</span><span class="n">non_empty_annotations_only</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="n">count</span><span class="p">)]</span></div>
                
 
                            -<div class="viewcode-block" id="DetectionDataset.get_random_sample"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset.get_random_sample">[docs]</a>    <span class="k">def</span> <span class="nf">get_random_sample</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">non_empty_annotations_only</span><span class="p">:</span> <span class="nb">bool</s
                
 
                            +<div class="viewcode-block" id="DetectionDataset.get_random_sample"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset.get_random_sample">[docs]</a>    <span class="k">def</span> <span class="nf">get_random_sample</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">non_empty_annotations_only</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span
                
 
                                     <span class="k">if</span> <span class="n">non_empty_annotations_only</span><span class="p">:</span>
                
 
                                         <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_sample</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_random_non_empty_annotation_available_indexes</span><span class="p">())</span>
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
@@ -475,23 +484,22 @@
 
                                             <span class="n">target_format</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">output_format</span>
                
 
                                     <span class="k">return</span> <span class="n">target_format</span>
                
 
                            -<div class="viewcode-block" id="DetectionDataset.plot"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.DetectionDataset.plot">[docs]</a>    <span class="k">def</span> <span class="nf">plot</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">max_samples_per_plot</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi"
                
 
                            +<div class="viewcode-block" id="DetectionDataset.plot"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.DetectionDataset.plot">[docs]</a>    <span class="k">def</span> <span class="nf">plot</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">max_samples_per_plot</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">16</span><span class="p">,<
                
 
                                     <span class="sd">&quot;&quot;&quot;Combine samples of images with bbox into plots and display the result.</span>
                
 
                            -<span class="sd">            :param max_samples_per_plot:    Maximum number of images to be displayed per plot</span>
                
 
                            -<span class="sd">            :param n_plots:                 Number of plots to display (each plot being a combination of img with bbox)</span>
                
 
                            -<span class="sd">            :param plot_transformed_data:   If True, the plot will be over samples after applying transforms (i.e. on __getitem__).</span>
                
 
                            -<span class="sd">                                            If False, the plot will be over the raw samples (i.e. on get_sample)</span>
                
 
                            -<span class="sd">            :return:</span>
                
 
                            +<span class="sd">        :param max_samples_per_plot:    Maximum number of images to be displayed per plot</span>
                
 
                            +<span class="sd">        :param n_plots:                 Number of plots to display (each plot being a combination of img with bbox)</span>
                
 
                            +<span class="sd">        :param plot_transformed_data:   If True, the plot will be over samples after applying transforms (i.e. on __getitem__).</span>
                
 
                            +<span class="sd">                                        If False, the plot will be over the raw samples (i.e. on get_sample)</span>
                
 
                            +<span class="sd">        :return:</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="n">plot_counter</span> <span class="o">=</span> <span class="mi">0</span>
                
 
                                     <span class="n">input_format</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_target_format</span> <span class="k">if</span> <span class="n">plot_transformed_data</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">original_target_format</span>
                
 
                            -        <span class="n">target_format_transform</span> <span class="o">=</span> <span class="n">DetectionTargetsFormatTransform</span><span class="p">(</span><span class="n">input_format</span><span class="o">=</span><span class="n">input_format</span><span class="p">,</span>
                
 
                            -                                                                  <span class="n">output_format</span><span class="o">=</span><span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">XYXY_LABEL</span><span class="p">)</span>
                
 
                            +        <span class="n">target_format_transform</span> <span class="o">=</span> <span class="n">DetectionTargetsFormatTransform</span><span class="p">(</span><span class="n">input_format</span><span class="o">=</span><span class="n">input_format</span><span class="p">,</span> <span class="n">output_format</span><span class="o">=</span><span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">XYXY_LABEL</span><span class="p">)</span>
                
 
                                     <span class="k">for</span> <span class="n">plot_i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">n_plots</span><span class="p">):</span>
                
 
                                         <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><span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">10</span><span class="p">,</span> <span class="mi">10</span><span class="p">))</span>
                
 
                            -            <span class="n">n_subplot</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">max_samples_per_plot</span> <span class="o">**</span> <span class="mf">0.5</span><span class="p">))</span>
                
 
                            +            <span class="n">n_subplot</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">max_samples_per_plot</span><span class="o">**</span><span class="mf">0.5</span><span class="p">))</span>
                
 
                                         <span class="k">for</span> <span class="n">img_i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">max_samples_per_plot</span><span class="p">):</span>
                
 
                                             <span class="n">index</span> <span class="o">=</span> <span class="n">img_i</span> <span class="o">+</span> <span class="n">plot_i</span> <span class="o">*</span> <span class="mi">16</span>
                
@@ -509,9 +517,9 @@
 
                                             <span class="c1"># shape = [n_box x 4] (We remove padded boxes, which corresponds to boxes with only 0)</span>
                
 
                                             <span class="n">boxes</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[(</span><span class="n">boxes</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)]</span>
                
 
                            -                <span class="n">plt</span><span class="o">.</span><span class="n">subplot</span><span class="p">(</span><span class="n">n_subplot</span><span class="p">,</span> <span class="n">n_subplot</span><span class="p">,</span> <span class="n">img_i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">imshow</span><span class="p">(</span><span class="n">image</span><span class="p">)</span>
                
 
                            -                <span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">boxes</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</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="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]]</span><span class="o">.</span><span class="n">T</span><span
                
 
                            -                <span class="n">plt</span><span class="o">.</span><span class="n">axis</span><span class="p">(</span><span class="s1">&#39;off&#39;</span><span class="p">)</span>
                
 
                            +                <span class="n">plt</span><span class="o">.</span><span class="n">subplot</span><span class="p">(</span><span class="n">n_subplot</span><span class="p">,</span> <span class="n">n_subplot</span><span class="p">,</span> <span class="n">img_i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">imshow</span><span class="p">(</span><span class="n">image</span><span class="p">[:,</span> <span class="p">:,</span> <spa
                
 
                            +                <span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">boxes</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</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="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]]</span><span class="o">.</span><span class="n">T</span><span
                
 
                            +                <span class="n">plt</span><span class="o">.</span><span class="n">axis</span><span class="p">(</span><span class="s2">&quot;off&quot;</span><span class="p">)</span>
                
 
                                         <span class="n">fig</span><span class="o">.</span><span class="n">tight_layout</span><span class="p">()</span>
                
 
                                         <span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
                
 
                                         <span class="n">plt</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
                
@@ -548,4 +556,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.detection_datasets.pascal_voc_detection &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.detection_datasets.pascal_voc_detection &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -89,11 +91,15 @@
 
                             <span></span><span class="kn">import</span> <span class="nn">os</span>
                
 
                             <span class="kn">import</span> <span class="nn">glob</span>
                
 
                             <span class="kn">from</span> <span class="nn">pathlib</span> <span class="kn">import</span> <span class="n">Path</span>
                
 
                            +<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span>
                
 
                             <span class="kn">from</span> <span class="nn">xml.etree</span> <span class="kn">import</span> <span class="n">ElementTree</span>
                
 
                            +
                
 
                            +<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">ConcatDataset</span>
                
 
                             <span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                             <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.transforms.transforms</span> <span class="kn">import</span> <span class="n">DetectionTransform</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">download_and_untar_from_url</span><span class="p">,</span> <span class="n">get_image_size_from_path</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.datasets.detection_datasets.detection_dataset</span> <span class="kn">import</span> <span class="n">DetectionDataset</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">DetectionTargetsFormat</span>
                
@@ -103,19 +109,25 @@
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="PascalVOCDetectionDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.PascalVOCDetectionDataset">[docs]</a><span class="k">class</span> <span class="nc">PascalVOCDetectionDataset</span><span class="p">(</span><span class="n">DetectionDataset</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="PascalVOCDetectionDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalVOCDetectionDataset">[docs]</a><span class="k">class</span> <span class="nc">PascalVOCDetectionDataset</span><span class="p">(</span><span class="n">DetectionDataset</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Dataset for Pascal VOC object detection&quot;&quot;&quot;</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">images_sub_directory</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">images_sub_directory</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">download</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Dataset for Pascal VOC object detection</span>
                
 
                             <span class="sd">        :param images_sub_directory:    Sub directory of data_dir that includes images.</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">images_sub_directory</span> <span class="o">=</span> <span class="n">images_sub_directory</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">img_and_target_path_list</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -
                
 
                            -        <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;all_classes_list&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">PASCAL_VOC_2012_CLASSES_LIST</span>
                
 
                            -        <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;original_target_format&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">XYXY_LABEL</span>
                
 
                            +        <span class="n">data_dir</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;data_dir&quot;</span><span class="p">)</span>
                
 
                            +        <span class="k">if</span> <span class="n">data_dir</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Must pass data_dir != None through **kwargs&quot;</span><span class="p">)</span>
                
 
                            +        <span class="k">if</span> <span class="n">download</span><span class="p">:</span>
                
 
                            +            <span class="n">PascalVOCDetectionDataset</span><span class="o">.</span><span class="n">download</span><span class="p">(</span><span class="n">data_dir</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;original_target_format&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">XYXY_LABEL</span>
                
 
                            +        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;all_classes_list&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">PASCAL_VOC_2012_CLASSES_LIST</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                                 <span class="k">def</span> <span class="nf">_setup_data_source</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
@@ -125,9 +137,11 @@
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="n">img_files_folder</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">images_sub_directory</span>
                
 
                                     <span class="k">if</span> <span class="ow">not</span> <span class="n">Path</span><span class="p">(</span><span class="n">img_files_folder</span><span class="p">)</span><span class="o">.</span><span class="n">exists</span><span class="p">():</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">FileNotFoundError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="si">}</span><span class="s2"> does not include </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">images_sub_directory</span><span class="si">}</span><span class="s2">.
                
 
                            -                                    <span class="sa">f</span><span class="s2">&quot;Please make sure that f</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="si">}</span><span class="s2"> refers to PascalVOC dataset and that &quot;</span>
                
 
                            -                                    <span class="s2">&quot;it was downloaded using PascalVOCDetectionDataSetV2.download()&quot;</span><span class="p">)</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">FileNotFoundError</span><span class="p">(</span>
                
 
                            +                <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="si">}</span><span class="s2"> does not include </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">images_sub_directory</span><span class="si">}</span><span class="s2">. &quot;</span>
                
 
                            +                <span class="sa">f</span><span class="s2">&quot;Please make sure that f</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">data_dir</span><span class="si">}</span><span class="s2"> refers to PascalVOC dataset and that &quot;</span>
                
 
                            +                <span class="s2">&quot;it was downloaded using PascalVOCDetectionDataSetV2.download()&quot;</span>
                
 
                            +            <span class="p">)</span>
                
 
                                     <span class="n">img_files</span> <span class="o">=</span> <span class="n">glob</span><span class="o">.</span><span class="n">glob</span><span class="p">(</span><span class="n">img_files_folder</span> <span class="o">+</span> <span class="s2">&quot;*.jpg&quot;</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">img_files</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                
@@ -135,15 +149,13 @@
 
                                     <span class="n">target_files</span> <span class="o">=</span> <span class="p">[</span><span class="n">img_file</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot;images&quot;</span><span class="p">,</span> <span class="s2">&quot;labels&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot;.jpg&quot;</span><span class="p">,</span> <span class="s2">&quo
                
 
                            -        <span class="n">img_and_target_path_list</span> <span class="o">=</span> <span class="p">[(</span><span class="n">img_file</span><span class="p">,</span> <span class="n">target_file</span><span class="p">)</span>
                
 
                            -                                    <span class="k">for</span> <span class="n">img_file</span><span class="p">,</span> <span class="n">target_file</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">img_files</span><span class="p">,</span> <span class="n">target_files</span><span class="p">)</span>
                
 
                            -                                    <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">target_file</span><span class="p">)]</span>
                
 
                            +        <span class="n">img_and_target_path_list</span> <span class="o">=</span> <span class="p">[(</span><span class="n">img_file</span><span class="p">,</span> <span class="n">target_file</span><span class="p">)</span> <span class="k">for</span> <span class="n">img_file</span><span class="p">,</span> <span class="n">target_file</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">img_files</span><span class="p">,</span> <span class="n">target_fi
                
 
                                     <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">img_and_target_path_list</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                
 
                                         <span class="k">raise</span> <span class="ne">FileNotFoundError</span><span class="p">(</span><span class="s2">&quot;No target file associated to the images was found&quot;</span><span class="p">)</span>
                
 
                                     <span class="n">num_missing_files</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">img_files</span><span class="p">)</span> <span class="o">-</span> <span class="nb">len</span><span class="p">(</span><span class="n">img_and_target_path_list</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="n">num_missing_files</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                
 
                            -            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">num_missing_files</span><span class="si">}</span><span class="s1"> labels files were not loaded our of </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">img_files</span><span class="p">)</span><span class="si">}</span><span class="s1"> 
                
 
                            +            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">num_missing_files</span><span class="si">}</span><span class="s2"> labels files were not loaded our of </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">img_files</span><span class="p">)</span><span class="si">}</span><span class="s2">
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">img_and_target_path_list</span> <span class="o">=</span> <span class="n">img_and_target_path_list</span>
                
 
                                     <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">img_and_target_path_list</span><span class="p">)</span>
                
@@ -156,38 +168,37 @@
 
                             <span class="sd">                    - img_path</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="n">img_path</span><span class="p">,</span> <span class="n">target_path</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">img_and_target_path_list</span><span class="p">[</span><span class="n">sample_id</span><span class="p">]</span>
                
 
                            -        <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">target_path</span><span class="p">,</span> <span class="s1">&#39;r&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">targets_file</span><span class="p">:</span>
                
 
                            +        <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">target_path</span><span class="p">,</span> <span class="s2">&quot;r&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">targets_file</span><span class="p">:</span>
                
 
                                         <span class="n">target</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="n">x</span><span class="o">.</span><span class="n">split</span><span class="p">()</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">targets_file</span><span class="o">.</span><span class="n">read</span><span class="p">()</span><span class="o">.</span><span class="n
                
 
                            -        <span class="n">width</span><span class="p">,</span> <span class="n">height</span> <span class="o">=</span> <span class="n">get_image_size_from_path</span><span class="p">(</span><span class="n">img_path</span><span class="p">)</span>
                
 
                            +        <span class="n">height</span><span class="p">,</span> <span class="n">width</span> <span class="o">=</span> <span class="n">get_image_size_from_path</span><span class="p">(</span><span class="n">img_path</span><span class="p">)</span>
                
 
                            -        <span class="c1"># We have to rescale the targets because the images will be rescaled.</span>
                
 
                            +        <span class="c1"># We have to rescale the targets because the images will be resized.</span>
                
 
                                     <span class="n">r</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">height</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">0</span><span class="p">]</s
                
 
                                     <span class="n">target</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">r</span>
                
 
                            -        <span class="n">initial_img_shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">width</span><span class="p">,</span> <span class="n">height</span><span class="p">)</span>
                
 
                            -        <span class="n">resized_img_shape</span> <span class="o">=</span> <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">width</span> <span class="o">*</span> <span class="n">r</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">height</span> <span class="o">*</span> <span class="n">r</span><span class="p">))</span>
                
 
                            +        <span class="n">resized_img_shape</span> <span class="o">=</span> <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">height</span> <span class="o">*</span> <span class="n">r</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">width</span> <span class="o">*</span> <span class="n">r</span><span class="p">))</span>
                
 
                            -        <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;img_path&quot;</span><span class="p">:</span> <span class="n">img_path</span><span class="p">,</span> <span class="s2">&quot;target&quot;</span><span class="p">:</span> <span class="n">target</span><span class="p">,</span>
                
 
                            -                <span class="s2">&quot;initial_img_shape&quot;</span><span class="p">:</span> <span class="n">initial_img_shape</span><span class="p">,</span> <span class="s2">&quot;resized_img_shape&quot;</span><span class="p">:</span> <span class="n">resized_img_shape</span><span class="p">}</span>
                
 
                            +        <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;img_path&quot;</span><span class="p">:</span> <span class="n">img_path</span><span class="p">,</span> <span class="s2">&quot;target&quot;</span><span class="p">:</span> <span class="n">target</span><span class="p">,</span> <span class="s2">&quot;resized_img_shape&quot;</span><span class="p">:</span> <span class="n">resized_img_shape</span><span class="p">}</span>
                
 
                            -<div class="viewcode-block" id="PascalVOCDetectionDataset.download"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.detection_datasets.html#super_gradients.training.datasets.PascalVOCDetectionDataset.download">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="PascalVOCDetectionDataset.download"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalVOCDetectionDataset.download">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">download</span><span class="p">(</span><span class="n">data_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Download Pascal dataset in XYXY_LABEL format.</span>
                
 
                             <span class="sd">        Data extracted form http://host.robots.ox.ac.uk/pascal/VOC/</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +
                
 
                                     <span class="k">def</span> <span class="nf">_parse_and_save_labels</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">new_label_path</span><span class="p">,</span> <span class="n">year</span><span class="p">,</span> <span class="n">image_id</span><span class="p">):</span>
                
 
                                         <span class="sd">&quot;&quot;&quot;Parse and save the labels of an image in XYXY_LABEL format.&quot;&quot;&quot;</span>
                
 
                            -            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">path</span><span class="si">}</span><span class="s1">/VOC</span><span class="si">{</span><span class="n">year</span><span class="si">}</span><span class="s1">/Annotations/</span><span class="si">{</span><span class="n">image_id</span><span class="si">}</span><span class="s1">.xml&#39;</span><span class="p">)</s
                
 
                            +            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">path</span><span class="si">}</span><span class="s2">/VOC</span><span class="si">{</span><span class="n">year</span><span class="si">}</span><span class="s2">/Annotations/</span><span class="si">{</span><span class="n">image_id</span><span class="si">}</span><span class="s2">.xml&quot;</span><span class="p">)<
                
 
                                             <span class="n">xml_parser</span> <span class="o">=</span> <span class="n">ElementTree</span><span class="o">.</span><span class="n">parse</span><span class="p">(</span><span class="n">f</span><span class="p">)</span><span class="o">.</span><span class="n">getroot</span><span class="p">()</span>
                
 
                                         <span class="n">labels</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                            -            <span class="k">for</span> <span class="n">obj</span> <span class="ow">in</span> <span class="n">xml_parser</span><span class="o">.</span><span class="n">iter</span><span class="p">(</span><span class="s1">&#39;object&#39;</span><span class="p">):</span>
                
 
                            -                <span class="bp">cls</span> <span class="o">=</span> <span class="n">obj</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s1">&#39;name&#39;</span><span class="p">)</span><span class="o">.</span><span class="n">text</span>
                
 
                            -                <span class="k">if</span> <span class="bp">cls</span> <span class="ow">in</span> <span class="n">PASCAL_VOC_2012_CLASSES_LIST</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">int</span><span class="p">(</span><span class="n">obj</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s1">&#39;difficult&#39;</span><span class="p">)</span><span class="o">.</span><span class="n">text</span><span class="p">)</span>
                
 
                            -                    <span class="n">xml_box</span> <span class="o">=</span> <span class="n">obj</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s1">&#39;bndbox&#39;</span><span class="p">)</span>
                
 
                            +            <span class="k">for</span> <span class="n">obj</span> <span class="ow">in</span> <span class="n">xml_parser</span><span class="o">.</span><span class="n">iter</span><span class="p">(</span><span class="s2">&quot;object&quot;</span><span class="p">):</span>
                
 
                            +                <span class="bp">cls</span> <span class="o">=</span> <span class="n">obj</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s2">&quot;name&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">text</span>
                
 
                            +                <span class="k">if</span> <span class="bp">cls</span> <span class="ow">in</span> <span class="n">PASCAL_VOC_2012_CLASSES_LIST</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">int</span><span class="p">(</span><span class="n">obj</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s2">&quot;difficult&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">text</span><span class="p">)</spa
                
 
                            +                    <span class="n">xml_box</span> <span class="o">=</span> <span class="n">obj</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="s2">&quot;bndbox&quot;</span><span class="p">)</span>
                
 
                                                 <span class="k">def</span> <span class="nf">get_coord</span><span class="p">(</span><span class="n">box_coord</span><span class="p">):</span>
                
 
                                                     <span class="k">return</span> <span class="n">xml_box</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">box_coord</span><span class="p">)</span><span class="o">.</span><span class="n">text</span>
                
@@ -195,33 +206,72 @@
 
                                                 <span class="n">xmin</span><span class="p">,</span> <span class="n">ymin</span><span class="p">,</span> <span class="n">xmax</span><span class="p">,</span> <span class="n">ymax</span> <span class="o">=</span> <span class="n">get_coord</span><span class="p">(</span><span class="s2">&quot;xmin&quot;</span><span class="p">),</span> <span class="n">get_coord</span><span class="p">(</span><span class="s2">&quot;ymin&quot;</span><span class="p">),</span> <span class="n">get_coord<
                
 
                                                 <span class="n">labels</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s2">&quot; &quot;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="n">xmin</span><span class="p">,</span> <span class="n">ymin</span><span class="p">,</span> <span class="n">xmax</span><span class="p">,</span> <span class="n">ymax</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><
                
 
                            -            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">new_label_path</span><span class="p">,</span> <span class="s1">&#39;w&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
                
 
                            +            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">new_label_path</span><span class="p">,</span> <span class="s2">&quot;w&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
                
 
                                             <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">labels</span><span class="p">))</span>
                
 
                            -        <span class="n">urls</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar&quot;</span><span class="p">,</span>  <span class="c1"># 439M 5011 images</span>
                
 
                            -                <span class="s2">&quot;http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar&quot;</span><span class="p">,</span>  <span class="c1"># 430M, 4952 images</span>
                
 
                            -                <span class="s2">&quot;http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar&quot;</span><span class="p">]</span>  <span class="c1"># 1.86G, 17125 images</span>
                
 
                            +        <span class="n">urls</span> <span class="o">=</span> <span class="p">[</span>
                
 
                            +            <span class="s2">&quot;http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar&quot;</span><span class="p">,</span>  <span class="c1"># 439M 5011 images</span>
                
 
                            +            <span class="s2">&quot;http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar&quot;</span><span class="p">,</span>  <span class="c1"># 430M, 4952 images</span>
                
 
                            +            <span class="s2">&quot;http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar&quot;</span><span class="p">,</span>
                
 
                            +        <span class="p">]</span>  <span class="c1"># 1.86G, 17125 images</span>
                
 
                                     <span class="n">data_dir</span> <span class="o">=</span> <span class="n">Path</span><span class="p">(</span><span class="n">data_dir</span><span class="p">)</span>
                
 
                            -        <span class="n">download_and_untar_from_url</span><span class="p">(</span><span class="n">urls</span><span class="p">,</span> <span class="nb">dir</span><span class="o">=</span><span class="n">data_dir</span> <span class="o">/</span> <span class="s1">&#39;images&#39;</span><span class="p">)</span>
                
 
                            +        <span class="n">download_and_untar_from_url</span><span class="p">(</span><span class="n">urls</span><span class="p">,</span> <span class="nb">dir</span><span class="o">=</span><span class="n">data_dir</span> <span class="o">/</span> <span class="s2">&quot;images&quot;</span><span class="p">)</span>
                
 
                                     <span class="c1"># Convert</span>
                
 
                            -        <span class="n">data_path</span> <span class="o">=</span> <span class="n">data_dir</span> <span class="o">/</span> <span class="s1">&#39;images&#39;</span> <span class="o">/</span> <span class="s1">&#39;VOCdevkit&#39;</span>
                
 
                            -        <span class="k">for</span> <span class="n">year</span><span class="p">,</span> <span class="n">image_set</span> <span class="ow">in</span> <span class="p">(</span><span class="s1">&#39;2012&#39;</span><span class="p">,</span> <span class="s1">&#39;train&#39;</span><span class="p">),</span> <span class="p">(</span><span class="s1">&#39;2012&#39;</span><span class="p">,</span> <span class="s1">&#39;val&#39;</span><span class="p">),</span> <span class="p">(</span><span class="s1">&#39;2007
                
 
                            -            <span class="n">dest_imgs_path</span> <span class="o">=</span> <span class="n">data_dir</span> <span class="o">/</span> <span class="s1">&#39;images&#39;</span> <span class="o">/</span> <span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">image_set</span><span class="si">}{</span><span class="n">year</span><span class="si">}</span><span class="s1">&#39;</span>
                
 
                            +        <span class="n">data_path</span> <span class="o">=</span> <span class="n">data_dir</span> <span class="o">/</span> <span class="s2">&quot;images&quot;</span> <span class="o">/</span> <span class="s2">&quot;VOCdevkit&quot;</span>
                
 
                            +        <span class="k">for</span> <span class="n">year</span><span class="p">,</span> <span class="n">image_set</span> <span class="ow">in</span> <span class="p">(</span><span class="s2">&quot;2012&quot;</span><span class="p">,</span> <span class="s2">&quot;train&quot;</span><span class="p">),</span> <span class="p">(</span><span class="s2">&quot;2012&quot;</span><span class="p">,</span> <span class="s2">&quot;val&quot;</span><span class="p">),</span> <span class="p">(</span><span class="s2">&
                
 
                            +            <span class="n">dest_imgs_path</span> <span class="o">=</span> <span class="n">data_dir</span> <span class="o">/</span> <span class="s2">&quot;images&quot;</span> <span class="o">/</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">image_set</span><span class="si">}{</span><span class="n">year</span><span class="si">}</span><span class="s2">&quot;</span>
                
 
                                         <span class="n">dest_imgs_path</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">parents</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            -            <span class="n">dest_labels_path</span> <span class="o">=</span> <span class="n">data_dir</span> <span class="o">/</span> <span class="s1">&#39;labels&#39;</span> <span class="o">/</span> <span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">image_set</span><span class="si">}{</span><span class="n">year</span><span class="si">}</span><span class="s1">&#39;</span>
                
 
                            +            <span class="n">dest_labels_path</span> <span class="o">=</span> <span class="n">data_dir</span> <span class="o">/</span> <span class="s2">&quot;labels&quot;</span> <span class="o">/</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">image_set</span><span class="si">}{</span><span class="n">year</span><span class="si">}</span><span class="s2">&quot;</span>
                
 
                                         <span class="n">dest_labels_path</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">parents</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            -            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">data_path</span> <span class="o">/</span> <span class="sa">f</span><span class="s1">&#39;VOC</span><span class="si">{</span><span class="n">year</span><span class="si">}</span><span class="s1">/ImageSets/Main/</span><span class="si">{</span><span class="n">image_set</span><span class="si">}</span><span class="s1">.txt&#39;</span><span class="p">)</span> <span class="k">as</span> <span cla
                
 
                            +            <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">data_path</span> <span class="o">/</span> <span class="sa">f</span><span class="s2">&quot;VOC</span><span class="si">{</span><span class="n">year</span><span class="si">}</span><span class="s2">/ImageSets/Main/</span><span class="si">{</span><span class="n">image_set</span><span class="si">}</span><span class="s2">.txt&quot;</span><span class="p">)</span> <span class="k">as</span> <span c
                
 
                                             <span class="n">image_ids</span> <span class="o">=</span> <span class="n">f</span><span class="o">.</span><span class="n">read</span><span class="p">()</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span><span class="o">.</span><span class="n">split</span><span class="p">()</span>
                
 
                            -            <span class="k">for</span> <span class="nb">id</span> <span class="ow">in</span> <span class="n">tqdm</span><span class="p">(</span><span class="n">image_ids</span><span class="p">,</span> <span class="n">desc</span><span class="o">=</span><span class="sa">f</span><span class="s1">&#39;</span><span class="si">{</span><span class="n">image_set</span><span class="si">}{</span><span class="n">year</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">):</span>
                
 
                            -                <span class="n">img_path</span> <span class="o">=</span> <span class="n">data_path</span> <span class="o">/</span> <span class="sa">f</span><span class="s1">&#39;VOC</span><span class="si">{</span><span class="n">year</span><span class="si">}</span><span class="s1">/JPEGImages/</span><span class="si">{</span><span class="nb">id</span><span class="si">}</span><span class="s1">.jpg&#39;</span>
                
 
                            +            <span class="k">for</span> <span class="nb">id</span> <span class="ow">in</span> <span class="n">tqdm</span><span class="p">(</span><span class="n">image_ids</span><span class="p">,</span> <span class="n">desc</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">image_set</span><span class="si">}{</span><span class="n">year</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">):</span>
                
 
                            +                <span class="n">img_path</span> <span class="o">=</span> <span class="n">data_path</span> <span class="o">/</span> <span class="sa">f</span><span class="s2">&quot;VOC</span><span class="si">{</span><span class="n">year</span><span class="si">}</span><span class="s2">/JPEGImages/</span><span class="si">{</span><span class="nb">id</span><span class="si">}</span><span class="s2">.jpg&quot;</span>
                
 
                                             <span class="n">new_img_path</span> <span class="o">=</span> <span class="n">dest_imgs_path</span> <span class="o">/</span> <span class="n">img_path</span><span class="o">.</span><span class="n">name</span>
                
 
                            -                <span class="n">new_label_path</span> <span class="o">=</span> <span class="p">(</span><span class="n">dest_labels_path</span> <span class="o">/</span> <span class="n">img_path</span><span class="o">.</span><span class="n">name</span><span class="p">)</span><span class="o">.</span><span class="n">with_suffix</span><span class="p">(</span><span class="s1">&#39;.txt&#39;</span><span class="p">)</span>
                
 
                            +                <span class="n">new_label_path</span> <span class="o">=</span> <span class="p">(</span><span class="n">dest_labels_path</span> <span class="o">/</span> <span class="n">img_path</span><span class="o">.</span><span class="n">name</span><span class="p">)</span><span class="o">.</span><span class="n">with_suffix</span><span class="p">(</span><span class="s2">&quot;.txt&quot;</span><span class="p">)</span>
                
 
                                             <span class="n">img_path</span><span class="o">.</span><span class="n">rename</span><span class="p">(</span><span class="n">new_img_path</span><span class="p">)</span>  <span class="c1"># Move image to dest folder</span>
                
 
                                             <span class="n">_parse_and_save_labels</span><span class="p">(</span><span class="n">data_path</span><span class="p">,</span> <span class="n">new_label_path</span><span class="p">,</span> <span class="n">year</span><span class="p">,</span> <span class="nb">id</span><span class="p">)</span></div></div>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">class</span> <span class="nc">PascalVOCUnifiedDetectionTrainDataset</span><span class="p">(</span><span class="n">ConcatDataset</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">data_dir</span><span class="p">,</span>
                
 
                            +        <span class="n">input_dim</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span>
                
 
                            +        <span class="n">cache</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">cache_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">transforms</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">DetectionTransform</span><span class="p">]</span> <span class="o">=</span> <span class="p">[],</span>
                
 
                            +        <span class="n">class_inclusion_list</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">max_num_samples</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">download</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                            +        <span class="k">if</span> <span class="n">download</span><span class="p">:</span>
                
 
                            +            <span class="n">PascalVOCDetectionDataset</span><span class="o">.</span><span class="n">download</span><span class="p">(</span><span class="n">data_dir</span><span class="o">=</span><span class="n">data_dir</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">train_dataset_names</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;train2007&quot;</span><span class="p">,</span> <span class="s2">&quot;val2007&quot;</span><span class="p">,</span> <span class="s2">&quot;train2012&quot;</span><span class="p">,</span> <span class="s2">&quot;val2012&quot;</span><span class="p">]</span>
                
 
                            +        <span class="c1"># We divide train_max_num_samples between the datasets</span>
                
 
                            +        <span class="k">if</span> <span class="n">max_num_samples</span><span class="p">:</span>
                
 
                            +            <span class="n">max_num_samples_per_train_dataset</span> <span class="o">=</span> <span class="p">[</span><span class="nb">len</span><span class="p">(</span><span class="n">segment</span><span class="p">)</span> <span class="k">for</span> <span class="n">segment</span> <span class="ow">in</span> <span class="n">np</span><span class="o">.</span><span class="n">array_split</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">max_num_sample
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="n">max_num_samples_per_train_dataset</span> <span class="o">=</span> <span class="p">[</span><span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">train_dataset_names</span><span class="p">)</span>
                
 
                            +        <span class="n">train_sets</span> <span class="o">=</span> <span class="p">[</span>
                
 
                            +            <span class="n">PascalVOCDetectionDataset</span><span class="p">(</span>
                
 
                            +                <span class="n">data_dir</span><span class="o">=</span><span class="n">data_dir</span><span class="p">,</span>
                
 
                            +                <span class="n">input_dim</span><span class="o">=</span><span class="n">input_dim</span><span class="p">,</span>
                
 
                            +                <span class="n">cache</span><span class="o">=</span><span class="n">cache</span><span class="p">,</span>
                
 
                            +                <span class="n">cache_dir</span><span class="o">=</span><span class="n">cache_dir</span><span class="p">,</span>
                
 
                            +                <span class="n">transforms</span><span class="o">=</span><span class="n">transforms</span><span class="p">,</span>
                
 
                            +                <span class="n">images_sub_directory</span><span class="o">=</span><span class="s2">&quot;images/&quot;</span> <span class="o">+</span> <span class="n">trainset_name</span> <span class="o">+</span> <span class="s2">&quot;/&quot;</span><span class="p">,</span>
                
 
                            +                <span class="n">class_inclusion_list</span><span class="o">=</span><span class="n">class_inclusion_list</span><span class="p">,</span>
                
 
                            +                <span class="n">max_num_samples</span><span class="o">=</span><span class="n">max_num_samples_per_train_dataset</span><span class="p">[</span><span class="n">i</span><span class="p">],</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +            <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">trainset_name</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">train_dataset_names</span><span class="p">)</span>
                
 
                            +        <span class="p">]</span>
                
 
                            +        <span class="nb">super</span><span class="p">(</span><span class="n">PascalVOCUnifiedDetectionTrainDataset</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><span class="n">train_sets</span><span class="p">)</span>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -251,4 +301,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.mixup &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.mixup</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.mixup</h1><div class="highlight"><pre>
<span></span><span class="sd">&quot;&quot;&quot; Mixup and Cutmix</span>

<span class="sd">Papers:</span>
<span class="sd">mixup: Beyond Empirical Risk Minimization (https://arxiv.org/abs/1710.09412)</span>

<span class="sd">CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features (https://arxiv.org/abs/1905.04899)</span>

<span class="sd">Code Reference:</span>
<span class="sd">CutMix: https://github.com/clovaai/CutMix-PyTorch</span>
<span class="sd">CutMix by timm: https://github.com/rwightman/pytorch-image-models/timm</span>

<span class="sd">&quot;&quot;&quot;</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Union</span>

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">torch</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.dataset_exceptions</span> <span class="kn">import</span> <span class="n">IllegalDatasetParameterException</span>


<div class="viewcode-block" id="one_hot"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.mixup.one_hot">[docs]</a><span class="k">def</span> <span class="nf">one_hot</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">,</span> <span class="n">on_value</span><span class="o">=</span><span class="mf">1.</span><span class="p">,</span> <span class="n">off_value</span><span class="o">=</span><span class="mf">0.</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cuda&#39;</span><span class="p">):</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">long</span><span class="p">()</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">full</span><span class="p">((</span><span class="n">x</span><span class="o">.</span><span class="n">size</span><span class="p">()[</span><span class="mi">0</span><span class="p">],</span> <span class="n">num_classes</span><span class="p">),</span> <span class="n">off_value</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="o">.</span><span class="n">scatter_</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">on_value</span><span class="p">)</span></div>


<div class="viewcode-block" id="mixup_target"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.mixup.mixup_target">[docs]</a><span class="k">def</span> <span class="nf">mixup_target</span><span class="p">(</span><span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">lam</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.</span><span class="p">,</span> <span class="n">smoothing</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span> <span class="n">device</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;cuda&#39;</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    generate a smooth target (label) two-hot tensor to support the mixed images with different labels</span>
<span class="sd">    :param target: the targets tensor</span>
<span class="sd">    :param num_classes: number of classes (to set the final tensor size)</span>
<span class="sd">    :param lam: percentage of label a range [0, 1] in the mixing</span>
<span class="sd">    :param smoothing: the smoothing multiplier</span>
<span class="sd">    :param device: usable device [&#39;cuda&#39;, &#39;cpu&#39;]</span>
<span class="sd">    :return:</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">off_value</span> <span class="o">=</span> <span class="n">smoothing</span> <span class="o">/</span> <span class="n">num_classes</span>
    <span class="n">on_value</span> <span class="o">=</span> <span class="mf">1.</span> <span class="o">-</span> <span class="n">smoothing</span> <span class="o">+</span> <span class="n">off_value</span>
    <span class="n">y1</span> <span class="o">=</span> <span class="n">one_hot</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">,</span> <span class="n">on_value</span><span class="o">=</span><span class="n">on_value</span><span class="p">,</span> <span class="n">off_value</span><span class="o">=</span><span class="n">off_value</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">y2</span> <span class="o">=</span> <span class="n">one_hot</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">flip</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span> <span class="n">num_classes</span><span class="p">,</span> <span class="n">on_value</span><span class="o">=</span><span class="n">on_value</span><span class="p">,</span> <span class="n">off_value</span><span class="o">=</span><span class="n">off_value</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="k">return</span> <span class="n">y1</span> <span class="o">*</span> <span class="n">lam</span> <span class="o">+</span> <span class="n">y2</span> <span class="o">*</span> <span class="p">(</span><span class="mf">1.</span> <span class="o">-</span> <span class="n">lam</span><span class="p">)</span></div>


<div class="viewcode-block" id="rand_bbox"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.mixup.rand_bbox">[docs]</a><span class="k">def</span> <span class="nf">rand_bbox</span><span class="p">(</span><span class="n">img_shape</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span> <span class="n">lam</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">margin</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.</span><span class="p">,</span> <span class="n">count</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot; Standard CutMix bounding-box</span>
<span class="sd">    Generates a random square bbox based on lambda value. This impl includes</span>
<span class="sd">    support for enforcing a border margin as percent of bbox dimensions.</span>

<span class="sd">    :param img_shape: Image shape as tuple</span>
<span class="sd">    :param lam: Cutmix lambda value</span>
<span class="sd">    :param margin: Percentage of bbox dimension to enforce as margin (reduce amount of box outside image)</span>
<span class="sd">    :param count: Number of bbox to generate</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">ratio</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">lam</span><span class="p">)</span>
    <span class="n">img_h</span><span class="p">,</span> <span class="n">img_w</span> <span class="o">=</span> <span class="n">img_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>
    <span class="n">cut_h</span><span class="p">,</span> <span class="n">cut_w</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">img_h</span> <span class="o">*</span> <span class="n">ratio</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">img_w</span> <span class="o">*</span> <span class="n">ratio</span><span class="p">)</span>
    <span class="n">margin_y</span><span class="p">,</span> <span class="n">margin_x</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">margin</span> <span class="o">*</span> <span class="n">cut_h</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">margin</span> <span class="o">*</span> <span class="n">cut_w</span><span class="p">)</span>
    <span class="n">cy</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span> <span class="o">+</span> <span class="n">margin_y</span><span class="p">,</span> <span class="n">img_h</span> <span class="o">-</span> <span class="n">margin_y</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="n">cx</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span> <span class="o">+</span> <span class="n">margin_x</span><span class="p">,</span> <span class="n">img_w</span> <span class="o">-</span> <span class="n">margin_x</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="n">yl</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">cy</span> <span class="o">-</span> <span class="n">cut_h</span> <span class="o">//</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">img_h</span><span class="p">)</span>
    <span class="n">yh</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">cy</span> <span class="o">+</span> <span class="n">cut_h</span> <span class="o">//</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">img_h</span><span class="p">)</span>
    <span class="n">xl</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">cx</span> <span class="o">-</span> <span class="n">cut_w</span> <span class="o">//</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">img_w</span><span class="p">)</span>
    <span class="n">xh</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">cx</span> <span class="o">+</span> <span class="n">cut_w</span> <span class="o">//</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">img_w</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">yl</span><span class="p">,</span> <span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xh</span></div>


<div class="viewcode-block" id="rand_bbox_minmax"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.mixup.rand_bbox_minmax">[docs]</a><span class="k">def</span> <span class="nf">rand_bbox_minmax</span><span class="p">(</span><span class="n">img_shape</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span> <span class="n">minmax</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">],</span> <span class="n">count</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot; Min-Max CutMix bounding-box</span>
<span class="sd">    Inspired by Darknet cutmix impl, generates a random rectangular bbox</span>
<span class="sd">    based on min/max percent values applied to each dimension of the input image.</span>

<span class="sd">    Typical defaults for minmax are usually in the  .2-.3 for min and .8-.9 range for max.</span>

<span class="sd">    :param img_shape: Image shape as tuple</span>
<span class="sd">    :param minmax: Min and max bbox ratios (as percent of image size)</span>
<span class="sd">    :param count: Number of bbox to generate</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">minmax</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span>
    <span class="n">img_h</span><span class="p">,</span> <span class="n">img_w</span> <span class="o">=</span> <span class="n">img_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>
    <span class="n">cut_h</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">img_h</span> <span class="o">*</span> <span class="n">minmax</span><span class="p">[</span><span class="mi">0</span><span class="p">]),</span> <span class="nb">int</span><span class="p">(</span><span class="n">img_h</span> <span class="o">*</span> <span class="n">minmax</span><span class="p">[</span><span class="mi">1</span><span class="p">]),</span> <span class="n">size</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="n">cut_w</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">img_w</span> <span class="o">*</span> <span class="n">minmax</span><span class="p">[</span><span class="mi">0</span><span class="p">]),</span> <span class="nb">int</span><span class="p">(</span><span class="n">img_w</span> <span class="o">*</span> <span class="n">minmax</span><span class="p">[</span><span class="mi">1</span><span class="p">]),</span> <span class="n">size</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="n">yl</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">img_h</span> <span class="o">-</span> <span class="n">cut_h</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="n">xl</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">img_w</span> <span class="o">-</span> <span class="n">cut_w</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="n">yu</span> <span class="o">=</span> <span class="n">yl</span> <span class="o">+</span> <span class="n">cut_h</span>
    <span class="n">xu</span> <span class="o">=</span> <span class="n">xl</span> <span class="o">+</span> <span class="n">cut_w</span>
    <span class="k">return</span> <span class="n">yl</span><span class="p">,</span> <span class="n">yu</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xu</span></div>


<div class="viewcode-block" id="cutmix_bbox_and_lam"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.mixup.cutmix_bbox_and_lam">[docs]</a><span class="k">def</span> <span class="nf">cutmix_bbox_and_lam</span><span class="p">(</span><span class="n">img_shape</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span> <span class="n">lam</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">ratio_minmax</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">correct_lam</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                        <span class="n">count</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Generate bbox and apply lambda correction.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">ratio_minmax</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">yl</span><span class="p">,</span> <span class="n">yu</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xu</span> <span class="o">=</span> <span class="n">rand_bbox_minmax</span><span class="p">(</span><span class="n">img_shape</span><span class="p">,</span> <span class="n">ratio_minmax</span><span class="p">,</span> <span class="n">count</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">yl</span><span class="p">,</span> <span class="n">yu</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xu</span> <span class="o">=</span> <span class="n">rand_bbox</span><span class="p">(</span><span class="n">img_shape</span><span class="p">,</span> <span class="n">lam</span><span class="p">,</span> <span class="n">count</span><span class="o">=</span><span class="n">count</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">correct_lam</span> <span class="ow">or</span> <span class="n">ratio_minmax</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">bbox_area</span> <span class="o">=</span> <span class="p">(</span><span class="n">yu</span> <span class="o">-</span> <span class="n">yl</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="n">xu</span> <span class="o">-</span> <span class="n">xl</span><span class="p">)</span>
        <span class="n">lam</span> <span class="o">=</span> <span class="mf">1.</span> <span class="o">-</span> <span class="n">bbox_area</span> <span class="o">/</span> <span class="nb">float</span><span class="p">(</span><span class="n">img_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="n">img_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">yl</span><span class="p">,</span> <span class="n">yu</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xu</span><span class="p">),</span> <span class="n">lam</span></div>


<div class="viewcode-block" id="CollateMixup"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.mixup.CollateMixup">[docs]</a><span class="k">class</span> <span class="nc">CollateMixup</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Collate with Mixup/Cutmix that applies different params to each element or whole batch</span>
<span class="sd">    A Mixup impl that&#39;s performed while collating the batches.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mixup_alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.</span><span class="p">,</span> <span class="n">cutmix_alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.</span><span class="p">,</span> <span class="n">cutmix_minmax</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                 <span class="n">prob</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">switch_prob</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                 <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;batch&#39;</span><span class="p">,</span> <span class="n">correct_lam</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> <span class="n">label_smoothing</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1000</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Mixup/Cutmix that applies different params to each element or whole batch</span>

<span class="sd">        :param mixup_alpha: mixup alpha value, mixup is active if &gt; 0.</span>
<span class="sd">        :param cutmix_alpha: cutmix alpha value, cutmix is active if &gt; 0.</span>
<span class="sd">        :param cutmix_minmax: cutmix min/max image ratio, cutmix is active and uses this vs alpha if not None.</span>
<span class="sd">        :param prob: probability of applying mixup or cutmix per batch or element</span>
<span class="sd">        :param switch_prob: probability of switching to cutmix instead of mixup when both are active</span>
<span class="sd">        :param mode: how to apply mixup/cutmix params (per &#39;batch&#39;, &#39;pair&#39; (pair of elements), &#39;elem&#39; (element)</span>
<span class="sd">        :param correct_lam: apply lambda correction when cutmix bbox clipped by image borders</span>
<span class="sd">        :param label_smoothing: apply label smoothing to the mixed target tensor</span>
<span class="sd">        :param num_classes: number of classes for target</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span> <span class="o">=</span> <span class="n">mixup_alpha</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span> <span class="o">=</span> <span class="n">cutmix_alpha</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_minmax</span> <span class="o">=</span> <span class="n">cutmix_minmax</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_minmax</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_minmax</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span>
            <span class="c1"># force cutmix alpha == 1.0 when minmax active to keep logic simple &amp; safe</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span> <span class="o">=</span> <span class="mf">1.0</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mix_prob</span> <span class="o">=</span> <span class="n">prob</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">switch_prob</span> <span class="o">=</span> <span class="n">switch_prob</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">label_smoothing</span> <span class="o">=</span> <span class="n">label_smoothing</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">num_classes</span> <span class="o">=</span> <span class="n">num_classes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">=</span> <span class="n">mode</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">correct_lam</span> <span class="o">=</span> <span class="n">correct_lam</span>  <span class="c1"># correct lambda based on clipped area for cutmix</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mixup_enabled</span> <span class="o">=</span> <span class="kc">True</span>  <span class="c1"># set to false to disable mixing (intended tp be set by train loop)</span>

    <span class="k">def</span> <span class="nf">_params_per_elem</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        generate two random masks to define which elements of the batch will be mixed and how (depending on the</span>
<span class="sd">        self.mixup_enabled, self.mixup_alpha, self.cutmix_alpha parameters</span>

<span class="sd">        :param batch_size:</span>
<span class="sd">        :return: two tensors with shape=batch_size - the first contains the lambda value per batch element</span>
<span class="sd">        and the second is a binary flag indicating use of cutmix per batch element</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">lam</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">batch_size</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">float32</span><span class="p">)</span>
        <span class="n">use_cutmix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">batch_size</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">bool</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_enabled</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span><span class="p">:</span>
                <span class="n">use_cutmix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">batch_size</span><span class="p">)</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">switch_prob</span>
                <span class="n">lam_mix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">where</span><span class="p">(</span>
                    <span class="n">use_cutmix</span><span class="p">,</span>
                    <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">(</span><span class="n">sample_shape</span><span class="o">=</span><span class="n">batch_size</span><span class="p">),</span>
                    <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">(</span><span class="n">sample_shape</span><span class="o">=</span><span class="n">batch_size</span><span class="p">))</span>
            <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span><span class="p">:</span>
                <span class="n">lam_mix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">(</span><span class="n">sample_shape</span><span class="o">=</span><span class="n">batch_size</span><span class="p">)</span>
            <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span><span class="p">:</span>
                <span class="n">use_cutmix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">batch_size</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">bool</span><span class="p">)</span>
                <span class="n">lam_mix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">(</span><span class="n">sample_shape</span><span class="o">=</span><span class="n">batch_size</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">raise</span> <span class="n">IllegalDatasetParameterException</span><span class="p">(</span><span class="s2">&quot;One of mixup_alpha &gt; 0., cutmix_alpha &gt; 0., &quot;</span>
                                                       <span class="s2">&quot;cutmix_minmax not None should be true.&quot;</span><span class="p">)</span>
            <span class="n">lam</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">batch_size</span><span class="p">)</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">mix_prob</span><span class="p">,</span> <span class="n">lam_mix</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">),</span> <span class="n">lam</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">lam</span><span class="p">,</span> <span class="n">use_cutmix</span>

    <span class="k">def</span> <span class="nf">_params_per_batch</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        generate two random parameters to define if batch will be mixed and how (depending on the</span>
<span class="sd">        self.mixup_enabled, self.mixup_alpha, self.cutmix_alpha parameters</span>

<span class="sd">        :return: two parameters - the first contains the lambda value for the whole batch</span>
<span class="sd">        and the second is a binary flag indicating use of cutmix for the batch</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">lam</span> <span class="o">=</span> <span class="mf">1.</span>
        <span class="n">use_cutmix</span> <span class="o">=</span> <span class="kc">False</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_enabled</span> <span class="ow">and</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">mix_prob</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span><span class="p">:</span>
                <span class="n">use_cutmix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">switch_prob</span>
                <span class="n">lam_mix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">()</span> <span class="k">if</span> <span class="n">use_cutmix</span> <span class="k">else</span> \
                    <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">()</span>
            <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span><span class="p">:</span>
                <span class="n">lam_mix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">mixup_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">()</span>
            <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span> <span class="o">&gt;</span> <span class="mf">0.</span><span class="p">:</span>
                <span class="n">use_cutmix</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="n">lam_mix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">distributions</span><span class="o">.</span><span class="n">beta</span><span class="o">.</span><span class="n">Beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cutmix_alpha</span><span class="p">)</span><span class="o">.</span><span class="n">sample</span><span class="p">()</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">raise</span> <span class="n">IllegalDatasetParameterException</span><span class="p">(</span><span class="s2">&quot;One of mixup_alpha &gt; 0., cutmix_alpha &gt; 0., &quot;</span>
                                                       <span class="s2">&quot;cutmix_minmax not None should be true.&quot;</span><span class="p">)</span>
            <span class="n">lam</span> <span class="o">=</span> <span class="nb">float</span><span class="p">(</span><span class="n">lam_mix</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">lam</span><span class="p">,</span> <span class="n">use_cutmix</span>

    <span class="k">def</span> <span class="nf">_mix_elem_collate</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">output</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">half</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This is the implementation for &#39;elem&#39; or &#39;half&#39; modes</span>
<span class="sd">        :param output: the output tensor to fill</span>
<span class="sd">        :param batch: list of thr batch items</span>
<span class="sd">        :return: a tensor containing the lambda values used for the mixing (this vector can be used for</span>
<span class="sd">        mixing the labels as well)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">batch_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>
        <span class="n">num_elem</span> <span class="o">=</span> <span class="n">batch_size</span> <span class="o">//</span> <span class="mi">2</span> <span class="k">if</span> <span class="n">half</span> <span class="k">else</span> <span class="n">batch_size</span>
        <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">output</span><span class="p">)</span> <span class="o">==</span> <span class="n">num_elem</span>
        <span class="n">lam_batch</span><span class="p">,</span> <span class="n">use_cutmix</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_params_per_elem</span><span class="p">(</span><span class="n">num_elem</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_elem</span><span class="p">):</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">batch_size</span> <span class="o">-</span> <span class="n">i</span> <span class="o">-</span> <span class="mi">1</span>
            <span class="n">lam</span> <span class="o">=</span> <span class="n">lam_batch</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
            <span class="n">mixed</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
            <span class="k">if</span> <span class="n">lam</span> <span class="o">!=</span> <span class="mf">1.</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">use_cutmix</span><span class="p">[</span><span class="n">i</span><span class="p">]:</span>
                    <span class="k">if</span> <span class="ow">not</span> <span class="n">half</span><span class="p">:</span>
                        <span class="n">mixed</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">clone</span><span class="p">(</span><span class="n">mixed</span><span class="p">)</span>
                    <span class="p">(</span><span class="n">yl</span><span class="p">,</span> <span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xh</span><span class="p">),</span> <span class="n">lam</span> <span class="o">=</span> <span class="n">cutmix_bbox_and_lam</span><span class="p">(</span>
                        <span class="n">output</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">lam</span><span class="p">,</span> <span class="n">ratio_minmax</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_minmax</span><span class="p">,</span> <span class="n">correct_lam</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">correct_lam</span><span class="p">)</span>
                    <span class="n">mixed</span><span class="p">[:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">]</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="n">j</span><span class="p">][</span><span class="mi">0</span><span class="p">][:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">]</span>
                    <span class="n">lam_batch</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">lam</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">mixed</span> <span class="o">=</span> <span class="n">mixed</span> <span class="o">*</span> <span class="n">lam</span> <span class="o">+</span> <span class="n">batch</span><span class="p">[</span><span class="n">j</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">lam</span><span class="p">)</span>
            <span class="n">output</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">mixed</span>
        <span class="k">if</span> <span class="n">half</span><span class="p">:</span>
            <span class="n">lam_batch</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">lam_batch</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">num_elem</span><span class="p">)))</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">lam_batch</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_mix_pair_collate</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">output</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="nb">list</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This is the implementation for &#39;pair&#39; mode</span>
<span class="sd">        :param output: the output tensor to fill</span>
<span class="sd">        :param batch: list of thr batch items</span>
<span class="sd">        :return: a tensor containing the lambda values used for the mixing (this vector can be used for</span>
<span class="sd">        mixing the labels as well)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">batch_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>
        <span class="n">lam_batch</span><span class="p">,</span> <span class="n">use_cutmix</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_params_per_elem</span><span class="p">(</span><span class="n">batch_size</span> <span class="o">//</span> <span class="mi">2</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">batch_size</span> <span class="o">//</span> <span class="mi">2</span><span class="p">):</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">batch_size</span> <span class="o">-</span> <span class="n">i</span> <span class="o">-</span> <span class="mi">1</span>
            <span class="n">lam</span> <span class="o">=</span> <span class="n">lam_batch</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
            <span class="n">mixed_i</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">mixed_j</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="n">j</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
            <span class="k">assert</span> <span class="mi">0</span> <span class="o">&lt;=</span> <span class="n">lam</span> <span class="o">&lt;=</span> <span class="mf">1.0</span>
            <span class="k">if</span> <span class="n">lam</span> <span class="o">&lt;</span> <span class="mf">1.</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">use_cutmix</span><span class="p">[</span><span class="n">i</span><span class="p">]:</span>
                    <span class="p">(</span><span class="n">yl</span><span class="p">,</span> <span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xh</span><span class="p">),</span> <span class="n">lam</span> <span class="o">=</span> <span class="n">cutmix_bbox_and_lam</span><span class="p">(</span>
                        <span class="n">output</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">lam</span><span class="p">,</span> <span class="n">ratio_minmax</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_minmax</span><span class="p">,</span> <span class="n">correct_lam</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">correct_lam</span><span class="p">)</span>
                    <span class="n">patch_i</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">clone</span><span class="p">(</span><span class="n">mixed_i</span><span class="p">[:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">])</span>
                    <span class="n">mixed_i</span><span class="p">[:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">]</span> <span class="o">=</span> <span class="n">mixed_j</span><span class="p">[:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">]</span>
                    <span class="n">mixed_j</span><span class="p">[:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">]</span> <span class="o">=</span> <span class="n">patch_i</span>
                    <span class="n">lam_batch</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">lam</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">mixed_temp</span> <span class="o">=</span> <span class="n">mixed_i</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span> <span class="o">*</span> <span class="n">lam</span> <span class="o">+</span> <span class="n">mixed_j</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">lam</span><span class="p">)</span>
                    <span class="n">mixed_j</span> <span class="o">=</span> <span class="n">mixed_j</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span> <span class="o">*</span> <span class="n">lam</span> <span class="o">+</span> <span class="n">mixed_i</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">lam</span><span class="p">)</span>
                    <span class="n">mixed_i</span> <span class="o">=</span> <span class="n">mixed_temp</span>
                    <span class="n">torch</span><span class="o">.</span><span class="n">rint</span><span class="p">(</span><span class="n">mixed_j</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mixed_j</span><span class="p">)</span>
                    <span class="n">torch</span><span class="o">.</span><span class="n">rint</span><span class="p">(</span><span class="n">mixed_i</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mixed_i</span><span class="p">)</span>
            <span class="n">output</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">mixed_i</span>
            <span class="n">output</span><span class="p">[</span><span class="n">j</span><span class="p">]</span> <span class="o">+=</span> <span class="n">mixed_j</span>
        <span class="n">lam_batch</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">lam_batch</span><span class="p">,</span> <span class="n">lam_batch</span><span class="p">[::</span><span class="o">-</span><span class="mi">1</span><span class="p">]))</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">lam_batch</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_mix_batch_collate</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">output</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="nb">list</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This is the implementation for &#39;batch&#39; mode</span>
<span class="sd">        :param output: the output tensor to fill</span>
<span class="sd">        :param batch: list of thr batch items</span>
<span class="sd">        :return: the lambda value used for the mixing</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">batch_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>
        <span class="n">lam</span><span class="p">,</span> <span class="n">use_cutmix</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_params_per_batch</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">use_cutmix</span><span class="p">:</span>
            <span class="p">(</span><span class="n">yl</span><span class="p">,</span> <span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">,</span> <span class="n">xh</span><span class="p">),</span> <span class="n">lam</span> <span class="o">=</span> <span class="n">cutmix_bbox_and_lam</span><span class="p">(</span>
                <span class="n">output</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">lam</span><span class="p">,</span> <span class="n">ratio_minmax</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">cutmix_minmax</span><span class="p">,</span> <span class="n">correct_lam</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">correct_lam</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">batch_size</span><span class="p">):</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">batch_size</span> <span class="o">-</span> <span class="n">i</span> <span class="o">-</span> <span class="mi">1</span>
            <span class="n">mixed</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
            <span class="k">if</span> <span class="n">lam</span> <span class="o">!=</span> <span class="mf">1.</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">use_cutmix</span><span class="p">:</span>
                    <span class="n">mixed</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">clone</span><span class="p">(</span><span class="n">mixed</span><span class="p">)</span>  <span class="c1"># don&#39;t want to modify the original while iterating</span>
                    <span class="n">mixed</span><span class="p">[:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">]</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="n">j</span><span class="p">][</span><span class="mi">0</span><span class="p">][:,</span> <span class="n">yl</span><span class="p">:</span><span class="n">yh</span><span class="p">,</span> <span class="n">xl</span><span class="p">:</span><span class="n">xh</span><span class="p">]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">mixed</span> <span class="o">=</span> <span class="n">mixed</span> <span class="o">*</span> <span class="n">lam</span> <span class="o">+</span> <span class="n">batch</span><span class="p">[</span><span class="n">j</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">lam</span><span class="p">)</span>
            <span class="n">output</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">mixed</span>
        <span class="k">return</span> <span class="n">lam</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">,</span> <span class="n">_</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">batch_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">batch_size</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="n">IllegalDatasetParameterException</span><span class="p">(</span><span class="s1">&#39;Batch size should be even when using this&#39;</span><span class="p">)</span>
        <span class="n">half</span> <span class="o">=</span> <span class="s1">&#39;half&#39;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span>
        <span class="k">if</span> <span class="n">half</span><span class="p">:</span>
            <span class="n">batch_size</span> <span class="o">//=</span> <span class="mi">2</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">batch_size</span><span class="p">,</span> <span class="o">*</span><span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">shape</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">float32</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">==</span> <span class="s1">&#39;elem&#39;</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">==</span> <span class="s1">&#39;half&#39;</span><span class="p">:</span>
            <span class="n">lam</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_mix_elem_collate</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">batch</span><span class="p">,</span> <span class="n">half</span><span class="o">=</span><span class="n">half</span><span class="p">)</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">==</span> <span class="s1">&#39;pair&#39;</span><span class="p">:</span>
            <span class="n">lam</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_mix_pair_collate</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">batch</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">lam</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_mix_batch_collate</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">batch</span><span class="p">)</span>
        <span class="n">target</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">b</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="k">for</span> <span class="n">b</span> <span class="ow">in</span> <span class="n">batch</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">int32</span><span class="p">)</span>
        <span class="n">target</span> <span class="o">=</span> <span class="n">mixup_target</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">lam</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">label_smoothing</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s1">&#39;cpu&#39;</span><span class="p">)</span>
        <span class="n">target</span> <span class="o">=</span> <span class="n">target</span><span class="p">[:</span><span class="n">batch_size</span><span class="p">]</span>

        <span class="k">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">target</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
        <script src="../../../../../_static/jquery.js"></script>
        <script src="../../../../../_static/underscore.js"></script>
        <script src="../../../../../_static/doctools.js"></script>
    <script src="../../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">cv2</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span><span class="p">,</span> <span class="n">ImageColor</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets.segmentation_dataset</span> <span class="kn">import</span> <span class="n">SegmentationDataSet</span>

<span class="c1"># TODO - ADD COARSE DATA - right now cityscapes dataset includes fine annotations. It&#39;s optional to use extra coarse</span>
<span class="c1">#  annotations.</span>

<span class="c1"># label for background and labels to ignore during training and evaluation.</span>
<span class="n">CITYSCAPES_IGNORE_LABEL</span> <span class="o">=</span> <span class="mi">19</span>


<div class="viewcode-block" id="CityscapesDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation.CityscapesDataset">[docs]</a><span class="k">class</span> <span class="nc">CityscapesDataset</span><span class="p">(</span><span class="n">SegmentationDataSet</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    CityscapesDataset - Segmentation Data Set Class for Cityscapes Segmentation Data Set,</span>
<span class="sd">    main resolution of dataset: (2048 x 1024).</span>
<span class="sd">    Not all the original labels are used for training and evaluation, according to cityscape paper:</span>
<span class="sd">    &quot;Classes that are too rare are excluded from our benchmark, leaving 19 classes for evaluation&quot;.</span>
<span class="sd">    For more details about the dataset labels format see:</span>
<span class="sd">    https://github.com/mcordts/cityscapesScripts/blob/master/cityscapesscripts/helpers/labels.py</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                 <span class="n">root_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                 <span class="n">list_file</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                 <span class="n">labels_csv_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                 <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param root: root directory to dataset.</span>
<span class="sd">        :param list_file: list file that contains names of images to load, line format: &lt;image_path&gt; &lt;label_path&gt;</span>
<span class="sd">        :param labels_csv_path: path to csv file, with labels metadata and mapping.</span>
<span class="sd">        :param kwargs: Any hyper params required for the dataset, i.e img_size, crop_size, cache_images</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span> <span class="o">=</span> <span class="n">root_dir</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">root_dir</span><span class="p">,</span> <span class="n">list_file</span><span class="o">=</span><span class="n">list_file</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="c1"># labels dataframe for labels metadata.</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">labels_data</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">recfromcsv</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span> <span class="n">labels_csv_path</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="s1">&#39;&lt;i8,U20,&lt;i8,&lt;i8,U12,&lt;i8,?,?,U7&#39;</span><span class="p">,</span> <span class="n">comments</span><span class="o">=</span><span class="s1">&#39;&amp;&#39;</span><span class="p">)</span>
        <span class="c1"># map vector to map ground-truth labels to train labels</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">labels_map</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels_data</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="s2">&quot;trainid&quot;</span><span class="p">)</span>
        <span class="c1"># class names</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels_data</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="s2">&quot;name&quot;</span><span class="p">)[</span><span class="n">np</span><span class="o">.</span><span class="n">logical_not</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">labels_data</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="s2">&quot;ignoreineval&quot;</span><span class="p">))]</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span>
        <span class="c1"># color palette for visualization</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">train_id_color_palette</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_create_color_palette</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">_generate_samples_and_targets</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        override _generate_samples_and_targets function, to parse list file.</span>
<span class="sd">        line format of list file: &lt;image_path&gt; &lt;label_path&gt;</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">list_file_path</span><span class="p">))</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
            <span class="n">img_list</span> <span class="o">=</span> <span class="p">[</span><span class="n">line</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span><span class="o">.</span><span class="n">split</span><span class="p">()</span> <span class="k">for</span> <span class="n">line</span> <span class="ow">in</span> <span class="n">f</span><span class="p">]</span>
        <span class="k">for</span> <span class="n">image_path</span><span class="p">,</span> <span class="n">label_path</span> <span class="ow">in</span> <span class="n">img_list</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">samples_targets_tuples_list</span><span class="o">.</span><span class="n">append</span><span class="p">((</span>
                <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root</span><span class="p">,</span> <span class="n">image_path</span><span class="p">),</span>
                <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root</span><span class="p">,</span> <span class="n">label_path</span><span class="p">)</span>
            <span class="p">))</span>

<div class="viewcode-block" id="CityscapesDataset.target_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation.CityscapesDataset.target_loader">[docs]</a>    <span class="k">def</span> <span class="nf">target_loader</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">label_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Image</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Override target_loader function, load the labels mask image.</span>
<span class="sd">            :param label_path:  Path to the label image.</span>
<span class="sd">            :return:                     The mask image created from the array, with converted class labels.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># assert that is a png file, other file types might alter the class labels value.</span>
        <span class="k">assert</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">splitext</span><span class="p">(</span><span class="n">label_path</span><span class="p">)[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span> <span class="o">==</span> <span class="s2">&quot;.png&quot;</span>

        <span class="n">label</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">imread</span><span class="p">(</span><span class="n">label_path</span><span class="p">,</span> <span class="n">cv2</span><span class="o">.</span><span class="n">IMREAD_GRAYSCALE</span><span class="p">)</span>
        <span class="c1"># map ground-truth ids to train ids</span>
        <span class="n">label</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels_map</span><span class="p">[</span><span class="n">label</span><span class="p">]</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">Image</span><span class="o">.</span><span class="n">fromarray</span><span class="p">(</span><span class="n">label</span><span class="p">,</span> <span class="s1">&#39;L&#39;</span><span class="p">)</span></div>

    <span class="k">def</span> <span class="nf">_create_color_palette</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Create color pallete for visualizing the segmentation masks</span>
<span class="sd">        :return: list of rgb color values</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">palette</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">hex_colors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels_data</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="s2">&quot;color&quot;</span><span class="p">)[</span><span class="n">np</span><span class="o">.</span><span class="n">logical_not</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">labels_data</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="s2">&quot;ignoreineval&quot;</span><span class="p">))]</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span>

        <span class="k">for</span> <span class="n">hex_color</span> <span class="ow">in</span> <span class="n">hex_colors</span><span class="p">:</span>
            <span class="n">rgb_color</span> <span class="o">=</span> <span class="n">ImageColor</span><span class="o">.</span><span class="n">getcolor</span><span class="p">(</span><span class="n">hex_color</span><span class="p">,</span> <span class="s2">&quot;RGB&quot;</span><span class="p">)</span>
            <span class="n">palette</span> <span class="o">+=</span> <span class="p">[</span><span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">rgb_color</span><span class="p">]</span>

        <span class="k">return</span> <span class="n">palette</span>

<div class="viewcode-block" id="CityscapesDataset.get_train_ids_color_palette"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation.CityscapesDataset.get_train_ids_color_palette">[docs]</a>    <span class="k">def</span> <span class="nf">get_train_ids_color_palette</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_id_color_palette</span></div>

<div class="viewcode-block" id="CityscapesDataset.target_transform"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation.CityscapesDataset.target_transform">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">target_transform</span><span class="p">(</span><span class="n">target</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        target_transform - Transforms the sample image</span>
<span class="sd">        This function overrides the original function from SegmentationDataSet and changes target pixels with value</span>
<span class="sd">        255 to value = CITYSCAPES_IGNORE_LABEL. This was done since current IoU metric from torchmetrics does not</span>
<span class="sd">        support such a high ignore label value (crashed on OOM)</span>

<span class="sd">            :param target: The target mask to transform</span>
<span class="sd">            :return:       The transformed target mask</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">out</span> <span class="o">=</span> <span class="n">SegmentationDataSet</span><span class="o">.</span><span class="n">target_transform</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
        <span class="n">out</span><span class="p">[</span><span class="n">out</span> <span class="o">==</span> <span class="mi">255</span><span class="p">]</span> <span class="o">=</span> <span class="n">CITYSCAPES_IGNORE_LABEL</span>
        <span class="k">return</span> <span class="n">out</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.segmentation_datasets.coco_segmentation &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.segmentation_datasets.coco_segmentation &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -87,14 +89,11 @@
 
                               <h1>Source code for super_gradients.training.datasets.segmentation_datasets.coco_segmentation</h1><div class="highlight"><pre>
                
 
                             <span></span><span class="kn">import</span> <span class="nn">os</span>
                
 
                            -<span class="kn">import</span> <span class="nn">torch</span>
                
 
                            +
                
 
                             <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
                
 
                            -<span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                            +<span class="kn">import</span> <span class="nn">torch</span>
                
 
                             <span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span>
                
 
                            -<span class="kn">import</span> <span class="nn">torchvision.transforms</span> <span class="k">as</span> <span class="nn">transform</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.transforms.transforms</span> <span class="kn">import</span> <span class="n">RandomFlip</span><span class="p">,</span> <span class="n">Rescale</span><span class="p">,</span> <span class="n">RandomRescale</span><span class="p">,</span> <span class="n">CropImageAndMask</span><span class="p">,</span> \
                
 
                            -    <span class="n">PadShortToCropSize</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">get_param</span>
                
 
                            +<span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                             <span class="k">try</span><span class="p">:</span>
                
 
                                 <span class="kn">from</span> <span class="nn">pycocotools.coco</span> <span class="kn">import</span> <span class="n">COCO</span>
                
@@ -106,36 +105,22 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets.segmentation_dataset</span> <span class="kn">import</span> <span class="n">SegmentationDataSet</span>
                
 
                            -<div class="viewcode-block" id="EmptyCoCoClassesSelectionException"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.EmptyCoCoClassesSelectionException">[docs]</a><span class="k">class</span> <span class="nc">EmptyCoCoClassesSelectionException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
                
 
                            -    <span class="k">pass</span></div>
                
 
                            +<span class="k">class</span> <span class="nc">EmptyCoCoClassesSelectionException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
                
 
                            +    <span class="k">pass</span>
                
 
                            -<div class="viewcode-block" id="CoCoSegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.CoCoSegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">CoCoSegmentationDataSet</span><span class="p">(</span><span class="n">SegmentationDataSet</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="CoCoSegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.CoCoSegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">CoCoSegmentationDataSet</span><span class="p">(</span><span class="n">SegmentationDataSet</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    CoCoSegmentationDataSet - Segmentation Data Set Class for COCO 2017 Segmentation Data Set</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_classes_inclusion_tuples_list</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">root_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                
 
                            +                 <span class="n">dataset_classes_inclusion_tuples_list</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                                     <span class="c1"># THERE ARE 91 CLASSES, INCLUDING BACKGROUND - BUT WE ENABLE THE USAGE OF SUBCLASSES, TO PARTIALLY USE THE DATA</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">dataset_classes_inclusion_tuples_list</span> <span class="o">=</span> <span class="n">dataset_classes_inclusion_tuples_list</span> <span class="ow">or</span> <span class="n">COCO_DEFAULT_CLASSES_TUPLES_LIST</span>
                
 
                            -        <span class="c1"># OVERRIDE DEFAULT AUGMENTATIONS, IMG_SIZE, CROP SIZE</span>
                
 
                            -        <span class="n">dataset_hyper_params</span> <span class="o">=</span> <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;dataset_hyper_params&#39;</span><span class="p">]</span>
                
 
                            -        <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;img_size&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">dataset_hyper_params</span><span class="p">,</span> <span class="s1">&#39;img_size&#39;</span><span class="p">,</span> <span class="mi">608</span><span class="p">)</span>
                
 
                            -        <span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;crop_size&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">dataset_hyper_params</span><span class="p">,</span> <span class="s1">&#39;crop_size&#39;</span><span class="p">,</span> <span class="mi">512</span><span class="p">)</span>
                
 
                            -        <span class="c1"># FIXME - Rescale before RandomRescale is kept for legacy support, consider removing it like most implementation</span>
                
 
                            -        <span class="c1">#  papers regimes.</span>
                
 
                            -        <span class="n">default_image_mask_transforms_aug</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span><span class="n">RandomFlip</span><span class="p">(),</span>
                
 
                            -                                                               <span class="n">Rescale</span><span class="p">(</span><span class="n">long_size</span><span class="o">=</span><span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;img_size&quot;</span><span class="p">]),</span>
                
 
                            -                                                               <span class="n">RandomRescale</span><span class="p">(</span><span class="n">scales</span><span class="o">=</span><span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">)),</span>
                
 
                            -                                                               <span class="n">PadShortToCropSize</span><span class="p">(</span><span class="n">crop_size</span><span class="o">=</span><span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;crop_size&#39;</span><span class="p">]),</span>
                
 
                            -                                                               <span class="n">CropImageAndMask</span><span class="p">(</span><span class="n">crop_size</span><span class="o">=</span><span class="n">kwargs</span><span class="p">[</span><span class="s1">&#39;crop_size&#39;</span><span class="p">],</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;random&quot;</span><span class="p">)])</span>
                
 
                            -        <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;image_mask_transforms_aug&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">dataset_hyper_params</span><span class="p">,</span> <span class="s2">&quot;image_mask_transforms_aug&quot;</span><span class="p">,</span>
                
 
                            -                                                        <span class="n">default_image_mask_transforms_aug</span><span class="p">)</span>
                
 
                            -        <span class="n">image_mask_transforms</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">dataset_hyper_params</span><span class="p">,</span> <span class="s1">&#39;image_mask_transforms&#39;</span><span class="p">)</span>
                
 
                            -        <span class="k">if</span> <span class="n">image_mask_transforms</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;image_mask_transforms&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image_mask_transforms</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">root_dir</span> <span class="o">=</span> <span class="n">root_dir</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">root_dir</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                                     <span class="n">_</span><span class="p">,</span> <span class="n">class_names</span> <span class="o">=</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_classes_inclusion_tuples_list</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">class_names</span>
                
@@ -162,7 +147,9 @@
 
                                         <span class="n">mask_metadata_tuple</span> <span class="o">=</span> <span class="p">(</span><span class="n">relevant_image_id</span><span class="p">,</span> <span class="n">img_metadata</span><span class="p">[</span><span class="s1">&#39;height&#39;</span><span class="p">],</span> <span class="n">img_metadata</span><span class="p">[</span><span class="s1">&#39;width&#39;</span><span class="p">])</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">samples_targets_tuples_list</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">image_path</span><span class="p">,</span> <span class="n">mask_metadata_tuple</span><span class="p">))</span>
                
 
                            -<div class="viewcode-block" id="CoCoSegmentationDataSet.target_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.CoCoSegmentationDataSet.target_loader">[docs]</a>    <span class="k">def</span> <span class="nf">target_loader</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mask_metadata_tuple</span><span class="p">)</span> <span class="o">-&gt;</spa
                
 
                            +        <span class="nb">super</span><span class="p">(</span><span class="n">CoCoSegmentationDataSet</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">_generate_samples_and_targets</span><span class="p">()</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="CoCoSegmentationDataSet.target_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.CoCoSegmentationDataSet.target_loader">[docs]</a>    <span class="k">def</span> <span class="nf">target_loader</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mask_metadata_tuple</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Image</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        target_loader</span>
                
 
                             <span class="sd">            :param mask_metadata_tuple:  A tuple of (coco_image_id, original_image_height, original_image_width)</span>
                
@@ -267,4 +254,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.datasets.segmentation_datasets.pascal_aug_segmentation &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
        <script src="../../../../../_static/jquery.js"></script>
        <script src="../../../../../_static/underscore.js"></script>
        <script src="../../../../../_static/doctools.js"></script>
    <script src="../../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.datasets.segmentation_datasets.pascal_aug_segmentation</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.datasets.segmentation_datasets.pascal_aug_segmentation</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">scipy.io</span>
<span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets.pascal_voc_segmentation</span> <span class="kn">import</span> <span class="n">PascalVOC2012SegmentationDataSet</span>

<span class="n">PASCAL_AUG_CLASSES</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s1">&#39;background&#39;</span><span class="p">,</span> <span class="s1">&#39;airplane&#39;</span><span class="p">,</span> <span class="s1">&#39;bicycle&#39;</span><span class="p">,</span> <span class="s1">&#39;bird&#39;</span><span class="p">,</span> <span class="s1">&#39;boat&#39;</span><span class="p">,</span> <span class="s1">&#39;bottle&#39;</span><span class="p">,</span>
    <span class="s1">&#39;bus&#39;</span><span class="p">,</span> <span class="s1">&#39;car&#39;</span><span class="p">,</span> <span class="s1">&#39;cat&#39;</span><span class="p">,</span> <span class="s1">&#39;chair&#39;</span><span class="p">,</span> <span class="s1">&#39;cow&#39;</span><span class="p">,</span> <span class="s1">&#39;diningtable&#39;</span><span class="p">,</span> <span class="s1">&#39;dog&#39;</span><span class="p">,</span> <span class="s1">&#39;horse&#39;</span><span class="p">,</span>
    <span class="s1">&#39;motorcycle&#39;</span><span class="p">,</span> <span class="s1">&#39;person&#39;</span><span class="p">,</span> <span class="s1">&#39;potted-plant&#39;</span><span class="p">,</span> <span class="s1">&#39;sheep&#39;</span><span class="p">,</span> <span class="s1">&#39;sofa&#39;</span><span class="p">,</span> <span class="s1">&#39;train&#39;</span><span class="p">,</span>
    <span class="s1">&#39;tv&#39;</span>
<span class="p">]</span>


<div class="viewcode-block" id="PascalAUG2012SegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.PascalAUG2012SegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">PascalAUG2012SegmentationDataSet</span><span class="p">(</span><span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    PascalAUG2012SegmentationDataSet - Segmentation Data Set Class for Pascal AUG 2012 Data Set</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sample_suffix</span> <span class="o">=</span> <span class="s1">&#39;.jpg&#39;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">target_suffix</span> <span class="o">=</span> <span class="s1">&#39;.mat&#39;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">sample_suffix</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">sample_suffix</span><span class="p">,</span> <span class="n">target_suffix</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">target_suffix</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>

        <span class="c1"># THERE ARE 21 CLASSES, INCLUDING BACKGROUND</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">PASCAL_AUG_CLASSES</span>

<div class="viewcode-block" id="PascalAUG2012SegmentationDataSet.target_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.PascalAUG2012SegmentationDataSet.target_loader">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">target_loader</span><span class="p">(</span><span class="n">target_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Image</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        target_loader</span>
<span class="sd">            :param target_path: The path to the target data</span>
<span class="sd">            :return:            The loaded target</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">mat</span> <span class="o">=</span> <span class="n">scipy</span><span class="o">.</span><span class="n">io</span><span class="o">.</span><span class="n">loadmat</span><span class="p">(</span><span class="n">target_path</span><span class="p">,</span> <span class="n">mat_dtype</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">squeeze_me</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                               <span class="n">struct_as_record</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">mat</span><span class="p">[</span><span class="s1">&#39;GTcls&#39;</span><span class="p">]</span><span class="o">.</span><span class="n">Segmentation</span>
        <span class="k">return</span> <span class="n">Image</span><span class="o">.</span><span class="n">fromarray</span><span class="p">(</span><span class="n">mask</span><span class="p">)</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.segmentation_datasets.pascal_voc_segmentation &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.segmentation_datasets.pascal_voc_segmentation &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -89,31 +91,71 @@
 
                             <span></span><span class="kn">import</span> <span class="nn">os</span>
                
 
                             <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
                
 
                            +<span class="kn">import</span> <span class="nn">scipy.io</span>
                
 
                            +<span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span>
                
 
                            +<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">ConcatDataset</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets.segmentation_dataset</span> <span class="kn">import</span> <span class="n">SegmentationDataSet</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                            +
                
 
                            +<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                             <span class="n">PASCAL_VOC_2012_CLASSES</span> <span class="o">=</span> <span class="p">[</span>
                
 
                            -    <span class="s1">&#39;background&#39;</span><span class="p">,</span> <span class="s1">&#39;aeroplane&#39;</span><span class="p">,</span> <span class="s1">&#39;bicycle&#39;</span><span class="p">,</span> <span class="s1">&#39;bird&#39;</span><span class="p">,</span> <span class="s1">&#39;boat&#39;</span><span class="p">,</span> <span class="s1">&#39;bottle&#39;</span><span class="p">,</span>
                
 
                            -    <span class="s1">&#39;bus&#39;</span><span class="p">,</span> <span class="s1">&#39;car&#39;</span><span class="p">,</span> <span class="s1">&#39;cat&#39;</span><span class="p">,</span> <span class="s1">&#39;chair&#39;</span><span class="p">,</span> <span class="s1">&#39;cow&#39;</span><span class="p">,</span> <span class="s1">&#39;diningtable&#39;</span><span class="p">,</span> <span class="s1">&#39;dog&#39;</span><span class="p">,</span> <span class="s1">&#39;horse&#39;</span><span class=
                
 
                            -    <span class="s1">&#39;motorbike&#39;</span><span class="p">,</span> <span class="s1">&#39;person&#39;</span><span class="p">,</span> <span class="s1">&#39;potted-plant&#39;</span><span class="p">,</span> <span class="s1">&#39;sheep&#39;</span><span class="p">,</span> <span class="s1">&#39;sofa&#39;</span><span class="p">,</span> <span class="s1">&#39;train&#39;</span><span class="p">,</span>
                
 
                            -    <span class="s1">&#39;tv/monitor&#39;</span><span class="p">,</span> <span class="s1">&#39;ambigious&#39;</span>
                
 
                            +    <span class="s2">&quot;background&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;aeroplane&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;bicycle&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;bird&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;boat&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;bottle&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;bus&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;car&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;cat&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;chair&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;cow&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;diningtable&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;dog&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;horse&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;motorbike&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;person&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;potted-plant&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;sheep&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;sofa&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;train&quot;</span><span class="p">,</span>
                
 
                            +    <span class="s2">&quot;tv/monitor&quot;</span><span class="p">,</span>
                
 
                             <span class="p">]</span>
                
 
                            -<div class="viewcode-block" id="PascalVOC2012SegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.PascalVOC2012SegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">PascalVOC2012SegmentationDataSet</span><span class="p">(</span><span class="n">SegmentationDataSet</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="PascalVOC2012SegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalVOC2012SegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">PascalVOC2012SegmentationDataSet</span><span class="p">(</span><span class="n">SegmentationDataSet</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    PascalVOC2012SegmentationDataSet - Segmentation Data Set Class for Pascal VOC 2012 Data Set</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +    <span class="n">IGNORE_LABEL</span> <span class="o">=</span> <span class="mi">21</span>
                
 
                            +    <span class="n">_ORIGINAL_IGNORE_LABEL</span> <span class="o">=</span> <span class="mi">255</span>
                
 
                            +
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample_suffix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">target_suffix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><sp
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">sample_suffix</span> <span class="o">=</span> <span class="s1">&#39;.jpg&#39;</span> <span class="k">if</span> <span class="n">sample_suffix</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">sample_suffix</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">target_suffix</span> <span class="o">=</span> <span class="s1">&#39;.png&#39;</span> <span class="k">if</span> <span class="n">target_suffix</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">target_suffix</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">sample_suffix</span> <span class="o">=</span> <span class="s2">&quot;.jpg&quot;</span> <span class="k">if</span> <span class="n">sample_suffix</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">sample_suffix</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">target_suffix</span> <span class="o">=</span> <span class="s2">&quot;.png&quot;</span> <span class="k">if</span> <span class="n">target_suffix</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">target_suffix</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            -        <span class="c1"># THERE ARE 21 CLASSES, AND BACKGROUND</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">PASCAL_VOC_2012_CLASSES</span>
                
 
                            -<div class="viewcode-block" id="PascalVOC2012SegmentationDataSet.decode_segmentation_mask"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.PascalVOC2012SegmentationDataSet.decode_segmentation_mask">[docs]</a>    <span class="k">def</span> <span class="nf">decode_segmentation_mask</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">label_mask</span><span cla
                
 
                            +<div class="viewcode-block" id="PascalVOC2012SegmentationDataSet.target_transform"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalVOC2012SegmentationDataSet.target_transform">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">target_transform</span><span class="p">(</span><span class="n">target</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        target_transform - Transforms the label mask</span>
                
 
                            +<span class="sd">        This function overrides the original function from SegmentationDataSet and changes target pixels with value</span>
                
 
                            +<span class="sd">        255 to value = IGNORE_LABEL. This was done since current IoU metric from torchmetrics does not</span>
                
 
                            +<span class="sd">        support such a high ignore label value (crashed on OOM)</span>
                
 
                            +
                
 
                            +<span class="sd">            :param target: The target mask to transform</span>
                
 
                            +<span class="sd">            :return:       The transformed target mask</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">out</span> <span class="o">=</span> <span class="n">SegmentationDataSet</span><span class="o">.</span><span class="n">target_transform</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
                
 
                            +        <span class="n">out</span><span class="p">[</span><span class="n">out</span> <span class="o">==</span> <span class="n">PascalVOC2012SegmentationDataSet</span><span class="o">.</span><span class="n">_ORIGINAL_IGNORE_LABEL</span><span class="p">]</span> <span class="o">=</span> <span class="n">PascalVOC2012SegmentationDataSet</span><span class="o">.</span><span class="n">IGNORE_LABEL</span>
                
 
                            +        <span class="k">return</span> <span class="n">out</span></div>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="PascalVOC2012SegmentationDataSet.decode_segmentation_mask"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalVOC2012SegmentationDataSet.decode_segmentation_mask">[docs]</a>    <span class="k">def</span> <span class="nf">decode_segmentation_mask</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">label_mask</span><span class="p">:</span> <span class="n"
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        decode_segmentation_mask - Decodes the colors for the Segmentation Mask</span>
                
 
                             <span class="sd">            :param: label_mask:  an (M,N) array of integer values denoting</span>
                
@@ -125,8 +167,7 @@
 
                                     <span class="n">g</span> <span class="o">=</span> <span class="n">label_mask</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                
 
                                     <span class="n">b</span> <span class="o">=</span> <span class="n">label_mask</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                
 
                            -        <span class="c1"># REMOVING THE BACKGROUND CLASS FROM THE PLOTS</span>
                
 
                            -        <span class="n">num_classes_to_plot</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">classes</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span>
                
 
                            +        <span class="n">num_classes_to_plot</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">classes</span><span class="p">)</span>
                
 
                                     <span class="k">for</span> <span class="n">ll</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">num_classes_to_plot</span><span class="p">):</span>
                
 
                                         <span class="n">r</span><span class="p">[</span><span class="n">label_mask</span> <span class="o">==</span> <span class="n">ll</span><span class="p">]</span> <span class="o">=</span> <span class="n">label_colours</span><span class="p">[</span><span class="n">ll</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
                
 
                                         <span class="n">g</span><span class="p">[</span><span class="n">label_mask</span> <span class="o">==</span> <span class="n">ll</span><span class="p">]</span> <span class="o">=</span> <span class="n">label_colours</span><span class="p">[</span><span class="n">ll</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span>
                
@@ -145,8 +186,8 @@
 
                                     <span class="c1"># GENERATE SAMPLES AND TARGETS HERE SPECIFICALLY FOR PASCAL VOC 2012</span>
                
 
                                     <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root</span> <span class="o">+</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">sep</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">list_file_path</span><span class="p">,</span> <span class="s2">&quot;r&quot;</span><sp
                
 
                                         <span class="k">for</span> <span class="n">line</span> <span class="ow">in</span> <span class="n">lines</span><span class="p">:</span>
                
 
                            -                <span class="n">image_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">samples_sub_directory</span><span class="p">,</span> <span class="n">line</span><span class="o">.</sp
                
 
                            -                <span class="n">mask_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">targets_sub_directory</span><span class="p">,</span> <span class="n">line</span><span class="o">.</spa
                
 
                            +                <span class="n">image_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">samples_sub_directory</span><span class="p">,</span> <span class="n">line</span><span class="o">.</sp
                
 
                            +                <span class="n">mask_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">root</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">targets_sub_directory</span><span class="p">,</span> <span class="n">line</span><span class="o">.</spa
                
 
                                             <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">mask_path</span><span class="p">)</span> <span class="ow">and</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">image_path</span><span class="p">):</span>
                
 
                                                 <span class="bp">self</span><span class="o">.</span><span class="n">samples_targets_tuples_list</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">image_path</span><span class="p">,</span> <span class="n">mask_path</span><span class="p">))</span>
                
@@ -184,6 +225,59 @@
 
                                             <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">64</span><span class="p">,</span> <span class="mi">128</span><span class="p">],</span>
                
 
                                         <span class="p">]</span>
                
 
                                     <span class="p">)</span></div>
                
 
                            +
                
 
                            +
                
 
                            +<div class="viewcode-block" id="PascalAUG2012SegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalAUG2012SegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">PascalAUG2012SegmentationDataSet</span><span class="p">(</span><span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">):</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">    PascalAUG2012SegmentationDataSet - Segmentation Data Set Class for Pascal AUG 2012 Data Set</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">sample_suffix</span> <span class="o">=</span> <span class="s2">&quot;.jpg&quot;</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">target_suffix</span> <span class="o">=</span> <span class="s2">&quot;.mat&quot;</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">sample_suffix</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">sample_suffix</span><span class="p">,</span> <span class="n">target_suffix</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">target_suffix</span><span class="p">,</span> <span class
                
 
                            +
                
 
                            +<div class="viewcode-block" id="PascalAUG2012SegmentationDataSet.target_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalAUG2012SegmentationDataSet.target_loader">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">target_loader</span><span class="p">(</span><span class="n">target_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Image</span><span class="p">:</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        target_loader</span>
                
 
                            +<span class="sd">            :param target_path: The path to the target data</span>
                
 
                            +<span class="sd">            :return:            The loaded target</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">mat</span> <span class="o">=</span> <span class="n">scipy</span><span class="o">.</span><span class="n">io</span><span class="o">.</span><span class="n">loadmat</span><span class="p">(</span><span class="n">target_path</span><span class="p">,</span> <span class="n">mat_dtype</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">squeeze_me</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span 
                
 
                            +        <span class="n">mask</span> <span class="o">=</span> <span class="n">mat</span><span class="p">[</span><span class="s2">&quot;GTcls&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">Segmentation</span>
                
 
                            +        <span class="k">return</span> <span class="n">Image</span><span class="o">.</span><span class="n">fromarray</span><span class="p">(</span><span class="n">mask</span><span class="p">)</span></div></div>
                
 
                            +
                
 
                            +
                
 
                            +<div class="viewcode-block" id="PascalVOCAndAUGUnifiedDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.PascalVOCAndAUGUnifiedDataset">[docs]</a><span class="k">class</span> <span class="nc">PascalVOCAndAUGUnifiedDataset</span><span class="p">(</span><span class="n">ConcatDataset</span><span class="p">):</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">    Pascal VOC + AUG train dataset, aka `SBD` dataset contributed in &quot;Semantic contours from inverse detectors&quot;.</span>
                
 
                            +<span class="sd">    This is class implement the common usage of the SBD and PascalVOC datasets as a unified augmented trainset.</span>
                
 
                            +<span class="sd">    The unified dataset includes a total of 10,582 samples and don&#39;t contains duplicate samples from the PascalVOC</span>
                
 
                            +<span class="sd">    validation set.</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            +        <span class="nb">print</span><span class="p">(</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            +        <span class="k">if</span> <span class="nb">any</span><span class="p">([</span><span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;list_file&quot;</span><span class="p">),</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;samples_sub_directory&quot;</span><span class="p">),</span> <span class="n">kwargs</span><span class="o">.</span>
                
 
                            +            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
                
 
                            +                <span class="s2">&quot;[list_file, samples_sub_directory, targets_sub_directory] arguments passed will not be used&quot;</span>
                
 
                            +                <span class="s2">&quot; when passed to `PascalVOCAndAUGUnifiedDataset`. Those values are predefined for initiating&quot;</span>
                
 
                            +                <span class="s2">&quot; the Pascal VOC + AUG training set.&quot;</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
                
 
                            +            <span class="n">datasets</span><span class="o">=</span><span class="p">[</span>
                
 
                            +                <span class="n">PascalVOC2012SegmentationDataSet</span><span class="p">(</span>
                
 
                            +                    <span class="n">list_file</span><span class="o">=</span><span class="s2">&quot;VOCdevkit/VOC2012/ImageSets/Segmentation/train.txt&quot;</span><span class="p">,</span>
                
 
                            +                    <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s2">&quot;VOCdevkit/VOC2012/JPEGImages&quot;</span><span class="p">,</span>
                
 
                            +                    <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s2">&quot;VOCdevkit/VOC2012/SegmentationClass&quot;</span><span class="p">,</span>
                
 
                            +                    <span class="o">**</span><span class="n">kwargs</span><span class="p">,</span>
                
 
                            +                <span class="p">),</span>
                
 
                            +                <span class="n">PascalAUG2012SegmentationDataSet</span><span class="p">(</span>
                
 
                            +                    <span class="n">list_file</span><span class="o">=</span><span class="s2">&quot;VOCaug/dataset/aug.txt&quot;</span><span class="p">,</span> <span class="n">samples_sub_directory</span><span class="o">=</span><span class="s2">&quot;VOCaug/dataset/img&quot;</span><span class="p">,</span> <span class="n">targets_sub_directory</span><span class="o">=</span><span class="s2">&quot;VOCaug/dataset/cls&quot;</span><span class="p">,</span> <span class="o">**</span><span class="n">kwarg
                
 
                            +                <span class="p">),</span>
                
 
                            +            <span class="p">]</span>
                
 
                            +        <span class="p">)</span></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -213,4 +307,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.segmentation_datasets.segmentation_dataset &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.segmentation_datasets.segmentation_dataset &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -87,71 +89,44 @@
 
                               <h1>Source code for super_gradients.training.datasets.segmentation_datasets.segmentation_dataset</h1><div class="highlight"><pre>
                
 
                             <span></span><span class="kn">import</span> <span class="nn">os</span>
                
 
                            -<span class="kn">import</span> <span class="nn">torch</span>
                
 
                            -<span class="kn">import</span> <span class="nn">random</span>
                
 
                            +<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">Iterable</span>
                
 
                            +
                
 
                             <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
                
 
                            -<span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                            -<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Callable</span>
                
 
                            +<span class="kn">import</span> <span class="nn">torch</span>
                
 
                             <span class="kn">import</span> <span class="nn">torchvision.transforms</span> <span class="k">as</span> <span class="nn">transform</span>
                
 
                             <span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span>
                
 
                            +<span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.decorators.factory_decorator</span> <span class="kn">import</span> <span class="n">resolve_param</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.factories.transforms_factory</span> <span class="kn">import</span> <span class="n">TransformsFactory</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.datasets.sg_dataset</span> <span class="kn">import</span> <span class="n">DirectoryDataSet</span><span class="p">,</span> <span class="n">ListDataset</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.transforms.transforms</span> <span class="kn">import</span> <span class="n">RandomFlip</span><span class="p">,</span> <span class="n">Rescale</span><span class="p">,</span> <span class="n">RandomRescale</span><span class="p">,</span> <span class="n">RandomRotate</span><span class="p">,</span> \
                
 
                            -    <span class="n">CropImageAndMask</span><span class="p">,</span> <span class="n">RandomGaussianBlur</span><span class="p">,</span> <span class="n">PadShortToCropSize</span>
                
 
                            -<div class="viewcode-block" id="SegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.SegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">SegmentationDataSet</span><span class="p">(</span><span class="n">DirectoryDataSet</span><span class="p">,</span> <span class="n">ListDataset</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="SegmentationDataSet"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.SegmentationDataSet">[docs]</a><span class="k">class</span> <span class="nc">SegmentationDataSet</span><span class="p">(</span><span class="n">DirectoryDataSet</span><span class="p">,</span> <span class="n">ListDataset</span><span class="p">):</span>
                
 
                            -    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s1">&#39;image_mask_transforms&#39;</span><span class="p">,</span> <span class="n">factory</span><span class="o">=</span><span class="n">TransformsFactory</span><span class="p">())</span>
                
 
                            -    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s1">&#39;image_mask_transforms_aug&#39;</span><span class="p">,</span> <span class="n">factory</span><span class="o">=</span><span class="n">TransformsFactory</span><span class="p">())</span>
                
 
                            +    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s1">&#39;transforms&#39;</span><span class="p">,</span> <span class="n">factory</span><span class="o">=</span><span class="n">TransformsFactory</span><span class="p">())</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">root</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">list_file</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">samples_sub_directory</span><span class="p">:</span> <span class="nb">str</s
                
 
                                              <span class="n">targets_sub_directory</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">img_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">608</span><span class="p">,</span> <span class="n">crop_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">512</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">16</span><span class="p">,</spa
                
 
                            -                 <span class="n">dataset_hyper_params</span><span class="p">:</span> <span class="nb">dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">cache_labels</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">cache_images</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">sample_loader</span><span class="p">:</span> <span class="n">Callable</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -                 <span class="n">target_loader</span><span class="p">:</span> <span class="n">Callable</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">collate_fn</span><span class="p">:</span> <span class="n">Callable</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">target_extension</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;.png
                
 
                            -                 <span class="n">image_mask_transforms</span><span class="p">:</span> <span class="n">transform</span><span class="o">.</span><span class="n">Compose</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">image_mask_transforms_aug</span><span class="p">:</span> <span class="n">transform</span><span class="o">.</span><span class="n">Compose</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +                 <span class="n">cache_labels</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">cache_images</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +                 <span class="n">collate_fn</span><span class="p">:</span> <span class="n">Callable</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">target_extension</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;.png&#39;</span><span class="p">,</span>
                
 
                            +                 <span class="n">transforms</span><span class="p">:</span> <span class="n">Iterable</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        SegmentationDataSet</span>
                
 
                            -<span class="sd">                                * Please use self.augment == True only for training</span>
                
 
                            -
                
 
                             <span class="sd">            :param root:                        Root folder of the Data Set</span>
                
 
                             <span class="sd">            :param list_file:                   Path to the file with the samples list</span>
                
 
                             <span class="sd">            :param samples_sub_directory:       name of the samples sub-directory</span>
                
 
                             <span class="sd">            :param targets_sub_directory:       name of the targets sub-directory</span>
                
 
                            -<span class="sd">            :param img_size:                    Image size of the Model that uses this Data Set</span>
                
 
                            -<span class="sd">            :param crop_size:                   The size of the cropped image</span>
                
 
                            -<span class="sd">            :param batch_size:                  Batch Size of the Model that uses this Data Set</span>
                
 
                            -<span class="sd">            :param augment:                     True / False flag to allow Augmentation</span>
                
 
                            -<span class="sd">            :param dataset_hyper_params:        Any hyper params required for the data set</span>
                
 
                             <span class="sd">            :param cache_labels:                &quot;Caches&quot; the labels -&gt; Pre-Loads to memory as a list</span>
                
 
                             <span class="sd">            :param cache_images:                &quot;Caches&quot; the images -&gt; Pre-Loads to memory as a list</span>
                
 
                            -<span class="sd">            :param sample_loader:               A function that specifies how to load a sample</span>
                
 
                            -<span class="sd">            :param target_loader:               A function that specifies how to load a target</span>
                
 
                             <span class="sd">            :param collate_fn:                  collate_fn func to process batches for the Data Loader</span>
                
 
                            -<span class="sd">            :param target_extension:            file extension of the targets (defualt is .png for PASCAL VOC 2012)</span>
                
 
                            -<span class="sd">            :param image_mask_transforms        transforms to be applied on image and mask when augment=False</span>
                
 
                            -<span class="sd">            :param image_mask_transforms_aug    transforms to be applied on image and mask when augment=True</span>
                
 
                            +<span class="sd">            :param target_extension:            file extension of the targets (default is .png for PASCAL VOC 2012)</span>
                
 
                            +<span class="sd">            :param transforms:                  transforms to be applied on image and mask</span>
                
 
                            +
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">samples_sub_directory</span> <span class="o">=</span> <span class="n">samples_sub_directory</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">targets_sub_directory</span> <span class="o">=</span> <span class="n">targets_sub_directory</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_hyperparams</span> <span class="o">=</span> <span class="n">dataset_hyper_params</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">cache_labels</span> <span class="o">=</span> <span class="n">cache_labels</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">cache_images</span> <span class="o">=</span> <span class="n">cache_images</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span> <span class="o">=</span> <span class="n">batch_size</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">img_size</span> <span class="o">=</span> <span class="n">img_size</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span> <span class="o">=</span> <span class="n">crop_size</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">augment</span> <span class="o">=</span> <span class="n">augment</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">batch_index</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">total_batches_num</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -
                
 
                            -        <span class="c1"># ENABLES USING CUSTOM SAMPLE/TARGET LOADERS</span>
                
 
                            -        <span class="k">if</span> <span class="n">sample_loader</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">sample_loader</span> <span class="o">=</span> <span class="n">sample_loader</span>
                
 
                            -        <span class="k">if</span> <span class="n">target_loader</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">target_loader</span> <span class="o">=</span> <span class="n">target_loader</span>
                
 
                                     <span class="c1"># CREATE A DIRECTORY DATASET OR A LIST DATASET BASED ON THE list_file INPUT VARIABLE</span>
                
 
                                     <span class="k">if</span> <span class="n">list_file</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
@@ -165,26 +140,8 @@
 
                                                                   <span class="n">sample_loader</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">sample_loader</span><span class="p">,</span> <span class="n">sample_transform</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">sample_transform</span><span class="p">,</span>
                
 
                                                                   <span class="n">target_loader</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">target_loader</span><span class="p">,</span> <span class="n">target_transform</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">target_transform</span><span class="p">,</span>
                
 
                                                                   <span class="n">collate_fn</span><span class="o">=</span><span class="n">collate_fn</span><span class="p">)</span>
                
 
                            -        <span class="c1"># DEFAULT TRANSFORMS</span>
                
 
                            -        <span class="c1"># FIXME - Rescale before RandomRescale is kept for legacy support, consider removing it like most implementation</span>
                
 
                            -        <span class="c1">#  papers regimes.</span>
                
 
                            -        <span class="n">default_image_mask_transforms_aug</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span><span class="n">RandomFlip</span><span class="p">(),</span>
                
 
                            -                                                               <span class="n">Rescale</span><span class="p">(</span><span class="n">short_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">img_size</span><span class="p">),</span>
                
 
                            -                                                               <span class="n">RandomRescale</span><span class="p">(</span><span class="n">scales</span><span class="o">=</span><span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">)),</span>
                
 
                            -                                                               <span class="n">RandomRotate</span><span class="p">(),</span>
                
 
                            -                                                               <span class="n">PadShortToCropSize</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">),</span>
                
 
                            -                                                               <span class="n">CropImageAndMask</span><span class="p">(</span><span class="n">crop_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">,</span>
                
 
                            -                                                                                <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;random&quot;</span><span class="p">),</span>
                
 
                            -                                                               <span class="n">RandomGaussianBlur</span><span class="p">()])</span>
                
 
                            -
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">image_mask_transforms_aug</span> <span class="o">=</span> <span class="n">image_mask_transforms_aug</span> <span class="ow">or</span> <span class="n">default_image_mask_transforms_aug</span>
                
 
                            -        <span class="c1"># FIXME: CROP SIZE CANNOT BE PASSED WHEN LIST</span>
                
 
                            -        <span class="k">if</span> <span class="n">image_mask_transforms</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="n">image_mask_transforms</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">Compose</span><span class="p">([</span><span class="n">Rescale</span><span class="p">(</span><span class="n">short_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">),</span>
                
 
                            -                                                       <span class="n">CropImageAndMask</span><span class="p">(</span><span class="n">crop_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">)</span>
                
 
                            -                                                       <span class="p">])</span>
                
 
                            -
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">image_mask_transforms</span> <span class="o">=</span> <span class="n">image_mask_transforms</span>
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span> <span class="o">=</span> <span class="n">transform</span><span class="o">.</span><span class="n">Compose</span><span class="p">(</span><span class="n">transforms</span> <span class="k">if</span> <span class="n">transforms</span> <span class="k">else</span> <span class="p">[])</span>
                
 
                                 <span class="k">def</span> <span class="fm">__getitem__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">index</span><span class="p">):</span>
                
 
                                     <span class="n">sample_path</span><span class="p">,</span> <span class="n">target_path</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">samples_targets_tuples_list</span><span class="p">[</span><span class="n">index</span><span class="p">]</span>
                
@@ -206,7 +163,7 @@
 
                                     <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">sample_transform</span><span class="p">(</span><span class="n">sample</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">target_transform</span><span class="p">(</span><span class="n">target</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="SegmentationDataSet.sample_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.SegmentationDataSet.sample_loader">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="SegmentationDataSet.sample_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.SegmentationDataSet.sample_loader">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">sample_loader</span><span class="p">(</span><span class="n">sample_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Image</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        sample_loader - Loads a dataset image from path using PIL</span>
                
@@ -216,7 +173,7 @@
 
                                     <span class="n">image</span> <span class="o">=</span> <span class="n">Image</span><span class="o">.</span><span class="n">open</span><span class="p">(</span><span class="n">sample_path</span><span class="p">)</span><span class="o">.</span><span class="n">convert</span><span class="p">(</span><span class="s1">&#39;RGB&#39;</span><span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="n">image</span></div>
                
 
                            -<div class="viewcode-block" id="SegmentationDataSet.sample_transform"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.SegmentationDataSet.sample_transform">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="SegmentationDataSet.sample_transform"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.SegmentationDataSet.sample_transform">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">sample_transform</span><span class="p">(</span><span class="n">image</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        sample_transform - Transforms the sample image</span>
                
@@ -230,7 +187,7 @@
 
                                     <span class="k">return</span> <span class="n">sample_transform</span><span class="p">(</span><span class="n">image</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="SegmentationDataSet.target_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.SegmentationDataSet.target_loader">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="SegmentationDataSet.target_loader"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.SegmentationDataSet.target_loader">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">target_loader</span><span class="p">(</span><span class="n">target_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Image</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        target_loader</span>
                
@@ -240,7 +197,7 @@
 
                                     <span class="n">target</span> <span class="o">=</span> <span class="n">Image</span><span class="o">.</span><span class="n">open</span><span class="p">(</span><span class="n">target_path</span><span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="n">target</span></div>
                
 
                            -<div class="viewcode-block" id="SegmentationDataSet.target_transform"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.SegmentationDataSet.target_transform">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +<div class="viewcode-block" id="SegmentationDataSet.target_transform"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.SegmentationDataSet.target_transform">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">target_transform</span><span class="p">(</span><span class="n">target</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        target_transform - Transforms the sample image</span>
                
@@ -258,9 +215,6 @@
 
                                     <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">samples_targets_tuples_list</span><span class="p">:</span>
                
 
                                         <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">_generate_samples_and_targets</span><span class="p">()</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">batch_index</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.<
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">total_batches_num</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_index</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span>
                
 
                            -
                
 
                                     <span class="c1"># EXTRACT THE LABELS FROM THE TUPLES LIST</span>
                
 
                                     <span class="n">image_files</span><span class="p">,</span> <span class="n">label_files</span> <span class="o">=</span> <span class="nb">map</span><span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="bp">self</span><span class="o">.</span><span class="n">samples_targets_tuples_list</span><span class="p">))</span>
                
 
                                     <span class="n">image_indices_to_remove</span> <span class="o">=</span> <span class="p">[]</span>
                
@@ -309,66 +263,13 @@
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">label_files</span> <span class="o">=</span> <span class="p">[</span><span class="n">e</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">e</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">label_files</span><span class="p">)</span> <span class="k">if</span> <span class="n">i</span> <span class="ow">not</span> <sp
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">labels</span> <span class="o">=</span> <span class="p">[</span><span class="n">e</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">e</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">)</span> <span class="k">if</span> <span class="n
                
 
                            -    <span class="k">def</span> <span class="nf">_calculate_short_size</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">img</span><span class="p">):</span>
                
 
                            -        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        _calculate_crop</span>
                
 
                            -<span class="sd">        :param img:</span>
                
 
                            -<span class="sd">        :return:</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">augment</span><span class="p">:</span>
                
 
                            -            <span class="c1"># RANDOM SCALE (SHORT EDGE FROM 480 TO 720)</span>
                
 
                            -            <span class="n">short_size</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">img_size</span> <span class="o">*</span> <span class="mf">0.5</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span
                
 
                            -        <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">short_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span>
                
 
                            -
                
 
                            -        <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">img</span><span class="o">.</span><span class="n">size</span>
                
 
                            -        <span class="k">if</span> <span class="n">w</span> <span class="o">&gt;</span> <span class="n">h</span><span class="p">:</span>
                
 
                            -            <span class="n">oh</span> <span class="o">=</span> <span class="n">short_size</span>
                
 
                            -            <span class="n">ow</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">*</span> <span class="n">w</span> <span class="o">*</span> <span class="n">oh</span> <span class="o">/</span> <span class="n">h</span><span class="p">)</span>
                
 
                            -        <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">ow</span> <span class="o">=</span> <span class="n">short_size</span>
                
 
                            -            <span class="n">oh</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">*</span> <span class="n">h</span> <span class="o">*</span> <span class="n">ow</span> <span class="o">/</span> <span class="n">w</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="k">return</span> <span class="n">oh</span><span class="p">,</span> <span class="n">ow</span><span class="p">,</span> <span class="n">short_size</span>
                
 
                            -
                
 
                            -    <span class="k">def</span> <span class="nf">_get_center_crop</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">):</span>
                
 
                            -        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -
                
 
                            -<span class="sd">        :param w:</span>
                
 
                            -<span class="sd">        :param h:</span>
                
 
                            -<span class="sd">        :return:</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="c1"># CENTER CROP</span>
                
 
                            -        <span class="n">x1</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">((</span><span class="n">w</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">)</span> <span class="o">/</span> <span class="mf">2.</span><span class="p">))</span>
                
 
                            -        <span class="n">y1</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">((</span><span class="n">h</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">)</span> <span class="o">/</span> <span class="mf">2.</span><span class="p">))</span>
                
 
                            -
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">augment</span><span class="p">:</span>
                
 
                            -            <span class="c1"># RANDOM CROP CROP_SIZE</span>
                
 
                            -            <span class="n">x1</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">w</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">)</span>
                
 
                            -            <span class="n">y1</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">h</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="k">return</span> <span class="n">x1</span><span class="p">,</span> <span class="n">y1</span>
                
 
                            -
                
 
                                 <span class="k">def</span> <span class="nf">_transform_image_and_mask</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">image</span><span class="p">,</span> <span class="n">mask</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        _transform -  Transforms the input (image, mask) in the following order:</span>
                
 
                            -<span class="sd">                                1. FLIP (if augment==true)</span>
                
 
                            -<span class="sd">                                2. RESIZE</span>
                
 
                            -<span class="sd">                                3. ROTATE (if augment==true)</span>
                
 
                            -<span class="sd">                                4. CROP</span>
                
 
                            -<span class="sd">                                5. GAUSSIAN BLUR (if augment==true)</span>
                
 
                            -
                
 
                            -<span class="sd">                            * Please use self.augment == True only for training</span>
                
 
                            -
                
 
                             <span class="sd">            :param image:           The input image</span>
                
 
                             <span class="sd">            :param mask:            The input mask</span>
                
 
                             <span class="sd">            :return:                The transformed image, mask</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">augment</span><span class="p">:</span>
                
 
                            -            <span class="n">transformed</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_mask_transforms_aug</span><span class="p">({</span><span class="s2">&quot;image&quot;</span><span class="p">:</span> <span class="n">image</span><span class="p">,</span> <span class="s2">&quot;mask&quot;</span><span class="p">:</span> <span class="n">mask</span><span class="p">})</span>
                
 
                            -        <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">transformed</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">image_mask_transforms</span><span class="p">({</span><span class="s2">&quot;image&quot;</span><span class="p">:</span> <span class="n">image</span><span class="p">,</span> <span class="s2">&quot;mask&quot;</span><span class="p">:</span> <span class="n">mask</span><span class="p">})</span>
                
 
                            -
                
 
                            +        <span class="n">transformed</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">transforms</span><span class="p">({</span><span class="s2">&quot;image&quot;</span><span class="p">:</span> <span class="n">image</span><span class="p">,</span> <span class="s2">&quot;mask&quot;</span><span class="p">:</span> <span class="n">mask</span><span class="p">})</span>
                
 
                                     <span class="k">return</span> <span class="n">transformed</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">transformed</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span></div>
                
 
                             </pre></div>
                
@@ -399,4 +300,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.segmentation_datasets.supervisely_persons_segmentation &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.segmentation_datasets.supervisely_persons_segmentation &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -92,7 +94,7 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.datasets.segmentation_datasets.segmentation_dataset</span> <span class="kn">import</span> <span class="n">SegmentationDataSet</span>
                
 
                            -<div class="viewcode-block" id="SuperviselyPersonsDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.datasets.segmentation_datasets.html#super_gradients.training.datasets.segmentation_datasets.SuperviselyPersonsDataset">[docs]</a><span class="k">class</span> <span class="nc">SuperviselyPersonsDataset</span><span class="p">(</span><span class="n">SegmentationDataSet</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="SuperviselyPersonsDataset"><a class="viewcode-back" href="../../../../../super_gradients.training.html#super_gradients.training.datasets.SuperviselyPersonsDataset">[docs]</a><span class="k">class</span> <span class="nc">SuperviselyPersonsDataset</span><span class="p">(</span><span class="n">SegmentationDataSet</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    SuperviselyPersonsDataset - Segmentation Data Set Class for Supervisely Persons Segmentation Data Set,</span>
                
 
                             <span class="sd">    main resolution of dataset: (600 x 800).</span>
                
@@ -126,7 +128,8 @@
 
                                                 <span class="bp">self</span><span class="o">.</span><span class="n">samples_targets_tuples_list</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">sample_path</span><span class="p">,</span> <span class="n">target_path</span><span class="p">))</span>
                
 
                                             <span class="k">else</span><span class="p">:</span>
                
 
                                                 <span class="k">raise</span> <span class="ne">AssertionError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Sample and/or target file(s) not found or in illegal format &quot;</span>
                
 
                            -                                         <span class="sa">f</span><span class="s2">&quot;(sample path: </span><span class="si">{</span><span class="n">sample_path</span><span class="si">}</span><span class="s2">, target path: </span><span class="si">{</span><span class="n">target_path</span><span class="si">}</span><span class="s2">)&quot;</span><span class="p">)</span></div>
                
 
                            +                                         <span class="sa">f</span><span class="s2">&quot;(sample path: </span><span class="si">{</span><span class="n">sample_path</span><span class="si">}</span><span class="s2">, target path: </span><span class="si">{</span><span class="n">target_path</span><span class="si">}</span><span class="s2">)&quot;</span><span class="p">)</span>
                
 
                            +        <span class="nb">super</span><span class="p">(</span><span class="n">SuperviselyPersonsDataset</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">_generate_samples_and_targets</span><span class="p">()</span></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -156,4 +159,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.datasets.sg_dataset &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.datasets.sg_dataset &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -97,7 +99,7 @@
 
                             <span class="n">IMG_EXTENSIONS</span> <span class="o">=</span> <span class="p">(</span><span class="s1">&#39;.jpg&#39;</span><span class="p">,</span> <span class="s1">&#39;.jpeg&#39;</span><span class="p">,</span> <span class="s1">&#39;.png&#39;</span><span class="p">,</span> <span class="s1">&#39;.ppm&#39;</span><span class="p">,</span> <span class="s1">&#39;.bmp&#39;</span><span class="p">,</span> <span class="s1">&#39;.pgm&#39;</span><span class="p">,</span> <span class="s1">&#39;.tif&#39;</
                
 
                            -<div class="viewcode-block" id="BaseSgVisionDataset"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.BaseSgVisionDataset">[docs]</a><span class="k">class</span> <span class="nc">BaseSgVisionDataset</span><span class="p">(</span><span class="n">VisionDataset</span><span class="p">):</span>
                
 
                            +<span class="k">class</span> <span class="nc">BaseSgVisionDataset</span><span class="p">(</span><span class="n">VisionDataset</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    BaseSgVisionDataset</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
@@ -160,16 +162,16 @@
 
                                     <span class="k">return</span> <span class="kc">False</span>
                
 
                            -<div class="viewcode-block" id="BaseSgVisionDataset.numpy_loader_func"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.BaseSgVisionDataset.numpy_loader_func">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">numpy_loader_func</span><span class="p">(</span><span class="n">path</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        _numpy_loader_func - Uses numpy load func</span>
                
 
                             <span class="sd">            :param path:</span>
                
 
                             <span class="sd">            :return:</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">path</span><span class="p">)</span></div>
                
 
                            +        <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">path</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="BaseSgVisionDataset.text_file_loader_func"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.BaseSgVisionDataset.text_file_loader_func">[docs]</a>    <span class="nd">@staticmethod</span>
                
 
                            +    <span class="nd">@staticmethod</span>
                
 
                                 <span class="k">def</span> <span class="nf">text_file_loader_func</span><span class="p">(</span><span class="n">text_file_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">inline_splitter</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39; &#39;</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        text_file_loader_func - Uses a line by line based code to get vectorized data from a text-based file</span>
                
@@ -184,10 +186,10 @@
 
                                     <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">text_file_path</span><span class="p">,</span> <span class="s2">&quot;r&quot;</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="s2">&quot;utf-8&quot;</span><span class="p">)</span> <span class="k">as</span> <span class="n">text_file</span><span class="p">:</span>
                
 
                                         <span class="n">targets_list</span> <span class="o">=</span> <span class="p">[</span><span class="nb">tuple</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="nb">float</span><span class="p">,</span> <span class="n">line</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">inline_splitter</span><span class="p">)))</span> <span class="k">for</span> <span class="n">line</span> <span class="ow">in</
                
 
                            -        <span class="k">return</span> <span class="n">targets_list</span></div></div>
                
 
                            +        <span class="k">return</span> <span class="n">targets_list</span>
                
 
                            -<div class="viewcode-block" id="DirectoryDataSet"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.DirectoryDataSet">[docs]</a><span class="k">class</span> <span class="nc">DirectoryDataSet</span><span class="p">(</span><span class="n">BaseSgVisionDataset</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="DirectoryDataSet"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.datasets.DirectoryDataSet">[docs]</a><span class="k">class</span> <span class="nc">DirectoryDataSet</span><span class="p">(</span><span class="n">BaseSgVisionDataset</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    DirectoryDataSet - A PyTorch Vision Data Set extension that receives a root Dir and two separate sub directories:</span>
                
 
                             <span class="sd">                        - Sub-Directory for Samples</span>
                
@@ -279,7 +281,7 @@
 
                                             <span class="nb">print</span><span class="p">(</span><span class="vm">__name__</span> <span class="o">+</span> <span class="s1">&#39; There are &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">missing_files_counter</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39; missing  &#39;</span> <span class="o">+</span> <span class="n">counter_name</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="ListDataset"><a class="viewcode-back" href="../../../../super_gradients.training.datasets.html#super_gradients.training.datasets.ListDataset">[docs]</a><span class="k">class</span> <span class="nc">ListDataset</span><span class="p">(</span><span class="n">BaseSgVisionDataset</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="ListDataset"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.datasets.ListDataset">[docs]</a><span class="k">class</span> <span class="nc">ListDataset</span><span class="p">(</span><span class="n">BaseSgVisionDataset</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    ListDataset - A PyTorch Vision Data Set extension that receives a file with FULL PATH to each of the samples.</span>
                
 
                             <span class="sd">                  Then, the assumption is that for every sample, there is a * matching target * in the same</span>
                
@@ -383,4 +385,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.exceptions.dataset_exceptions &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.exceptions.dataset_exceptions</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.exceptions.dataset_exceptions</h1><div class="highlight"><pre>
<span></span>
<div class="viewcode-block" id="IllegalDatasetParameterException"><a class="viewcode-back" href="../../../../super_gradients.training.exceptions.html#super_gradients.training.exceptions.dataset_exceptions.IllegalDatasetParameterException">[docs]</a><span class="k">class</span> <span class="nc">IllegalDatasetParameterException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Exception raised illegal dataset param.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        message -- explanation of the error</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">desc</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">message</span> <span class="o">=</span> <span class="s2">&quot;Unsupported dataset parameter format: &quot;</span> <span class="o">+</span> <span class="n">desc</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">message</span><span class="p">)</span></div>


<div class="viewcode-block" id="EmptyDatasetException"><a class="viewcode-back" href="../../../../super_gradients.training.exceptions.html#super_gradients.training.exceptions.dataset_exceptions.EmptyDatasetException">[docs]</a><span class="k">class</span> <span class="nc">EmptyDatasetException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Exception raised when a dataset does not have any image for a specific config</span>

<span class="sd">    Attributes:</span>
<span class="sd">        message -- explanation of the error</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">desc</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">message</span> <span class="o">=</span> <span class="s2">&quot;Empty Dataset: &quot;</span> <span class="o">+</span> <span class="n">desc</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">message</span><span class="p">)</span></div>


<div class="viewcode-block" id="UnsupportedBatchItemsFormat"><a class="viewcode-back" href="../../../../super_gradients.training.exceptions.html#super_gradients.training.exceptions.dataset_exceptions.UnsupportedBatchItemsFormat">[docs]</a><span class="k">class</span> <span class="nc">UnsupportedBatchItemsFormat</span><span class="p">(</span><span class="ne">ValueError</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Exception raised illegal batch items returned from data loader.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        message -- explanation of the error</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">message</span> <span class="o">=</span> <span class="s2">&quot;Batch items returned by the data loader expected format: </span><span class="se">\n</span><span class="s2">&quot;</span> \
                       <span class="s2">&quot;1. torch.Tensor or tuple, s.t inputs = batch_items[0], targets = batch_items[1] and len(&quot;</span> \
                       <span class="s2">&quot;batch_items) = 2 </span><span class="se">\n</span><span class="s2">&quot;</span> \
                       <span class="s2">&quot;2. tuple: (inputs, targets, additional_batch_items)&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">message</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.exceptions.sg_model_exceptions &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.exceptions.sg_model_exceptions</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.exceptions.sg_model_exceptions</h1><div class="highlight"><pre>
<span></span>
<div class="viewcode-block" id="UnsupportedTrainingParameterFormat"><a class="viewcode-back" href="../../../../super_gradients.training.exceptions.html#super_gradients.training.exceptions.sg_model_exceptions.UnsupportedTrainingParameterFormat">[docs]</a><span class="k">class</span> <span class="nc">UnsupportedTrainingParameterFormat</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Exception raised illegal training param format.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        message -- explanation of the error</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">desc</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">message</span> <span class="o">=</span> <span class="s2">&quot;Unsupported training parameter format: &quot;</span> <span class="o">+</span> <span class="n">desc</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">message</span><span class="p">)</span></div>


<div class="viewcode-block" id="UnsupportedOptimizerFormat"><a class="viewcode-back" href="../../../../super_gradients.training.exceptions.html#super_gradients.training.exceptions.sg_model_exceptions.UnsupportedOptimizerFormat">[docs]</a><span class="k">class</span> <span class="nc">UnsupportedOptimizerFormat</span><span class="p">(</span><span class="n">UnsupportedTrainingParameterFormat</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Exception raised illegal optimizer format.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        message -- explanation of the error</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
            <span class="s2">&quot;optimizer parameter expected one of [&#39;Adam&#39;,&#39;SGD&#39;,&#39;RMSProp&#39;], or torch.optim.Optimizer object&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="IllegalDataloaderInitialization"><a class="viewcode-back" href="../../../../super_gradients.training.exceptions.html#super_gradients.training.exceptions.sg_model_exceptions.IllegalDataloaderInitialization">[docs]</a><span class="k">class</span> <span class="nc">IllegalDataloaderInitialization</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Exception raised illegal data loaders.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
            <span class="s2">&quot;train_loader, valid_loader and class parameters are required when initializing SgModel with data loaders&quot;</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.kd_model.kd_model &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.kd_trainer.kd_trainer &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -76,7 +78,7 @@
 
                               <ul class="wy-breadcrumbs">
                
 
                                   <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
                
 
                                       <li><a href="../../../index.html">Module code</a> &raquo;</li>
                
 
                            -      <li>super_gradients.training.kd_model.kd_model</li>
                
 
                            +      <li>super_gradients.training.kd_trainer.kd_trainer</li>
                
 
                                   <li class="wy-breadcrumbs-aside">
                
 
                                   </li>
                
 
                               </ul>
                
@@ -85,101 +87,85 @@
 
                                       <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
                
 
                                        <div itemprop="articleBody">
                
 
                            -  <h1>Source code for super_gradients.training.kd_model.kd_model</h1><div class="highlight"><pre>
                
 
                            -<span></span><span class="kn">import</span> <span class="nn">torch.nn</span>
                
 
                            -
                
 
                            +  <h1>Source code for super_gradients.training.kd_trainer.kd_trainer</h1><div class="highlight"><pre>
                
 
                            +<span></span><span class="kn">import</span> <span class="nn">hydra</span>
                
 
                            +<span class="kn">import</span> <span class="nn">torch.nn</span>
                
 
                            +<span class="kn">from</span> <span class="nn">omegaconf</span> <span class="kn">import</span> <span class="n">DictConfig</span>
                
 
                            +<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">DataLoader</span>
                
 
                            +
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common</span> <span class="kn">import</span> <span class="n">MultiGPUMode</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.dataloaders</span> <span class="kn">import</span> <span class="n">dataloaders</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.models</span> <span class="kn">import</span> <span class="n">SgModule</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.models.all_architectures</span> <span class="kn">import</span> <span class="n">KD_ARCHITECTURES</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.models.kd_modules.kd_module</span> <span class="kn">import</span> <span class="n">KDModule</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.sg_model</span> <span class="kn">import</span> <span class="n">SgModel</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.sg_trainer</span> <span class="kn">import</span> <span class="n">Trainer</span>
                
 
                             <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span><span class="p">,</span> <span class="n">models</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.pretrained_models</span> <span class="kn">import</span> <span class="n">PRETRAINED_NUM_CLASSES</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">get_param</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">get_param</span><span class="p">,</span> <span class="n">HpmStruct</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.checkpoint_utils</span> <span class="kn">import</span> <span class="n">read_ckpt_state_dict</span><span class="p">,</span> \
                
 
                                 <span class="n">load_checkpoint_to_model</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.kd_model_exceptions</span> <span class="kn">import</span> <span class="n">ArchitectureKwargsException</span><span class="p">,</span> \
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.kd_trainer_exceptions</span> <span class="kn">import</span> <span class="n">ArchitectureKwargsException</span><span class="p">,</span> \
                
 
                                 <span class="n">UnsupportedKDArchitectureException</span><span class="p">,</span> <span class="n">InconsistentParamsException</span><span class="p">,</span> <span class="n">UnsupportedKDModelArgException</span><span class="p">,</span> \
                
 
                                 <span class="n">TeacherKnowledgeException</span><span class="p">,</span> <span class="n">UndefinedNumClassesException</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.callbacks</span> <span class="kn">import</span> <span class="n">KDModelMetricsUpdateCallback</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.ema</span> <span class="kn">import</span> <span class="n">KDModelEMA</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils.sg_trainer_utils</span> <span class="kn">import</span> <span class="n">parse_args</span>
                
 
                            +
                
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="KDModel"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.KDModel">[docs]</a><span class="k">class</span> <span class="nc">KDModel</span><span class="p">(</span><span class="n">SgModel</span><span class="p">):</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            -        <span class="nb">super</span><span class="p">(</span><span class="n">KDModel</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><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            +<div class="viewcode-block" id="KDTrainer"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.KDTrainer">[docs]</a><span class="k">class</span> <span class="nc">KDTrainer</span><span class="p">(</span><span class="n">Trainer</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">device</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">multi_gpu</span><span class="p">:</span> <span class="n">Union</span
                
 
                            +                 <span class="n">ckpt_root_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">experiment_name</span><span class="o">=</span><span class="n">experiment_name</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">multi_gpu</span><span class="o">=</span><span class="n">multi_gpu</span><span class="p">,</span> <span class=
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">student_architecture</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">teacher_architecture</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">student_arch_params</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">teacher_arch_params</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -<div class="viewcode-block" id="KDModel.build_model"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.KDModel.build_model">[docs]</a>    <span class="k">def</span> <span class="nf">build_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                            -                    <span class="c1"># noqa: C901 - too complex</span>
                
 
                            -                    <span class="n">architecture</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">KDModule</span><span class="p">]</span> <span class="o">=</span> <span class="s1">&#39;kd_module&#39;</span><span class="p">,</span>
                
 
                            -                    <span class="n">arch_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">checkpoint_params</span><span class="o">=</span><span class="p">{},</span>
                
 
                            -                    <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="KDTrainer.train_from_config"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.KDTrainer.train_from_config">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">train_from_config</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">cfg</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">DictConfig</span><span class="p">,</span> <span class="nb">dict</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        :param architecture: (Union[str, KDModule]) Defines the network&#39;s architecture from models/KD_ARCHITECTURES</span>
                
 
                            -<span class="sd">         (default=&#39;kd_module&#39;)</span>
                
 
                            -
                
 
                            -<span class="sd">        :param arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc to be passed to kd</span>
                
 
                            -<span class="sd">            architecture class (discarded when architecture is KDModule instance)</span>
                
 
                            -
                
 
                            -<span class="sd">        :param checkpoint_params: (dict) A dictionary like object with the following keys/values:</span>
                
 
                            -
                
 
                            -<span class="sd">              student_pretrained_weights:   String describing the dataset of the pretrained weights (for example</span>
                
 
                            -<span class="sd">              &quot;imagenent&quot;) for the student network.</span>
                
 
                            -
                
 
                            -<span class="sd">              teacher_pretrained_weights:   String describing the dataset of the pretrained weights (for example</span>
                
 
                            -<span class="sd">              &quot;imagenent&quot;) for the teacher network.</span>
                
 
                            -
                
 
                            -<span class="sd">              teacher_checkpoint_path:    Local path to the teacher&#39;s checkpoint. Note that when passing pretrained_weights</span>
                
 
                            -<span class="sd">                                   through teacher_arch_params these weights will be overridden by the</span>
                
 
                            -<span class="sd">                                   pretrained checkpoint. (default=None)</span>
                
 
                            -
                
 
                            -<span class="sd">              load_kd_model_checkpoint:   Whether to load an entire KDModule checkpoint (used to continue KD training)</span>
                
 
                            -<span class="sd">               (default=False)</span>
                
 
                            -
                
 
                            -<span class="sd">              kd_model_source_ckpt_folder_name: Folder name to load an entire KDModule checkpoint from</span>
                
 
                            -<span class="sd">                (self.experiment_name if none is given) to resume KD training (default=None)</span>
                
 
                            +<span class="sd">        Trains according to cfg recipe configuration.</span>
                
 
                            -<span class="sd">              kd_model_external_checkpoint_path: The path to the external checkpoint to be loaded. Can be absolute or relative</span>
                
 
                            -<span class="sd">                                               (ie: path/to/checkpoint.pth). If provided, will automatically attempt to</span>
                
 
                            -<span class="sd">                                               load the checkpoint even if the load_checkpoint flag is not provided.</span>
                
 
                            -<span class="sd">                                               (deafult=None)</span>
                
 
                            -
                
 
                            -<span class="sd">        :keyword student_architecture: (Union[str, SgModule]) Defines the student&#39;s architecture from</span>
                
 
                            -<span class="sd">            models/ALL_ARCHITECTURES (when str), or directly defined the student network (when SgModule).</span>
                
 
                            -
                
 
                            -<span class="sd">        :keyword teacher_architecture: (Union[str, SgModule]) Defines the teacher&#39;s architecture from</span>
                
 
                            -<span class="sd">            models/ALL_ARCHITECTURES (when str), or directly defined the teacher network (when SgModule).</span>
                
 
                            -
                
 
                            -<span class="sd">        :keyword student_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for student</span>
                
 
                            -<span class="sd">            net. (deafult={})</span>
                
 
                            +<span class="sd">        @param cfg: The parsed DictConfig from yaml recipe files</span>
                
 
                            +<span class="sd">        @return: output of kd_trainer.train(...) (i.e results tuple)</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="c1"># INSTANTIATE ALL OBJECTS IN CFG</span>
                
 
                            +        <span class="n">cfg</span> <span class="o">=</span> <span class="n">hydra</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">instantiate</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
                
 
                            -<span class="sd">        :keyword teacher_arch_params: (dict) Architecture H.P. e.g.: block, num_blocks, num_classes, etc for teacher</span>
                
 
                            -<span class="sd">            net. (deafult={})</span>
                
 
                            +        <span class="n">kwargs</span> <span class="o">=</span> <span class="n">parse_args</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="bp">cls</span><span class="o">.</span><span class="fm">__init__</span><span class="p">)</span>
                
 
                            -<span class="sd">        :keyword run_teacher_on_eval: (bool)- whether to run self.teacher at eval mode regardless of self.train(mode)</span>
                
 
                            +        <span class="n">trainer</span> <span class="o">=</span> <span class="n">KDTrainer</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            +        <span class="c1"># INSTANTIATE DATA LOADERS</span>
                
 
                            +        <span class="n">train_dataloader</span> <span class="o">=</span> <span class="n">dataloaders</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">train_dataloader</span><span class="p">,</span>
                
 
                            +                                           <span class="n">dataset_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_dataset_params</span><span class="p">,</span>
                
 
                            +                                           <span class="n">dataloader_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_dataloader_params</span><span class="p">)</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">kwargs</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s2">&quot;student_architecture&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
                
 
                            -        <span class="n">kwargs</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s2">&quot;teacher_architecture&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
                
 
                            -        <span class="n">kwargs</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s2">&quot;student_arch_params&quot;</span><span class="p">,</span> <span class="p">{})</span>
                
 
                            -        <span class="n">kwargs</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s2">&quot;teacher_arch_params&quot;</span><span class="p">,</span> <span class="p">{})</span>
                
 
                            -        <span class="n">kwargs</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s2">&quot;run_teacher_on_eval&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
                
 
                            +        <span class="n">val_dataloader</span> <span class="o">=</span> <span class="n">dataloaders</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">val_dataloader</span><span class="p">,</span>
                
 
                            +                                         <span class="n">dataset_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_dataset_params</span><span class="p">,</span>
                
 
                            +                                         <span class="n">dataloader_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_dataloader_params</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">_validate_args</span><span class="p">(</span><span class="n">arch_params</span><span class="p">,</span> <span class="n">architecture</span><span class="p">,</span> <span class="n">checkpoint_params</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            +        <span class="n">student</span> <span class="o">=</span> <span class="n">models</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">student_architecture</span><span class="p">,</span> <span class="n">arch_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">student_arch_params</span><span class="p">,</span>
                
 
                            +                             <span class="n">strict_load</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">student_checkpoint_params</span><span class="o">.</span><span class="n">strict_load</span><span class="p">,</span>
                
 
                            +                             <span class="n">pretrained_weights</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">student_checkpoint_params</span><span class="o">.</span><span class="n">pretrained_weights</span><span class="p">,</span>
                
 
                            +                             <span class="n">checkpoint_path</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">student_checkpoint_params</span><span class="o">.</span><span class="n">checkpoint_path</span><span class="p">,</span>
                
 
                            +                             <span class="n">load_backbone</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">student_checkpoint_params</span><span class="o">.</span><span class="n">load_backbone</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">student_architecture</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;student_architecture&quot;</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">teacher_architecture</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;teacher_architecture&quot;</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">student_arch_params</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;student_arch_params&quot;</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">teacher_arch_params</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;teacher_arch_params&quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">teacher</span> <span class="o">=</span> <span class="n">models</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">cfg</span><span class="o">.</span><span class="n">teacher_architecture</span><span class="p">,</span> <span class="n">arch_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">teacher_arch_params</span><span class="p">,</span>
                
 
                            +                             <span class="n">strict_load</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">teacher_checkpoint_params</span><span class="o">.</span><span class="n">strict_load</span><span class="p">,</span>
                
 
                            +                             <span class="n">pretrained_weights</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">teacher_checkpoint_params</span><span class="o">.</span><span class="n">pretrained_weights</span><span class="p">,</span>
                
 
                            +                             <span class="n">checkpoint_path</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">teacher_checkpoint_params</span><span class="o">.</span><span class="n">checkpoint_path</span><span class="p">,</span>
                
 
                            +                             <span class="n">load_backbone</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">teacher_checkpoint_params</span><span class="o">.</span><span class="n">load_backbone</span><span class="p">)</span>
                
 
                            -        <span class="nb">super</span><span class="p">(</span><span class="n">KDModel</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">build_model</span><span class="p">(</span><span class="n">architecture</span><span class="o">=</span><span class="n">architecture</span><span class="p">,</span> <span class="n">arch_params</span><span class="o">=</span><span class="n">arch_params</span><span class="p">,</span>
                
 
                            -                                         <span class="n">checkpoint_params</span><span class="o">=</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>
                
 
                            +        <span class="c1"># TRAIN</span>
                
 
                            +        <span class="n">trainer</span><span class="o">.</span><span class="n">train</span><span class="p">(</span><span class="n">training_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">training_hyperparams</span><span class="p">,</span> <span class="n">student</span><span class="o">=</span><span class="n">student</span><span class="p">,</span> <span class="n">teacher</span><span class="o">=</span><span class="n">teacher</span><span class=
                
 
                            +                      <span class="n">kd_architecture</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">architecture</span><span class="p">,</span> <span class="n">kd_arch_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">arch_params</span><span class="p">,</span>
                
 
                            +                      <span class="n">run_teacher_on_eval</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">run_teacher_on_eval</span><span class="p">,</span>
                
 
                            +                      <span class="n">train_loader</span><span class="o">=</span><span class="n">train_dataloader</span><span class="p">,</span> <span class="n">valid_loader</span><span class="o">=</span><span class="n">val_dataloader</span><span class="p">)</span></div>
                
 
                                 <span class="k">def</span> <span class="nf">_validate_args</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">arch_params</span><span class="p">,</span> <span class="n">architecture</span><span class="p">,</span> <span class="n">checkpoint_params</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                                     <span class="n">student_architecture</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">kwargs</span><span class="p">,</span> <span class="s2">&quot;student_architecture&quot;</span><span class="p">)</span>
                
@@ -215,7 +201,8 @@
 
                                     <span class="n">load_kd_model_checkpoint</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s2">&quot;load_checkpoint&quot;</span><span class="p">)</span>
                
 
                                     <span class="c1"># CHECK THAT TEACHER NETWORK HOLDS KNOWLEDGE FOR THE STUDENT TO LEARN FROM OR THAT WE ARE LOADING AN ENTIRE KD</span>
                
 
                            -        <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">teacher_pretrained_weights</span> <span class="ow">or</span> <span class="n">teacher_checkpoint_path</span> <span class="ow">or</span> <span class="n">load_kd_model_checkpoint</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">teacher_architecture</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span cla
                
 
                            +        <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">teacher_pretrained_weights</span> <span class="ow">or</span> <span class="n">teacher_checkpoint_path</span> <span class="ow">or</span> <span class="n">load_kd_model_checkpoint</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span>
                
 
                            +                <span class="n">teacher_architecture</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">)):</span>
                
 
                                         <span class="k">raise</span> <span class="n">TeacherKnowledgeException</span><span class="p">()</span>
                
 
                                 <span class="k">def</span> <span class="nf">_validate_num_classes</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">student_arch_params</span><span class="p">,</span> <span class="n">teacher_arch_params</span><span class="p">):</span>
                
@@ -277,6 +264,9 @@
 
                                     <span class="n">run_teacher_on_eval</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">kwargs</span><span class="p">,</span> <span class="s2">&quot;run_teacher_on_eval&quot;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
                
 
                            +        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_instantiate_kd_net</span><span class="p">(</span><span class="n">arch_params</span><span class="p">,</span> <span class="n">architecture</span><span class="p">,</span> <span class="n">run_teacher_on_eval</span><span class="p">,</span> <span class="n">student</span><span class="p">,</span> <span class="n">teacher</span><span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_instantiate_kd_net</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">arch_params</span><span class="p">,</span> <span class="n">architecture</span><span class="p">,</span> <span class="n">run_teacher_on_eval</span><span class="p">,</span> <span class="n">student</span><span class="p">,</span> <span class="n">teacher</span><span class="p">):</span>
                
 
                                     <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">architecture</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
                
 
                                         <span class="n">architecture_cls</span> <span class="o">=</span> <span class="n">KD_ARCHITECTURES</span><span class="p">[</span><span class="n">architecture</span><span class="p">]</span>
                
 
                                         <span class="n">net</span> <span class="o">=</span> <span class="n">architecture_cls</span><span class="p">(</span><span class="n">arch_params</span><span class="o">=</span><span class="n">arch_params</span><span class="p">,</span> <span class="n">student</span><span class="o">=</span><span class="n">student</span><span class="p">,</span> <span class="n">teacher</span><span class="o">=</span><span class="n">teacher</span><span class="p">,</span>
                
@@ -286,13 +276,12 @@
 
                                                            <span class="n">run_teacher_on_eval</span><span class="o">=</span><span class="n">run_teacher_on_eval</span><span class="p">)</span>
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="n">net</span> <span class="o">=</span> <span class="n">architecture</span>
                
 
                            -
                
 
                                     <span class="k">return</span> <span class="n">net</span>
                
 
                                 <span class="k">def</span> <span class="nf">_load_checkpoint_to_model</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Initializes teacher weights with teacher_checkpoint_path if needed, then handles checkpoint loading for</span>
                
 
                            -<span class="sd">         the entire KD network following the same logic as in SgModel.</span>
                
 
                            +<span class="sd">         the entire KD network following the same logic as in Trainer.</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="n">teacher_checkpoint_path</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s2">&quot;teacher_checkpoint_path&quot;</span><span class="p">)</span>
                
 
                                     <span class="n">teacher_net</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">teacher</span>
                
@@ -315,7 +304,7 @@
 
                                                                  <span class="n">load_weights_only</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                
 
                                                                  <span class="n">load_ema_as_net</span><span class="o">=</span><span class="n">load_teachers_ema</span><span class="p">)</span>
                
 
                            -        <span class="nb">super</span><span class="p">(</span><span class="n">KDModel</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">_load_checkpoint_to_model</span><span class="p">()</span>
                
 
                            +        <span class="nb">super</span><span class="p">(</span><span class="n">KDTrainer</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">_load_checkpoint_to_model</span><span class="p">()</span>
                
 
                                 <span class="k">def</span> <span class="nf">_add_metrics_update_callback</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
@@ -337,7 +326,8 @@
 
                                                                <span class="p">})</span>
                
 
                                     <span class="k">return</span> <span class="n">hyper_param_config</span>
                
 
                            -    <span class="k">def</span> <span class="nf">_instantiate_ema_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">decay</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.9999</span><span class="p">,</span> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">15</span><span class="p">,</span> <span class="n">exp_acti
                
 
                            +    <span class="k">def</span> <span class="nf">_instantiate_ema_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">decay</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.9999</span><span class="p">,</span> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">15</span><span class="p">,</span>
                
 
                            +                               <span class="n">exp_activation</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">KDModelEMA</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Instantiate KD ema model for KDModule.</span>
                
 
                             <span class="sd">        If the model is of class KDModule, the instance will be adapted to work on knowledge distillation.</span>
                
@@ -360,7 +350,37 @@
 
                                         <span class="n">best_net</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">WrappedModel</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">student</span><span class="p">)</span>
                
 
                                     <span class="n">state</span><span class="p">[</span><span class="s2">&quot;net&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">best_net</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ckpt_best_name</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="n">state</span><span class="p">,</span> <span class="n">global_step</span><span class="o
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ckpt_best_name</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="n">state</span><span class="p">,</span> <span class="n">global_step</span><span class="o
                
 
                            +
                
 
                            +<div class="viewcode-block" id="KDTrainer.train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.KDTrainer.train">[docs]</a>    <span class="k">def</span> <span class="nf">train</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">:</span> <span class="n">KDModule</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">training_
                
 
                            +              <span class="n">teacher</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">kd_architecture</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">KDModule</span><span class="o">.</span><span class="vm">__class__</span><span class="p">,<
                
 
                            +              <span class="n">kd_arch_params</span><span class="p">:</span> <span class="nb">dict</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(),</span> <span class="n">run_teacher_on_eval</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">train_loader</span><span class="p">:</span> <span class="n">DataLoader</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +              <span class="n">valid_loader</span><span class="p">:</span> <span class="n">DataLoader</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Trains the student network (wrapped in KDModule network).</span>
                
 
                            +
                
 
                            +<span class="sd">        :param model: KDModule, network to train. When none is given will initialize KDModule according to kd_architecture,</span>
                
 
                            +<span class="sd">            student and teacher (default=None)</span>
                
 
                            +<span class="sd">        :param training_params: dict, Same as in Trainer.train()</span>
                
 
                            +<span class="sd">        :param student: SgModule - the student trainer</span>
                
 
                            +<span class="sd">        :param teacher: torch.nn.Module- the teacher trainer</span>
                
 
                            +<span class="sd">        :param kd_architecture: KDModule architecture to use, currently only &#39;kd_module&#39; is supported (default=&#39;kd_module&#39;).</span>
                
 
                            +<span class="sd">        :param kd_arch_params: architecture params to pas to kd_architecture constructor.</span>
                
 
                            +<span class="sd">        :param run_teacher_on_eval: bool- whether to run self.teacher at eval mode regardless of self.train(mode)</span>
                
 
                            +<span class="sd">        :param train_loader: Dataloader for train set.</span>
                
 
                            +<span class="sd">        :param valid_loader: Dataloader for validation.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">kd_net</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">net</span> <span class="ow">or</span> <span class="n">model</span>
                
 
                            +        <span class="k">if</span> <span class="n">kd_net</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="k">if</span> <span class="n">student</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">teacher</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Must pass student and teacher models or net (KDModule).&quot;</span><span class="p">)</span>
                
 
                            +            <span class="n">kd_net</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_instantiate_kd_net</span><span class="p">(</span><span class="n">arch_params</span><span class="o">=</span><span class="n">HpmStruct</span><span class="p">(</span><span class="o">**</span><span class="n">kd_arch_params</span><span class="p">),</span>
                
 
                            +                                              <span class="n">architecture</span><span class="o">=</span><span class="n">kd_architecture</span><span class="p">,</span>
                
 
                            +                                              <span class="n">run_teacher_on_eval</span><span class="o">=</span><span class="n">run_teacher_on_eval</span><span class="p">,</span>
                
 
                            +                                              <span class="n">student</span><span class="o">=</span><span class="n">student</span><span class="p">,</span>
                
 
                            +                                              <span class="n">teacher</span><span class="o">=</span><span class="n">teacher</span><span class="p">)</span>
                
 
                            +        <span class="nb">super</span><span class="p">(</span><span class="n">KDTrainer</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">train</span><span class="p">(</span><span class="n">model</span><span class="o">=</span><span class="n">kd_net</span><span class="p">,</span> <span class="n">training_params</span><span class="o">=</span><span class="n">training_params</span><span class="p">,</span>
                
 
                            +                                     <span class="n">train_loader</span><span class="o">=</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">valid_loader</span><span class="o">=</span><span class="n">valid_loader</span><span class="p">)</span></div></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -390,4 +410,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.legacy.utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.legacy.utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.legacy.utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
<span class="kn">import</span> <span class="nn">torch.nn.init</span> <span class="k">as</span> <span class="nn">init</span>


<div class="viewcode-block" id="prefetch_dataset"><a class="viewcode-back" href="../../../../super_gradients.training.legacy.html#super_gradients.training.legacy.utils.prefetch_dataset">[docs]</a><span class="k">def</span> <span class="nf">prefetch_dataset</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="n">num_workers</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span> <span class="n">batch_size</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">half</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">,</span> <span class="nb">list</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dataset</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="n">dataset</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">dataloader</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">(</span>
            <span class="n">dataset</span><span class="p">,</span>
            <span class="n">batch_size</span><span class="o">=</span><span class="n">batch_size</span><span class="p">,</span>
            <span class="n">shuffle</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">drop_last</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
            <span class="n">num_workers</span><span class="o">=</span><span class="n">num_workers</span><span class="p">,</span> <span class="n">pin_memory</span><span class="o">=</span><span class="kc">False</span>
        <span class="p">)</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">t</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">dataloader</span><span class="p">]</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">)]</span>

    <span class="k">if</span> <span class="n">device</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">to</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="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">tensors</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">half</span><span class="p">:</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">half</span><span class="p">()</span> <span class="k">if</span> <span class="n">t</span><span class="o">.</span><span class="n">is_floating_point</span><span class="p">()</span> <span class="k">else</span> <span class="n">t</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">tensors</span><span class="p">]</span>

    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">TensorDataset</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">)</span></div>


<div class="viewcode-block" id="PrefetchDataLoader"><a class="viewcode-back" href="../../../../super_gradients.training.legacy.html#super_gradients.training.legacy.utils.PrefetchDataLoader">[docs]</a><span class="k">class</span> <span class="nc">PrefetchDataLoader</span><span class="p">:</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataloader</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">half</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">loader</span> <span class="o">=</span> <span class="n">dataloader</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">iter</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">=</span> <span class="n">device</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">float16</span> <span class="k">if</span> <span class="n">half</span> <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">float32</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stream</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">Stream</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="fm">__len__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loader</span><span class="p">)</span>

<div class="viewcode-block" id="PrefetchDataLoader.async_prefech"><a class="viewcode-back" href="../../../../super_gradients.training.legacy.html#super_gradients.training.legacy.utils.PrefetchDataLoader.async_prefech">[docs]</a>    <span class="k">def</span> <span class="nf">async_prefech</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span> <span class="o">=</span> <span class="nb">next</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">iter</span><span class="p">)</span>
        <span class="k">except</span> <span class="ne">StopIteration</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">return</span>

        <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">stream</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">stream</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">next_data</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">non_blocking</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
            <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">next_data</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span> <span class="o">=</span> <span class="p">[</span>
                    <span class="n">t</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">non_blocking</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="k">if</span> <span class="n">t</span><span class="o">.</span><span class="n">is_floating_point</span><span class="p">()</span> <span class="k">else</span> <span class="n">t</span><span class="o">.</span><span class="n">to</span><span class="p">(</span>
                        <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">non_blocking</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span>
                <span class="p">]</span></div>

    <span class="k">def</span> <span class="fm">__iter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">iter</span> <span class="o">=</span> <span class="nb">iter</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loader</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">async_prefech</span><span class="p">()</span>
        <span class="k">while</span> <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">current_stream</span><span class="p">()</span><span class="o">.</span><span class="n">wait_stream</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">stream</span><span class="p">)</span>
            <span class="n">data</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">next_data</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">async_prefech</span><span class="p">()</span>
            <span class="k">yield</span> <span class="n">data</span></div>


<div class="viewcode-block" id="init_params"><a class="viewcode-back" href="../../../../super_gradients.training.legacy.html#super_gradients.training.legacy.utils.init_params">[docs]</a><span class="k">def</span> <span class="nf">init_params</span><span class="p">(</span><span class="n">net</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Init layer parameters.&quot;&quot;&quot;</span>
    <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">net</span><span class="o">.</span><span class="n">modules</span><span class="p">():</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span><span class="p">):</span>
            <span class="n">init</span><span class="o">.</span><span class="n">kaiming_normal</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s1">&#39;fan_out&#39;</span><span class="p">)</span>
            <span class="c1"># if m.bias:</span>
            <span class="c1">#    init.constant(m.bias, -5)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">BatchNorm2d</span><span class="p">):</span>
            <span class="n">init</span><span class="o">.</span><span class="n">constant</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">init</span><span class="o">.</span><span class="n">constant</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">bias</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">):</span>
            <span class="n">init</span><span class="o">.</span><span class="n">normal</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="n">std</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">m</span><span class="o">.</span><span class="n">bias</span><span class="p">:</span>
                <span class="n">init</span><span class="o">.</span><span class="n">constant</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">bias</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span></div>


<div class="viewcode-block" id="format_time"><a class="viewcode-back" href="../../../../super_gradients.training.legacy.html#super_gradients.training.legacy.utils.format_time">[docs]</a><span class="k">def</span> <span class="nf">format_time</span><span class="p">(</span><span class="n">seconds</span><span class="p">):</span>
    <span class="n">days</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">seconds</span> <span class="o">/</span> <span class="mi">3600</span> <span class="o">/</span> <span class="mi">24</span><span class="p">)</span>
    <span class="n">seconds</span> <span class="o">=</span> <span class="n">seconds</span> <span class="o">-</span> <span class="n">days</span> <span class="o">*</span> <span class="mi">3600</span> <span class="o">*</span> <span class="mi">24</span>
    <span class="n">hours</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">seconds</span> <span class="o">/</span> <span class="mi">3600</span><span class="p">)</span>
    <span class="n">seconds</span> <span class="o">=</span> <span class="n">seconds</span> <span class="o">-</span> <span class="n">hours</span> <span class="o">*</span> <span class="mi">3600</span>
    <span class="n">minutes</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">seconds</span> <span class="o">/</span> <span class="mi">60</span><span class="p">)</span>
    <span class="n">seconds</span> <span class="o">=</span> <span class="n">seconds</span> <span class="o">-</span> <span class="n">minutes</span> <span class="o">*</span> <span class="mi">60</span>
    <span class="n">secondsf</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">seconds</span><span class="p">)</span>
    <span class="n">seconds</span> <span class="o">=</span> <span class="n">seconds</span> <span class="o">-</span> <span class="n">secondsf</span>
    <span class="n">millis</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">seconds</span> <span class="o">*</span> <span class="mi">1000</span><span class="p">)</span>

    <span class="n">f</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span>
    <span class="n">i</span> <span class="o">=</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="n">days</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">f</span> <span class="o">+=</span> <span class="nb">str</span><span class="p">(</span><span class="n">days</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;D&#39;</span>
        <span class="n">i</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="n">hours</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">i</span> <span class="o">&lt;=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">f</span> <span class="o">+=</span> <span class="nb">str</span><span class="p">(</span><span class="n">hours</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;h&#39;</span>
        <span class="n">i</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="n">minutes</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">i</span> <span class="o">&lt;=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">f</span> <span class="o">+=</span> <span class="nb">str</span><span class="p">(</span><span class="n">minutes</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;m&#39;</span>
        <span class="n">i</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="n">secondsf</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">i</span> <span class="o">&lt;=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">f</span> <span class="o">+=</span> <span class="nb">str</span><span class="p">(</span><span class="n">secondsf</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;s&#39;</span>
        <span class="n">i</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="n">millis</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">i</span> <span class="o">&lt;=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">f</span> <span class="o">+=</span> <span class="nb">str</span><span class="p">(</span><span class="n">millis</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;ms&#39;</span>
        <span class="n">i</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="n">f</span> <span class="o">==</span> <span class="s1">&#39;&#39;</span><span class="p">:</span>
        <span class="n">f</span> <span class="o">=</span> <span class="s1">&#39;0ms&#39;</span>
    <span class="k">return</span> <span class="n">f</span></div>


<div class="viewcode-block" id="is_better"><a class="viewcode-back" href="../../../../super_gradients.training.legacy.html#super_gradients.training.legacy.utils.is_better">[docs]</a><span class="k">def</span> <span class="nf">is_better</span><span class="p">(</span><span class="n">new_metric</span><span class="p">,</span> <span class="n">current_best_metric</span><span class="p">,</span> <span class="n">metric_to_watch</span><span class="o">=</span><span class="s1">&#39;acc&#39;</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Determines which of the two metrics is better, the higher if watching acc or lower when watching loss</span>
<span class="sd">    :param new_metric:                 the new metric</span>
<span class="sd">    :param current_best_metric:        the compared to metric</span>
<span class="sd">    :param metric_to_watch:             acc or loss</span>
<span class="sd">    :return: bool, True if new metric is better than current</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">metric_to_watch</span> <span class="o">==</span> <span class="s1">&#39;acc&#39;</span> <span class="ow">and</span> <span class="n">new_metric</span> <span class="o">&gt;</span> <span class="n">current_best_metric</span> <span class="ow">or</span> <span class="p">(</span><span class="n">metric_to_watch</span> <span class="o">==</span> <span class="s1">&#39;loss&#39;</span> <span class="ow">and</span> <span class="n">current_best_metric</span> <span class="o">&gt;</span> <span class="n">new_metric</span><span class="p">)</span></div>


<div class="viewcode-block" id="makedirs_if_not_exists"><a class="viewcode-back" href="../../../../super_gradients.training.legacy.html#super_gradients.training.legacy.utils.makedirs_if_not_exists">[docs]</a><span class="k">def</span> <span class="nf">makedirs_if_not_exists</span><span class="p">(</span><span class="n">dir_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    make new directory in dir_path if it doesn&#39;t exists</span>
<span class="sd">        :param dir_path - full path of directory</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">dir_path</span><span class="p">):</span>
        <span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="n">dir_path</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.bce_dice_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.bce_dice_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -92,7 +94,7 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.losses.dice_loss</span> <span class="kn">import</span> <span class="n">BinaryDiceLoss</span>
                
 
                            -<div class="viewcode-block" id="BCEDiceLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.BCEDiceLoss">[docs]</a><span class="k">class</span> <span class="nc">BCEDiceLoss</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="BCEDiceLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.BCEDiceLoss">[docs]</a><span class="k">class</span> <span class="nc">BCEDiceLoss</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Binary Cross Entropy + Dice Loss</span>
                
@@ -108,7 +110,7 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">bce</span> <span class="o">=</span> <span class="n">BCE</span><span class="p">()</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">dice</span> <span class="o">=</span> <span class="n">BinaryDiceLoss</span><span class="p">(</span><span class="n">apply_sigmoid</span><span class="o">=</span><span class="n">logits</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="BCEDiceLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.BCEDiceLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="nb">input</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <spa
                
 
                            +<div class="viewcode-block" id="BCEDiceLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.BCEDiceLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="nb">input</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        @param input: Network&#39;s raw output shaped (N,1,H,W)</span>
                
@@ -145,4 +147,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.losses.ddrnet_loss &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.losses.ddrnet_loss</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.losses.ddrnet_loss</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.losses.ohem_ce_loss</span> <span class="kn">import</span> <span class="n">OhemCELoss</span>


<div class="viewcode-block" id="DDRNetLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ddrnet_loss.DDRNetLoss">[docs]</a><span class="k">class</span> <span class="nc">DDRNetLoss</span><span class="p">(</span><span class="n">OhemCELoss</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.7</span><span class="p">,</span>
                 <span class="n">ohem_percentage</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                 <span class="n">weights</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="p">[</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.4</span><span class="p">],</span>
                 <span class="n">ignore_label</span><span class="o">=</span><span class="mi">255</span><span class="p">,</span>
                 <span class="n">num_pixels_exclude_ignored</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This loss is an extension of the Ohem (Online Hard Example Mining Cross Entropy) Loss.</span>

<span class="sd">        as define in paper:</span>
<span class="sd">        Accurate Semantic Segmentation of Road Scenes ( https://arxiv.org/pdf/2101.06085.pdf )</span>

<span class="sd">        :param threshold: threshold to th hard example mining algorithm</span>
<span class="sd">        :param ohem_percentage: minimum percentage of total pixels for the hard example mining algorithm</span>
<span class="sd">        (taking only the largest) losses</span>
<span class="sd">        :param weights: weights per each input of the loss. This loss supports a multi output (like in DDRNet with</span>
<span class="sd">        an auxiliary head). the losses of each head can be weighted.</span>
<span class="sd">        :param ignore_label: targets label to be ignored</span>
<span class="sd">        :param num_pixels_exclude_ignored: whether to exclude ignore pixels when calculating the mining percentage.</span>
<span class="sd">        see OhemCELoss doc for more details.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span> <span class="n">mining_percent</span><span class="o">=</span><span class="n">ohem_percentage</span><span class="p">,</span> <span class="n">ignore_lb</span><span class="o">=</span><span class="n">ignore_label</span><span class="p">,</span>
                         <span class="n">num_pixels_exclude_ignored</span><span class="o">=</span><span class="n">num_pixels_exclude_ignored</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">weights</span> <span class="o">=</span> <span class="n">weights</span>

<div class="viewcode-block" id="DDRNetLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ddrnet_loss.DDRNetLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions_list</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span>
                <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">predictions_list</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
            <span class="n">predictions_list</span> <span class="o">=</span> <span class="p">(</span><span class="n">predictions_list</span><span class="p">,)</span>

        <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">predictions_list</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">weights</span><span class="p">),</span> <span class="s2">&quot;num of prediction must be the same as num of loss weights&quot;</span>

        <span class="n">losses</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">unweighted_losses</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">predictions</span><span class="p">,</span> <span class="n">weight</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">predictions_list</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">weights</span><span class="p">):</span>
            <span class="n">unweighted_loss</span> <span class="o">=</span> <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">predictions</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span>
            <span class="n">unweighted_losses</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">unweighted_loss</span><span class="p">)</span>
            <span class="n">losses</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">unweighted_loss</span> <span class="o">*</span> <span class="n">weight</span><span class="p">)</span>
        <span class="n">total_loss</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span><span class="n">losses</span><span class="p">)</span>
        <span class="n">unweighted_losses</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">total_loss</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">total_loss</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">unweighted_losses</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.dice_ce_edge_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.dice_ce_edge_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -95,7 +97,7 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.losses.mask_loss</span> <span class="kn">import</span> <span class="n">MaskAttentionLoss</span>
                
 
                            -<div class="viewcode-block" id="DiceCEEdgeLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.DiceCEEdgeLoss">[docs]</a><span class="k">class</span> <span class="nc">DiceCEEdgeLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="DiceCEEdgeLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.DiceCEEdgeLoss">[docs]</a><span class="k">class</span> <span class="nc">DiceCEEdgeLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                                              <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                                              <span class="n">num_aux_heads</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">2</span><span class="p">,</span>
                
@@ -151,7 +153,22 @@
 
                                     <span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">dice_loss</span> <span class="o">=</span> <span class="n">DiceLoss</span><span class="p">(</span><span class="n">apply_softmax</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="DiceCEEdgeLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.DiceCEEdgeLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><
                
 
                            +    <span class="nd">@property</span>
                
 
                            +    <span class="k">def</span> <span class="nf">component_names</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Component names for logging during training.</span>
                
 
                            +<span class="sd">        These correspond to 2nd item in the tuple returned in self.forward(...).</span>
                
 
                            +<span class="sd">        See super_gradients.Trainer.train() docs for more info.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">names</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;main_loss&quot;</span><span class="p">]</span>
                
 
                            +        <span class="c1"># Append aux losses names</span>
                
 
                            +        <span class="n">names</span> <span class="o">+=</span> <span class="p">[</span><span class="sa">f</span><span class="s2">&quot;aux_loss</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">&quot;</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">num_aux_heads</span><span class="p">)]</span
                
 
                            +        <span class="c1"># Append detail losses names</span>
                
 
                            +        <span class="n">names</span> <span class="o">+=</span> <span class="p">[</span><span class="sa">f</span><span class="s2">&quot;detail_loss</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">&quot;</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">num_detail_heads</span><span class="p">)]
                
 
                            +        <span class="n">names</span> <span class="o">+=</span> <span class="p">[</span><span class="s2">&quot;loss&quot;</span><span class="p">]</span>
                
 
                            +        <span class="k">return</span> <span class="n">names</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="DiceCEEdgeLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.DiceCEEdgeLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span cl
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        :param preds: Model output predictions, must be in the followed format:</span>
                
 
                             <span class="sd">         [Main-feats, Aux-feats[0], ..., Aux-feats[num_auxs-1], Detail-feats[0], ..., Detail-feats[num_details-1]</span>
                
@@ -214,4 +231,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.focal_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.focal_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -91,7 +93,7 @@
 
                             <span class="kn">from</span> <span class="nn">torch.nn.modules.loss</span> <span class="kn">import</span> <span class="n">_Loss</span>
                
 
                            -<div class="viewcode-block" id="FocalLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.FocalLoss">[docs]</a><span class="k">class</span> <span class="nc">FocalLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="FocalLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.FocalLoss">[docs]</a><span class="k">class</span> <span class="nc">FocalLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)&quot;&quot;&quot;</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">loss_fcn</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">BCEWithLogitsLoss</span><span class="p">,</span> <span class="n">gamma</span><span class="o">=</span><span class="mf">1.5</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.25</span><span
                
@@ -102,7 +104,7 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">reduction</span> <span class="o">=</span> <span class="n">loss_fcn</span><span class="o">.</span><span class="n">reduction</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">loss_fcn</span><span class="o">.</span><span class="n">reduction</span> <span class="o">=</span> <span class="s1">&#39;none&#39;</span>  <span class="c1"># required to apply FocalLoss to each element</span>
                
 
                            -<div class="viewcode-block" id="FocalLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.FocalLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pred</span><span class="p">,</span> <span class="n">true</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="FocalLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.FocalLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pred</span><span class="p">,</span> <span class="n">true</span><span class="p">):</span>
                
 
                                     <span class="n">loss</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_fcn</span><span class="p">(</span><span class="n">pred</span><span class="p">,</span> <span class="n">true</span><span class="p">)</span>
                
 
                                     <span class="n">pred_prob</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">(</span><span class="n">pred</span><span class="p">)</span>  <span class="c1"># prob from logits</span>
                
@@ -146,4 +148,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.kd_losses &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.kd_losses &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -100,7 +102,7 @@
 
                                                                             <span class="n">torch</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">teacher_output</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">))</span>
                
 
                            -<div class="viewcode-block" id="KDLogitsLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.KDLogitsLoss">[docs]</a><span class="k">class</span> <span class="nc">KDLogitsLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="KDLogitsLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.KDLogitsLoss">[docs]</a><span class="k">class</span> <span class="nc">KDLogitsLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot; Knowledge distillation loss, wraps the task loss and distillation loss &quot;&quot;&quot;</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">task_loss_fn</span><span class="p">:</span> <span class="n">_Loss</span><span class="p">,</span> <span class="n">distillation_loss_fn</span><span class="p">:</span> <span class="n">_Loss</span> <span class="o">=</span> <span class="n">KDklDivLoss</span><span class="p">(),</span> <span class="n">distillation_loss_coeff</span><span class="p">:
                
 
                                     <span class="sd">&#39;&#39;&#39;</span>
                
@@ -114,7 +116,16 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">distillation_loss_fn</span> <span class="o">=</span> <span class="n">distillation_loss_fn</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">distillation_loss_coeff</span> <span class="o">=</span> <span class="n">distillation_loss_coeff</span>
                
 
                            -<div class="viewcode-block" id="KDLogitsLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.KDLogitsLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">kd_module_output</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
                
 
                            +    <span class="nd">@property</span>
                
 
                            +    <span class="k">def</span> <span class="nf">component_names</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Component names for logging during training.</span>
                
 
                            +<span class="sd">        These correspond to 2nd item in the tuple returned in self.forward(...).</span>
                
 
                            +<span class="sd">        See super_gradients.Trainer.train() docs for more info.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="k">return</span> <span class="p">[</span><span class="s2">&quot;Loss&quot;</span><span class="p">,</span> <span class="s2">&quot;Task Loss&quot;</span><span class="p">,</span> <span class="s2">&quot;Distillation Loss&quot;</span><span class="p">]</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="KDLogitsLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.KDLogitsLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">kd_module_output</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
                
 
                                     <span class="n">task_loss</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">task_loss_fn</span><span class="p">(</span><span class="n">kd_module_output</span><span class="o">.</span><span class="n">student_output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">task_loss</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>  <span class="c1"># SOME LOSS FUNCTIONS RETURNS LOSS AND LOG_ITEMS</span>
                
 
                                         <span class="n">task_loss</span> <span class="o">=</span> <span class="n">task_loss</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                
@@ -151,4 +162,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.label_smoothing_cross_entropy_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.label_smoothing_cross_entropy_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -91,7 +93,7 @@
 
                             <span class="kn">import</span> <span class="nn">torch.nn.functional</span> <span class="k">as</span> <span class="nn">F</span>
                
 
                            -<div class="viewcode-block" id="onehot"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.onehot">[docs]</a><span class="k">def</span> <span class="nf">onehot</span><span class="p">(</span><span class="n">indexes</span><span class="p">,</span> <span class="n">N</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="kc">None</span><span
                
 
                            +<span class="k">def</span> <span class="nf">onehot</span><span class="p">(</span><span class="n">indexes</span><span class="p">,</span> <span class="n">N</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Creates a one-hot representation of indexes with N possible entries</span>
                
 
                             <span class="sd">    if N is not specified, it will suit the maximum index appearing.</span>
                
@@ -105,7 +107,7 @@
 
                                 <span class="n">output</span><span class="o">.</span><span class="n">scatter_</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">indexes</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="mi">1</span><span class="p">)</span>
                
 
                                 <span class="k">if</span> <span class="n">ignore_index</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">ignore_index</span> <span class="o">&gt;=</span> <span class="mi">0</span><span class="p">:</span>
                
 
                                     <span class="n">output</span><span class="o">.</span><span class="n">masked_fill_</span><span class="p">(</span><span class="n">indexes</span><span class="o">.</span><span class="n">eq</span><span class="p">(</span><span class="n">ignore_index</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="mi">0</span><span class="p">)</span>
                
 
                            -    <span class="k">return</span> <span class="n">output</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">output</span>
                
 
                             <span class="k">def</span> <span class="nf">_is_long</span><span class="p">(</span><span class="n">x</span><span class="p">):</span>
                
@@ -114,7 +116,7 @@
 
                                 <span class="k">return</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">LongTensor</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class
                
 
                            -<div class="viewcode-block" id="cross_entropy"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.cross_entropy">[docs]</a><span class="k">def</span> <span class="nf">cross_entropy</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="
                
 
                            +<span class="k">def</span> <span class="nf">cross_entropy</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=-</span><span class="mi">100</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="s1">&#39;mean&#3
                
 
                                               <span class="n">smooth_eps</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">smooth_dist</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">from_logits</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;cross entropy loss, with support for target distributions and label smoothing https://arxiv.org/abs/1512.00567&quot;&quot;&quot;</span>
                
 
                                 <span class="n">smooth_eps</span> <span class="o">=</span> <span class="n">smooth_eps</span> <span class="ow">or</span> <span class="mi">0</span>
                
@@ -166,10 +168,10 @@
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span class="o">/</span> <span class="nb">float</span><span class="p">(</span><span class="n">loss</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">-</span> <span class="n">masked_indices</span><span class="o">.</span><span class="n"
                
 
                            -    <span class="k">return</span> <span class="n">loss</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">loss</span>
                
 
                            -<div class="viewcode-block" id="LabelSmoothingCrossEntropyLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.LabelSmoothingCrossEntropyLoss">[docs]</a><span class="k">class</span> <span class="nc">LabelSmoothingCrossEntropyLoss</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">CrossEntropyLoss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="LabelSmoothingCrossEntropyLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.LabelSmoothingCrossEntropyLoss">[docs]</a><span class="k">class</span> <span class="nc">LabelSmoothingCrossEntropyLoss</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">CrossEntropyLoss</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;CrossEntropyLoss - with ability to recieve distrbution as targets, and optional label smoothing&quot;&quot;&quot;</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">weight</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=-</span><span class="mi">100</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="s1">&#39;mean&#39;</span><span class="p">,</span> <span class="n">smooth
                
@@ -180,7 +182,7 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">smooth_dist</span> <span class="o">=</span> <span class="n">smooth_dist</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">from_logits</span> <span class="o">=</span> <span class="n">from_logits</span>
                
 
                            -<div class="viewcode-block" id="LabelSmoothingCrossEntropyLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.LabelSmoothingCrossEntropyLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">sm
                
 
                            +<div class="viewcode-block" id="LabelSmoothingCrossEntropyLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.LabelSmoothingCrossEntropyLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">smooth_di
                
 
                                     <span class="k">if</span> <span class="n">smooth_dist</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                         <span class="n">smooth_dist</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">smooth_dist</span>
                
 
                                     <span class="n">loss</span> <span class="o">=</span> <span class="n">cross_entropy</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">weight</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span>
                
@@ -218,4 +220,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.losses.ohem_ce_loss &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.losses.ohem_ce_loss</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.losses.ohem_ce_loss</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>
<span class="kn">from</span> <span class="nn">torch.nn.modules.loss</span> <span class="kn">import</span> <span class="n">_Loss</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.loss_exceptions</span> <span class="kn">import</span> <span class="n">IllegalRangeForLossAttributeException</span><span class="p">,</span> <span class="n">RequiredLossComponentReductionException</span>


<div class="viewcode-block" id="OhemLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ohem_ce_loss.OhemLoss">[docs]</a><span class="k">class</span> <span class="nc">OhemLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    OhemLoss - Online Hard Example Mining Cross Entropy Loss</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                 <span class="n">mining_percent</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                 <span class="n">ignore_lb</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span>
                 <span class="n">num_pixels_exclude_ignored</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                 <span class="n">criteria</span><span class="p">:</span> <span class="n">_Loss</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param threshold: Sample below probability threshold, is considered hard.</span>
<span class="sd">        :param num_pixels_exclude_ignored: How to calculate total pixels from which extract mining percent of the</span>
<span class="sd">         samples.</span>
<span class="sd">        :param ignore_lb: label index to be ignored in loss calculation.</span>
<span class="sd">        :param criteria: loss to mine the examples from.</span>

<span class="sd">         i.e for num_pixels=100, ignore_pixels=30, mining_percent=0.1:</span>
<span class="sd">         num_pixels_exclude_ignored=False =&gt; num_mining = 100 * 0.1 = 10</span>
<span class="sd">         num_pixels_exclude_ignored=True  =&gt; num_mining = (100 - 30) * 0.1 = 7</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="k">if</span> <span class="n">mining_percent</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">mining_percent</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">raise</span> <span class="n">IllegalRangeForLossAttributeException</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="s2">&quot;mining percent&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">thresh</span> <span class="o">=</span> <span class="o">-</span><span class="n">torch</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">threshold</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">float</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mining_percent</span> <span class="o">=</span> <span class="n">mining_percent</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ignore_lb</span> <span class="o">=</span> <span class="n">ignore_lb</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">num_pixels_exclude_ignored</span> <span class="o">=</span> <span class="n">num_pixels_exclude_ignored</span>

        <span class="k">if</span> <span class="n">criteria</span><span class="o">.</span><span class="n">reduction</span> <span class="o">!=</span> <span class="s1">&#39;none&#39;</span><span class="p">:</span>
            <span class="k">raise</span> <span class="n">RequiredLossComponentReductionException</span><span class="p">(</span><span class="s2">&quot;criteria&quot;</span><span class="p">,</span> <span class="n">criteria</span><span class="o">.</span><span class="n">reduction</span><span class="p">,</span> <span class="s1">&#39;none&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">criteria</span> <span class="o">=</span> <span class="n">criteria</span>

<div class="viewcode-block" id="OhemLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ohem_ce_loss.OhemLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">logits</span><span class="p">,</span> <span class="n">labels</span><span class="p">):</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">criteria</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">labels</span><span class="p">)</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_pixels_exclude_ignored</span><span class="p">:</span>
            <span class="c1"># remove ignore label elements</span>
            <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span><span class="p">[</span><span class="n">labels</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="o">!=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ignore_lb</span><span class="p">]</span>
            <span class="c1"># num pixels in a batch -&gt; num_pixels = batch_size * width * height - ignore_pixels</span>
            <span class="n">num_pixels</span> <span class="o">=</span> <span class="n">loss</span><span class="o">.</span><span class="n">numel</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">num_pixels</span> <span class="o">=</span> <span class="n">labels</span><span class="o">.</span><span class="n">numel</span><span class="p">()</span>
        <span class="c1"># if all pixels are ignore labels, return empty loss tensor</span>
        <span class="k">if</span> <span class="n">num_pixels</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="mf">0.</span><span class="p">])</span><span class="o">.</span><span class="n">requires_grad_</span><span class="p">(</span><span class="kc">True</span><span class="p">)</span>

        <span class="n">num_mining</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mining_percent</span> <span class="o">*</span> <span class="n">num_pixels</span><span class="p">)</span>
        <span class="c1"># in case mining_percent=1, prevent out of bound exception</span>
        <span class="n">num_mining</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">num_mining</span><span class="p">,</span> <span class="n">num_pixels</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">thresh</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">thresh</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">logits</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
        <span class="n">loss</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">loss</span><span class="p">,</span> <span class="n">descending</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">loss</span><span class="p">[</span><span class="n">num_mining</span><span class="p">]</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">thresh</span><span class="p">:</span>
            <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span><span class="p">[</span><span class="n">loss</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">thresh</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span><span class="p">[:</span><span class="n">num_mining</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="OhemCELoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ohem_ce_loss.OhemCELoss">[docs]</a><span class="k">class</span> <span class="nc">OhemCELoss</span><span class="p">(</span><span class="n">OhemLoss</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    OhemLoss - Online Hard Example Mining Cross Entropy Loss</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                 <span class="n">mining_percent</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                 <span class="n">ignore_lb</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span>
                 <span class="n">num_pixels_exclude_ignored</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">):</span>
        <span class="n">ignore_lb</span> <span class="o">=</span> <span class="o">-</span><span class="mi">100</span> <span class="k">if</span> <span class="n">ignore_lb</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">ignore_lb</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">ignore_lb</span>
        <span class="n">criteria</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">CrossEntropyLoss</span><span class="p">(</span><span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_lb</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="s1">&#39;none&#39;</span><span class="p">)</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">OhemCELoss</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><span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span>
                                         <span class="n">mining_percent</span><span class="o">=</span><span class="n">mining_percent</span><span class="p">,</span>
                                         <span class="n">ignore_lb</span><span class="o">=</span><span class="n">ignore_lb</span><span class="p">,</span>
                                         <span class="n">num_pixels_exclude_ignored</span><span class="o">=</span><span class="n">num_pixels_exclude_ignored</span><span class="p">,</span>
                                         <span class="n">criteria</span><span class="o">=</span><span class="n">criteria</span><span class="p">)</span></div>


<div class="viewcode-block" id="OhemBCELoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ohem_ce_loss.OhemBCELoss">[docs]</a><span class="k">class</span> <span class="nc">OhemBCELoss</span><span class="p">(</span><span class="n">OhemLoss</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    OhemBCELoss - Online Hard Example Mining Binary Cross Entropy Loss</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                 <span class="n">mining_percent</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                 <span class="n">ignore_lb</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span>
                 <span class="n">num_pixels_exclude_ignored</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> <span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">OhemBCELoss</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><span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span>
                                          <span class="n">mining_percent</span><span class="o">=</span><span class="n">mining_percent</span><span class="p">,</span>
                                          <span class="n">ignore_lb</span><span class="o">=</span><span class="n">ignore_lb</span><span class="p">,</span>
                                          <span class="n">num_pixels_exclude_ignored</span><span class="o">=</span><span class="n">num_pixels_exclude_ignored</span><span class="p">,</span>
                                          <span class="n">criteria</span><span class="o">=</span><span class="n">nn</span><span class="o">.</span><span class="n">BCEWithLogitsLoss</span><span class="p">(</span><span class="n">reduction</span><span class="o">=</span><span class="s1">&#39;none&#39;</span><span class="p">))</span>

<div class="viewcode-block" id="OhemBCELoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ohem_ce_loss.OhemBCELoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">logits</span><span class="p">,</span> <span class="n">labels</span><span class="p">):</span>

        <span class="c1"># REMOVE SINGLE CLASS CHANNEL WHEN DEALING WITH BINARY DATA</span>
        <span class="k">if</span> <span class="n">logits</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">logits</span> <span class="o">=</span> <span class="n">logits</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="k">return</span> <span class="nb">super</span><span class="p">(</span><span class="n">OhemBCELoss</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">labels</span><span class="o">.</span><span class="n">float</span><span class="p">())</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.r_squared_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.r_squared_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -95,9 +97,9 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">convert_to_tensor</span>
                
 
                            -<div class="viewcode-block" id="RSquaredLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.RSquaredLoss">[docs]</a><span class="k">class</span> <span class="nc">RSquaredLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="RSquaredLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.RSquaredLoss">[docs]</a><span class="k">class</span> <span class="nc">RSquaredLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            -<div class="viewcode-block" id="RSquaredLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.RSquaredLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="RSquaredLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.RSquaredLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
                
 
                                     <span class="c1"># FIXME - THIS NEEDS TO BE CHANGED SUCH THAT THIS CLASS INHERETS FROM _Loss (TAKE A LOOK AT YoLoV3DetectionLoss)</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Computes the R-squared for the output and target values</span>
                
 
                             <span class="sd">        :param output: Tensor / Numpy / List</span>
                
@@ -140,4 +142,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.shelfnet_ohem_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.shelfnet_ohem_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -91,7 +93,7 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.losses.ohem_ce_loss</span> <span class="kn">import</span> <span class="n">OhemCELoss</span>
                
 
                            -<div class="viewcode-block" id="ShelfNetOHEMLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ShelfNetOHEMLoss">[docs]</a><span class="k">class</span> <span class="nc">ShelfNetOHEMLoss</span><span class="p">(</span><span class="n">OhemCELoss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="ShelfNetOHEMLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.ShelfNetOHEMLoss">[docs]</a><span class="k">class</span> <span class="nc">ShelfNetOHEMLoss</span><span class="p">(</span><span class="n">OhemCELoss</span><span class="p">):</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.7</span><span class="p">,</span> <span class="n">mining_percent</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-4</span><span class="p">,</span> <span class="n">ignore_lb
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        This loss is an extension of the Ohem (Online Hard Example Mining Cross Entropy) Loss.</span>
                
@@ -104,14 +106,23 @@
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span> <span class="n">mining_percent</span><span class="o">=</span><span class="n">mining_percent</span><span class="p">,</span> <span class="n">ignore_lb</span><span class="o">=</span><span class="n">ignore_lb</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="ShelfNetOHEMLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ShelfNetOHEMLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions_list</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">targets</span><span c
                
 
                            +<div class="viewcode-block" id="ShelfNetOHEMLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.ShelfNetOHEMLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions_list</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">targets</span><span class="p
                
 
                                     <span class="n">losses</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                                     <span class="k">for</span> <span class="n">predictions</span> <span class="ow">in</span> <span class="n">predictions_list</span><span class="p">:</span>
                
 
                                         <span class="n">losses</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">forward</span><span class="p">(</span><span class="n">predictions</span><span class="p">,</span> <span class="n">targets</span><span class="p">))</span>
                
 
                                     <span class="n">total_loss</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span><span class="n">losses</span><span class="p">)</span>
                
 
                                     <span class="n">losses</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">total_loss</span><span class="p">)</span>
                
 
                            -        <span class="k">return</span> <span class="n">total_loss</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">losses</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span></div></div>
                
 
                            +        <span class="k">return</span> <span class="n">total_loss</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">losses</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span></div>
                
 
                            +
                
 
                            +    <span class="nd">@property</span>
                
 
                            +    <span class="k">def</span> <span class="nf">component_names</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Component names for logging during training.</span>
                
 
                            +<span class="sd">        These correspond to 2nd item in the tuple returned in self.forward(...).</span>
                
 
                            +<span class="sd">        See super_gradients.Trainer.train() docs for more info.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="k">return</span> <span class="p">[</span><span class="s2">&quot;Loss1/4&quot;</span><span class="p">,</span> <span class="s2">&quot;Loss1/8&quot;</span><span class="p">,</span> <span class="s2">&quot;Loss1/16&quot;</span><span class="p">,</span> <span class="s2">&quot;Loss&quot;</span><span class="p">]</span></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -141,4 +152,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.shelfnet_semantic_encoding_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.shelfnet_semantic_encoding_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -91,7 +93,7 @@
 
                             <span class="kn">from</span> <span class="nn">torch.autograd</span> <span class="kn">import</span> <span class="n">Variable</span>
                
 
                            -<div class="viewcode-block" id="ShelfNetSemanticEncodingLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ShelfNetSemanticEncodingLoss">[docs]</a><span class="k">class</span> <span class="nc">ShelfNetSemanticEncodingLoss</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">CrossEntropyLoss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="ShelfNetSemanticEncodingLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.ShelfNetSemanticEncodingLoss">[docs]</a><span class="k">class</span> <span class="nc">ShelfNetSemanticEncodingLoss</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">CrossEntropyLoss</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;2D Cross Entropy Loss with Auxilary Loss&quot;&quot;&quot;</span>
                
 
                                 <span class="c1"># FIXME - THIS LOSS SHOULD BE CHANGED TO SUPPORT APEX</span>
                
@@ -104,7 +106,7 @@
 
                                     <span class="c1"># FIXME - TEST CODE LOTEM, CHANGED IN ORDER TO WORK WITH apex.amp</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">bcewithlogitsloss</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">BCELoss</span><span class="p">(</span><span class="n">weight</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="ShelfNetSemanticEncodingLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.ShelfNetSemanticEncodingLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">logits</span><span class="p">,</span> <span class="n">labels</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="ShelfNetSemanticEncodingLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.ShelfNetSemanticEncodingLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">logits</span><span class="p">,</span> <span class="n">labels</span><span class="p">):</span>
                
 
                                     <span class="n">pred1</span><span class="p">,</span> <span class="n">se_pred</span><span class="p">,</span> <span class="n">pred2</span> <span class="o">=</span> <span class="n">logits</span>
                
 
                                     <span class="n">batch</span> <span class="o">=</span> <span class="n">labels</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
                
@@ -122,7 +124,16 @@
 
                                     <span class="n">loss3</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">bcewithlogitsloss</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">(</span><span class="n">se_pred</span><span class="p">),</span> <span class="n">se_target</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">cuda</span><span class="p">())</span
                
 
                                     <span class="n">total_loss</span> <span class="o">=</span> <span class="n">loss1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">aux_weight</span> <span class="o">*</span> <span class="n">loss2</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">se_weight</span> <span class="o">*</span> <span class="n">loss3</span>
                
 
                                     <span class="n">losses</span> <span class="o">=</span> <span class="p">[</span><span class="n">loss1</span><span class="p">,</span> <span class="n">loss2</span><span class="p">,</span> <span class="n">loss3</span><span class="p">,</span> <span class="n">total_loss</span><span class="p">]</span>
                
 
                            -        <span class="k">return</span> <span class="n">total_loss</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">losses</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span></div></div>
                
 
                            +        <span class="k">return</span> <span class="n">total_loss</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">losses</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span></div>
                
 
                            +
                
 
                            +    <span class="nd">@property</span>
                
 
                            +    <span class="k">def</span> <span class="nf">component_names</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Component names for logging during training.</span>
                
 
                            +<span class="sd">        These correspond to 2nd item in the tuple returned in self.forward(...).</span>
                
 
                            +<span class="sd">        See super_gradients.Trainer.train() docs for more info.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="k">return</span> <span class="p">[</span><span class="s2">&quot;loss1&quot;</span><span class="p">,</span> <span class="s2">&quot;loss2&quot;</span><span class="p">,</span> <span class="s2">&quot;loss3&quot;</span><span class="p">,</span> <span class="s2">&quot;total_loss&quot;</span><span class="p">]</span></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -152,4 +163,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.ssd_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.ssd_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -96,7 +98,7 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.ssd_utils</span> <span class="kn">import</span> <span class="n">DefaultBoxes</span>
                
 
                            -<div class="viewcode-block" id="HardMiningCrossEntropyLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.HardMiningCrossEntropyLoss">[docs]</a><span class="k">class</span> <span class="nc">HardMiningCrossEntropyLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            +<span class="k">class</span> <span class="nc">HardMiningCrossEntropyLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    L_cls = [CE of all positives] + [CE of the hardest backgrounds]</span>
                
 
                             <span class="sd">    where the second term is built from [neg_pos_ratio * positive pairs] background cells with the highest CE</span>
                
@@ -113,7 +115,7 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">neg_pos_ratio</span> <span class="o">=</span> <span class="n">neg_pos_ratio</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">ce</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">CrossEntropyLoss</span><span class="p">(</span><span class="n">reduce</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="HardMiningCrossEntropyLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.HardMiningCrossEntropyLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pred_labels</span><span class="p">,</span> <span class="n">target_labels</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pred_labels</span><span class="p">,</span> <span class="n">target_labels</span><span class="p">):</span>
                
 
                                     <span class="n">mask</span> <span class="o">=</span> <span class="n">target_labels</span> <span class="o">&gt;</span> <span class="mi">0</span>  <span class="c1"># not background</span>
                
 
                                     <span class="n">pos_num</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
                
@@ -135,10 +137,10 @@
 
                                     <span class="n">neg_mask</span> <span class="o">=</span> <span class="n">con_rank</span> <span class="o">&lt;</span> <span class="n">neg_num</span>
                
 
                                     <span class="n">closs</span> <span class="o">=</span> <span class="p">(</span><span class="n">con</span> <span class="o">*</span> <span class="p">(</span><span class="n">mask</span><span class="o">.</span><span class="n">float</span><span class="p">()</span> <span class="o">+</span> <span class="n">neg_mask</span><span class="o">.</span><span class="n">float</span><span class="p">()))</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</sp
                
 
                            -        <span class="k">return</span> <span class="n">closs</span></div></div>
                
 
                            +        <span class="k">return</span> <span class="n">closs</span>
                
 
                            -<div class="viewcode-block" id="SSDLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.SSDLoss">[docs]</a><span class="k">class</span> <span class="nc">SSDLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="SSDLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.SSDLoss">[docs]</a><span class="k">class</span> <span class="nc">SSDLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Implements the loss as the sum of the followings:</span>
                
 
                             <span class="sd">        1. Confidence Loss: All labels, with hard negative mining</span>
                
@@ -167,6 +169,15 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">con_loss</span> <span class="o">=</span> <span class="n">HardMiningCrossEntropyLoss</span><span class="p">(</span><span class="n">neg_pos_ratio</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">iou_thresh</span> <span class="o">=</span> <span class="n">iou_thresh</span>
                
 
                            +    <span class="nd">@property</span>
                
 
                            +    <span class="k">def</span> <span class="nf">component_names</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Component names for logging during training.</span>
                
 
                            +<span class="sd">        These correspond to 2nd item in the tuple returned in self.forward(...).</span>
                
 
                            +<span class="sd">        See super_gradients.Trainer.train() docs for more info.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="k">return</span> <span class="p">[</span><span class="s2">&quot;smooth_l1&quot;</span><span class="p">,</span> <span class="s2">&quot;closs&quot;</span><span class="p">,</span> <span class="s2">&quot;Loss&quot;</span><span class="p">]</span>
                
 
                            +
                
 
                                 <span class="k">def</span> <span class="nf">_norm_relative_bbox</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">loc</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        convert bbox locations into relative locations (relative to the dboxes)</span>
                
@@ -176,7 +187,7 @@
 
                                     <span class="n">gwh</span> <span class="o">=</span> <span class="p">(</span><span class="n">loc</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:,</span> <span class="p">:]</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:,</span> <span class="p">:])</span><span class="o">.</span><span class="n">log</span><span class="p">()</span>
                
 
                                     <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">gxy</span><span class="p">,</span> <span class="n">gwh</span><span class="p">),</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span>
                
 
                            -<div class="viewcode-block" id="SSDLoss.match_dboxes"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.SSDLoss.match_dboxes">[docs]</a>    <span class="k">def</span> <span class="nf">match_dboxes</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">targets</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="SSDLoss.match_dboxes"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.SSDLoss.match_dboxes">[docs]</a>    <span class="k">def</span> <span class="nf">match_dboxes</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">targets</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        creates tensors with target boxes and labels for each dboxes, so with the same len as dboxes.</span>
                
@@ -220,7 +231,7 @@
 
                                     <span class="k">return</span> <span class="n">each_cell_target_locations</span><span class="p">,</span> <span class="n">each_cell_target_labels</span></div>
                
 
                            -<div class="viewcode-block" id="SSDLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.SSDLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">targets</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="SSDLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.SSDLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">targets</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Compute the loss</span>
                
 
                             <span class="sd">            :param predictions - predictions tensor coming from the network,</span>
                
@@ -289,4 +300,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.losses.yolo_v3_loss &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.losses.yolo_v3_loss</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.losses.yolo_v3_loss</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>
<span class="kn">from</span> <span class="nn">torch.nn.modules.loss</span> <span class="kn">import</span> <span class="n">_Loss</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">build_detection_targets</span><span class="p">,</span> <span class="n">calculate_bbox_iou_elementwise</span>


<div class="viewcode-block" id="YoLoV3DetectionLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.yolo_v3_loss.YoLoV3DetectionLoss">[docs]</a><span class="k">class</span> <span class="nc">YoLoV3DetectionLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    YoLoV3DetectionLoss - Loss Class for Object Detection</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">cls_pw</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.</span><span class="p">,</span> <span class="n">obj_pw</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.</span><span class="p">,</span> <span class="n">giou</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">3.54</span><span class="p">,</span> <span class="n">obj</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">64.3</span><span class="p">,</span>
                 <span class="bp">cls</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">37.4</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">YoLoV3DetectionLoss</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">model</span> <span class="o">=</span> <span class="n">model</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cls_pw</span> <span class="o">=</span> <span class="n">cls_pw</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">obj_pw</span> <span class="o">=</span> <span class="n">obj_pw</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">giou</span> <span class="o">=</span> <span class="n">giou</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">obj</span> <span class="o">=</span> <span class="n">obj</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cls</span> <span class="o">=</span> <span class="bp">cls</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes_num</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">model</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">num_classes</span>

<div class="viewcode-block" id="YoLoV3DetectionLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.yolo_v3_loss.YoLoV3DetectionLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_output</span><span class="p">,</span> <span class="n">targets</span><span class="p">):</span>

        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model_output</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">model_output</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
            <span class="c1"># in test/eval mode the Yolo v3 model output a tuple where the second item is the raw predictions</span>
            <span class="n">_</span><span class="p">,</span> <span class="n">predictions</span> <span class="o">=</span> <span class="n">model_output</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">predictions</span> <span class="o">=</span> <span class="n">model_output</span>

        <span class="n">detection_targets</span> <span class="o">=</span> <span class="n">build_detection_targets</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">model</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span>

        <span class="n">float_tensor</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">FloatTensor</span> <span class="k">if</span> <span class="n">predictions</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">is_cuda</span> <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span>
        <span class="n">class_loss</span><span class="p">,</span> <span class="n">giou_loss</span><span class="p">,</span> <span class="n">objectness_loss</span> <span class="o">=</span> <span class="n">float_tensor</span><span class="p">([</span><span class="mi">0</span><span class="p">]),</span> <span class="n">float_tensor</span><span class="p">([</span><span class="mi">0</span><span class="p">]),</span> <span class="n">float_tensor</span><span class="p">([</span><span class="mi">0</span><span class="p">])</span>

        <span class="n">target_class</span><span class="p">,</span> <span class="n">target_box</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="n">anchor_vec</span> <span class="o">=</span> <span class="n">detection_targets</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="s1">&#39;mean&#39;</span>  <span class="c1"># Loss reduction (sum or mean)</span>

        <span class="c1"># DEFINE CRITERIA</span>
        <span class="n">BCEcls</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">BCEWithLogitsLoss</span><span class="p">(</span><span class="n">pos_weight</span><span class="o">=</span><span class="n">float_tensor</span><span class="p">([</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_pw</span><span class="p">]),</span> <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">)</span>
        <span class="n">BCEobj</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">BCEWithLogitsLoss</span><span class="p">(</span><span class="n">pos_weight</span><span class="o">=</span><span class="n">float_tensor</span><span class="p">([</span><span class="bp">self</span><span class="o">.</span><span class="n">obj_pw</span><span class="p">]),</span> <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">)</span>

        <span class="c1"># COMPUTE THE LOSSES BASED ON EACH ONE OF THE YOLO LAYERS PREDICTIONS</span>
        <span class="n">grid_points_num</span><span class="p">,</span> <span class="n">targets_num</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">yolo_layer_index</span><span class="p">,</span> <span class="n">yolo_layer_prediction</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">predictions</span><span class="p">):</span>
            <span class="n">image</span><span class="p">,</span> <span class="n">anchor</span><span class="p">,</span> <span class="n">grid_y</span><span class="p">,</span> <span class="n">grid_x</span> <span class="o">=</span> <span class="n">indices</span><span class="p">[</span><span class="n">yolo_layer_index</span><span class="p">]</span>
            <span class="n">target_object</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">yolo_layer_prediction</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">0</span><span class="p">])</span>
            <span class="n">grid_points_num</span> <span class="o">+=</span> <span class="n">target_object</span><span class="o">.</span><span class="n">numel</span><span class="p">()</span>

            <span class="c1"># COMPUTE LOSSES</span>
            <span class="n">nb</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">image</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">nb</span><span class="p">:</span>  <span class="c1"># number of targets</span>
                <span class="n">targets_num</span> <span class="o">+=</span> <span class="n">nb</span>
                <span class="n">predictions_for_targets</span> <span class="o">=</span> <span class="n">yolo_layer_prediction</span><span class="p">[</span><span class="n">image</span><span class="p">,</span> <span class="n">anchor</span><span class="p">,</span> <span class="n">grid_y</span><span class="p">,</span> <span class="n">grid_x</span><span class="p">]</span>
                <span class="n">target_object</span><span class="p">[</span><span class="n">image</span><span class="p">,</span> <span class="n">anchor</span><span class="p">,</span> <span class="n">grid_y</span><span class="p">,</span> <span class="n">grid_x</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span>

                <span class="c1"># GIoU LOSS CALCULATION</span>
                <span class="n">pxy</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">(</span>
                    <span class="n">predictions_for_targets</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">:</span><span class="mi">2</span><span class="p">])</span>  <span class="c1"># pxy = pxy * s - (s - 1) / 2,  s = 1.5  (scale_xy)</span>
                <span class="n">bbox_prediction</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
                    <span class="p">(</span><span class="n">pxy</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">predictions_for_targets</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:</span><span class="mi">4</span><span class="p">])</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="nb">max</span><span class="o">=</span><span class="mf">1E3</span><span class="p">)</span> <span class="o">*</span> <span class="n">anchor_vec</span><span class="p">[</span><span class="n">yolo_layer_index</span><span class="p">]),</span> <span class="mi">1</span><span class="p">)</span>
                <span class="n">giou</span> <span class="o">=</span> <span class="mf">1.0</span> <span class="o">-</span> <span class="n">calculate_bbox_iou_elementwise</span><span class="p">(</span><span class="n">bbox_prediction</span><span class="o">.</span><span class="n">t</span><span class="p">(),</span> <span class="n">target_box</span><span class="p">[</span><span class="n">yolo_layer_index</span><span class="p">],</span>
                                                            <span class="n">x1y1x2y2</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">GIoU</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                <span class="n">giou_loss</span> <span class="o">+=</span> <span class="n">giou</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s1">&#39;sum&#39;</span> <span class="k">else</span> <span class="n">giou</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>

                <span class="c1"># ONLY RELEVANT TO MULTIPLE CLASSES</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes_num</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
                    <span class="n">class_targets</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">predictions_for_targets</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">:])</span>
                    <span class="n">class_targets</span><span class="p">[</span><span class="nb">range</span><span class="p">(</span><span class="n">nb</span><span class="p">),</span> <span class="n">target_class</span><span class="p">[</span><span class="n">yolo_layer_index</span><span class="p">]]</span> <span class="o">=</span> <span class="mf">1.0</span>
                    <span class="n">class_loss</span> <span class="o">+=</span> <span class="n">BCEcls</span><span class="p">(</span><span class="n">predictions_for_targets</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">:],</span> <span class="n">class_targets</span><span class="p">)</span>

            <span class="n">objectness_loss</span> <span class="o">+=</span> <span class="n">BCEobj</span><span class="p">(</span><span class="n">yolo_layer_prediction</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span> <span class="n">target_object</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s1">&#39;sum&#39;</span><span class="p">:</span>
            <span class="n">giou_loss</span> <span class="o">*=</span> <span class="mi">3</span> <span class="o">/</span> <span class="n">targets_num</span>
            <span class="n">objectness_loss</span> <span class="o">*=</span> <span class="mi">3</span> <span class="o">/</span> <span class="n">grid_points_num</span>
            <span class="n">class_loss</span> <span class="o">*=</span> <span class="mi">3</span> <span class="o">/</span> <span class="n">targets_num</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes_num</span>

        <span class="n">loss</span> <span class="o">=</span> <span class="n">giou_loss</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">giou</span> <span class="o">+</span> <span class="n">objectness_loss</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">obj</span> <span class="o">+</span> <span class="n">class_loss</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls</span>
        <span class="k">return</span> <span class="n">loss</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">giou_loss</span><span class="p">,</span> <span class="n">objectness_loss</span><span class="p">,</span> <span class="n">class_loss</span><span class="p">,</span> <span class="n">loss</span><span class="p">))</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.losses.yolo_v5_loss &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.losses.yolo_v5_loss</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.losses.yolo_v5_loss</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>
<span class="kn">from</span> <span class="nn">torch.nn.modules.loss</span> <span class="kn">import</span> <span class="n">_Loss</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.losses.focal_loss</span> <span class="kn">import</span> <span class="n">FocalLoss</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">calculate_bbox_iou_elementwise</span><span class="p">,</span> <span class="n">Anchors</span>


<div class="viewcode-block" id="YoLoV5DetectionLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoLoV5DetectionLoss">[docs]</a><span class="k">class</span> <span class="nc">YoLoV5DetectionLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculate YOLO V5 loss:</span>
<span class="sd">    L = L_objectivness + L_boxes + L_classification</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">anchors</span><span class="p">:</span> <span class="n">Anchors</span><span class="p">,</span>
                 <span class="n">cls_pos_weight</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">obj_pos_weight</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
                 <span class="n">obj_loss_gain</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">box_loss_gain</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.05</span><span class="p">,</span> <span class="n">cls_loss_gain</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                 <span class="n">focal_loss_gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
                 <span class="n">cls_objectness_weights</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">anchor_threshold</span><span class="o">=</span><span class="mf">4.0</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param anchors:                 the anchors of the model (same anchors used for training)</span>
<span class="sd">        :param cls_pos_weight:          pos_weight for BCE in L_classification,</span>
<span class="sd">                                        can be one value for all positives or a list of weights for each class</span>
<span class="sd">        :param obj_pos_weight:          pos_weight for BCE in L_objectivness</span>
<span class="sd">        :param obj_loss_gain:           coef for L_objectivness</span>
<span class="sd">        :param box_loss_gain:           coef for L_boxes</span>
<span class="sd">        :param cls_loss_gain:           coef for L_classification</span>
<span class="sd">        :param focal_loss_gamma:        gamma for a focal loss, 0 to train with a usual BCE</span>
<span class="sd">        :param cls_objectness_weights:  class-based weight for L_objectivness that will be applied in each cell that</span>
<span class="sd">                                        has a GT assigned to it.</span>
<span class="sd">                                        Note: default weight for objectness loss in each cell is 1.</span>
<span class="sd">        :param anchor_threshold:                ratio defining a size range of an appropriate anchor.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">YoLoV5DetectionLoss</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>

        <span class="bp">self</span><span class="o">.</span><span class="n">cls_pos_weight</span> <span class="o">=</span> <span class="n">cls_pos_weight</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">obj_pos_weight</span> <span class="o">=</span> <span class="n">obj_pos_weight</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">obj_loss_gain</span> <span class="o">=</span> <span class="n">obj_loss_gain</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">box_loss_gain</span> <span class="o">=</span> <span class="n">box_loss_gain</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cls_loss_gain</span> <span class="o">=</span> <span class="n">cls_loss_gain</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">focal_loss_gamma</span> <span class="o">=</span> <span class="n">focal_loss_gamma</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">anchor_threshold</span> <span class="o">=</span> <span class="n">anchor_threshold</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">anchors</span> <span class="o">=</span> <span class="n">anchors</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">cls_obj_weights</span> <span class="o">=</span> <span class="n">cls_objectness_weights</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">cls_objectness_weights</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">cls_obj_weights</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">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">cls_objectness_weights</span><span class="p">))</span>

<div class="viewcode-block" id="YoLoV5DetectionLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoLoV5DetectionLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_output</span><span class="p">,</span> <span class="n">targets</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model_output</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">model_output</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
            <span class="c1"># in test/eval mode the Yolo v5 model output a tuple where the second item is the raw predictions</span>
            <span class="n">_</span><span class="p">,</span> <span class="n">predictions</span> <span class="o">=</span> <span class="n">model_output</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">predictions</span> <span class="o">=</span> <span class="n">model_output</span>

        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">compute_loss</span><span class="p">(</span><span class="n">predictions</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span></div>

<div class="viewcode-block" id="YoLoV5DetectionLoss.build_targets"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoLoV5DetectionLoss.build_targets">[docs]</a>    <span class="k">def</span> <span class="nf">build_targets</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> \
            <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">List</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]],</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]]:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Assign targets to anchors to use in L_boxes &amp; L_classification calculation:</span>
<span class="sd">            * each target can be assigned to a few anchors,</span>
<span class="sd">            all anchors that are within [1/self.anchor_threshold, self.anchor_threshold] times target size range</span>
<span class="sd">            * each anchor can be assigned to a few targets</span>

<span class="sd">        :param predictions:         Yolo predictions</span>
<span class="sd">        :param targets:             ground truth targets</span>
<span class="sd">        :return:                    each of 4 outputs contains one element for each Yolo output,</span>
<span class="sd">                                    correspondences are raveled over the whole batch and all anchors:</span>
<span class="sd">                                        * classes of the targets;</span>
<span class="sd">                                        * boxes of the targets;</span>
<span class="sd">                                        * image id in a batch, anchor id, grid y, grid x coordinates;</span>
<span class="sd">                                        * anchor sizes.</span>
<span class="sd">                                    All the above can be indexed in parallel to get the selected correspondences</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">num_anchors</span><span class="p">,</span> <span class="n">num_targets</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">anchors</span><span class="o">.</span><span class="n">num_anchors</span><span class="p">,</span> <span class="n">targets</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">target_classes</span><span class="p">,</span> <span class="n">target_boxes</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="n">anchors</span> <span class="o">=</span> <span class="p">[],</span> <span class="p">[],</span> <span class="p">[],</span> <span class="p">[]</span>

        <span class="n">gain</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">7</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">targets</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>  <span class="c1"># normalized to gridspace gain</span>
        <span class="n">anchor_indices</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">num_anchors</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">targets</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
        <span class="n">anchor_indices</span> <span class="o">=</span> <span class="n">anchor_indices</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">num_anchors</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_targets</span><span class="p">)</span>
        <span class="c1"># repeat all targets for each anchor and append a corresponding anchor index</span>
        <span class="n">targets</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">targets</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="n">num_anchors</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">anchor_indices</span><span class="p">[:,</span> <span class="p">:,</span> <span class="kc">None</span><span class="p">]),</span> <span class="mi">2</span><span class="p">)</span>

        <span class="n">bias</span> <span class="o">=</span> <span class="mf">0.5</span>
        <span class="n">off</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="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
                            <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">],</span>  <span class="c1"># j,k,l,m</span>
                            <span class="p">],</span> <span class="n">device</span><span class="o">=</span><span class="n">targets</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span> <span class="o">*</span> <span class="n">bias</span>  <span class="c1"># offsets</span>

        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">anchors</span><span class="o">.</span><span class="n">detection_layers_num</span><span class="p">):</span>
            <span class="n">anch</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">anchors</span><span class="o">.</span><span class="n">anchors</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
            <span class="n">gain</span><span class="p">[</span><span class="mi">2</span><span class="p">:</span><span class="mi">6</span><span class="p">]</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">predictions</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">)[[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span>  <span class="c1"># xyxy gain</span>

            <span class="c1"># Convert target coordinates from [0, 1] range to coordinates in [0, GridY], [0, GridX] ranges</span>
            <span class="n">t</span> <span class="o">=</span> <span class="n">targets</span> <span class="o">*</span> <span class="n">gain</span>
            <span class="k">if</span> <span class="n">num_targets</span><span class="p">:</span>
                <span class="c1"># Match: filter targets by anchor size ratio</span>
                <span class="n">r</span> <span class="o">=</span> <span class="n">t</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">4</span><span class="p">:</span><span class="mi">6</span><span class="p">]</span> <span class="o">/</span> <span class="n">anch</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span>  <span class="c1"># wh ratio</span>
                <span class="n">filtered_targets_ids</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">r</span><span class="p">,</span> <span class="mf">1.</span> <span class="o">/</span> <span class="n">r</span><span class="p">)</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="mi">2</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">anchor_threshold</span>  <span class="c1"># compare</span>
                <span class="n">t</span> <span class="o">=</span> <span class="n">t</span><span class="p">[</span><span class="n">filtered_targets_ids</span><span class="p">]</span>

                <span class="c1"># Find coordinates of targets on a grid</span>
                <span class="n">gxy</span> <span class="o">=</span> <span class="n">t</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:</span><span class="mi">4</span><span class="p">]</span>  <span class="c1"># grid xy</span>
                <span class="n">gxi</span> <span class="o">=</span> <span class="n">gain</span><span class="p">[[</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">]]</span> <span class="o">-</span> <span class="n">gxy</span>  <span class="c1"># inverse</span>
                <span class="n">j</span><span class="p">,</span> <span class="n">k</span> <span class="o">=</span> <span class="p">((</span><span class="n">gxy</span> <span class="o">%</span> <span class="mf">1.</span> <span class="o">&lt;</span> <span class="n">bias</span><span class="p">)</span> <span class="o">&amp;</span> <span class="p">(</span><span class="n">gxy</span> <span class="o">&gt;</span> <span class="mf">1.</span><span class="p">))</span><span class="o">.</span><span class="n">T</span>
                <span class="n">l</span><span class="p">,</span> <span class="n">m</span> <span class="o">=</span> <span class="p">((</span><span class="n">gxi</span> <span class="o">%</span> <span class="mf">1.</span> <span class="o">&lt;</span> <span class="n">bias</span><span class="p">)</span> <span class="o">&amp;</span> <span class="p">(</span><span class="n">gxi</span> <span class="o">&gt;</span> <span class="mf">1.</span><span class="p">))</span><span class="o">.</span><span class="n">T</span>
                <span class="n">j</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">((</span><span class="n">torch</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">j</span><span class="p">),</span> <span class="n">j</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">l</span><span class="p">,</span> <span class="n">m</span><span class="p">))</span>
                <span class="n">t</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">repeat</span><span class="p">((</span><span class="mi">5</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">))[</span><span class="n">j</span><span class="p">]</span>
                <span class="n">offsets</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">gxy</span><span class="p">)[</span><span class="kc">None</span><span class="p">]</span> <span class="o">+</span> <span class="n">off</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">])[</span><span class="n">j</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">t</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">offsets</span> <span class="o">=</span> <span class="mi">0</span>

            <span class="c1"># Define</span>
            <span class="n">b</span><span class="p">,</span> <span class="n">c</span> <span class="o">=</span> <span class="n">t</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">long</span><span class="p">()</span><span class="o">.</span><span class="n">T</span>  <span class="c1"># image, class</span>
            <span class="n">gxy</span> <span class="o">=</span> <span class="n">t</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:</span><span class="mi">4</span><span class="p">]</span>  <span class="c1"># grid xy</span>
            <span class="n">gwh</span> <span class="o">=</span> <span class="n">t</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">:</span><span class="mi">6</span><span class="p">]</span>  <span class="c1"># grid wh</span>
            <span class="n">gij</span> <span class="o">=</span> <span class="p">(</span><span class="n">gxy</span> <span class="o">-</span> <span class="n">offsets</span><span class="p">)</span><span class="o">.</span><span class="n">long</span><span class="p">()</span>
            <span class="n">gi</span><span class="p">,</span> <span class="n">gj</span> <span class="o">=</span> <span class="n">gij</span><span class="o">.</span><span class="n">T</span>  <span class="c1"># grid xy indices</span>
            <span class="c1"># prevent coordinates from going out of bounds</span>
            <span class="n">gi</span><span class="p">,</span> <span class="n">gj</span> <span class="o">=</span> <span class="n">gi</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">gain</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="mi">1</span><span class="p">),</span> <span class="n">gj</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">gain</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>

            <span class="c1"># Append</span>
            <span class="n">a</span> <span class="o">=</span> <span class="n">t</span><span class="p">[:,</span> <span class="mi">6</span><span class="p">]</span><span class="o">.</span><span class="n">long</span><span class="p">()</span>  <span class="c1"># anchor indices</span>
            <span class="n">indices</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">b</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">gj</span><span class="p">,</span> <span class="n">gi</span><span class="p">))</span>  <span class="c1"># image, anchor, grid indices</span>
            <span class="n">target_boxes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">gxy</span> <span class="o">-</span> <span class="n">gij</span><span class="p">,</span> <span class="n">gwh</span><span class="p">),</span> <span class="mi">1</span><span class="p">))</span>  <span class="c1"># box</span>
            <span class="n">anchors</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">anch</span><span class="p">[</span><span class="n">a</span><span class="p">])</span>  <span class="c1"># anchors</span>
            <span class="n">target_classes</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">c</span><span class="p">)</span>  <span class="c1"># class</span>

        <span class="k">return</span> <span class="n">target_classes</span><span class="p">,</span> <span class="n">target_boxes</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="n">anchors</span></div>

<div class="viewcode-block" id="YoLoV5DetectionLoss.compute_loss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoLoV5DetectionLoss.compute_loss">[docs]</a>    <span class="k">def</span> <span class="nf">compute_loss</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">giou_loss_ratio</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">)</span> \
            <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        L = L_objectivness + L_boxes + L_classification</span>
<span class="sd">        where:</span>
<span class="sd">            * L_boxes and L_classification are calculated only between anchors and targets that suit them;</span>
<span class="sd">            * L_objectivness is calculated on all anchors.</span>

<span class="sd">        L_classification:</span>
<span class="sd">            for anchors that have suitable ground truths in their grid locations add BCEs</span>
<span class="sd">            to force max probability for each GT class in a multi-label way</span>
<span class="sd">            Coef: self.cls_loss_gain</span>
<span class="sd">        L_boxes:</span>
<span class="sd">            for anchors that have suitable ground truths in their grid locations</span>
<span class="sd">            add (1 - IoU), IoU between a predicted box and each GT box, force maximum IoU</span>
<span class="sd">            Coef: self.box_loss_gain</span>
<span class="sd">        L_objectness:</span>
<span class="sd">            for each anchor add BCE to force a prediction of (1 - giou_loss_ratio) + giou_loss_ratio * IoU,</span>
<span class="sd">            IoU between a predicted box and random GT in it</span>
<span class="sd">            Coef: self.obj_loss_gain, loss from each YOLO grid is additionally multiplied by balance = [4.0, 1.0, 0.4]</span>
<span class="sd">                  to balance different contributions coming from different numbers of grid cells</span>

<span class="sd">        :param predictions:     output from all Yolo levels, each of shape</span>
<span class="sd">                                [Batch x Num_Anchors x GridSizeY x GridSizeX x (4 + 1 + Num_classes)]</span>
<span class="sd">        :param targets:         [Num_targets x (4 + 2)], values on dim 1 are: image id in a batch, class, box x y w h</span>
<span class="sd">        :param giou_loss_ratio: a coef in L_objectness defining what should be predicted as objecness</span>
<span class="sd">                                in a call with a target: can be a value in [IoU, 1] range</span>
<span class="sd">        :return:                loss, all losses separately in a detached tensor</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">device</span> <span class="o">=</span> <span class="n">targets</span><span class="o">.</span><span class="n">device</span>
        <span class="n">loss_classification</span><span class="p">,</span> <span class="n">loss_boxes</span><span class="p">,</span> <span class="n">loss_objectivness</span> <span class="o">=</span> \
            <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</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">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</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">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</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">target_classes</span><span class="p">,</span> <span class="n">target_boxes</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="n">anchors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">build_targets</span><span class="p">(</span><span class="n">predictions</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span>  <span class="c1"># targets</span>

        <span class="c1"># Define criteria</span>
        <span class="n">BCEcls</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">BCEWithLogitsLoss</span><span class="p">(</span><span class="n">pos_weight</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="bp">self</span><span class="o">.</span><span class="n">cls_pos_weight</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>
        <span class="n">BCEobj</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">BCEWithLogitsLoss</span><span class="p">(</span><span class="n">pos_weight</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="bp">self</span><span class="o">.</span><span class="n">obj_pos_weight</span><span class="p">]),</span> <span class="n">reduction</span><span class="o">=</span><span class="s1">&#39;none&#39;</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>

        <span class="c1"># Focal loss</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">focal_loss_gamma</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">BCEcls</span><span class="p">,</span> <span class="n">BCEobj</span> <span class="o">=</span> <span class="n">FocalLoss</span><span class="p">(</span><span class="n">BCEcls</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">focal_loss_gamma</span><span class="p">),</span> <span class="n">FocalLoss</span><span class="p">(</span><span class="n">BCEobj</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">focal_loss_gamma</span><span class="p">)</span>

        <span class="c1"># Losses</span>
        <span class="n">num_targets</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="n">num_predictions</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">predictions</span><span class="p">)</span>
        <span class="n">balance</span> <span class="o">=</span> <span class="p">[</span><span class="mf">4.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.4</span><span class="p">]</span> <span class="k">if</span> <span class="n">num_predictions</span> <span class="o">==</span> <span class="mi">3</span> <span class="k">else</span> <span class="p">[</span><span class="mf">4.0</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.4</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">]</span>  <span class="c1"># P3-5 or P3-6</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">prediction</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">predictions</span><span class="p">):</span>  <span class="c1"># layer index, layer predictions</span>
            <span class="n">image</span><span class="p">,</span> <span class="n">anchor</span><span class="p">,</span> <span class="n">grid_y</span><span class="p">,</span> <span class="n">grid_x</span> <span class="o">=</span> <span class="n">indices</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
            <span class="n">target_obj</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">prediction</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">0</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">weight_obj</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">prediction</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">0</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">n</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>  <span class="c1"># number of targets</span>
            <span class="k">if</span> <span class="n">n</span><span class="p">:</span>
                <span class="n">num_targets</span> <span class="o">+=</span> <span class="n">n</span>  <span class="c1"># cumulative targets</span>
                <span class="n">ps</span> <span class="o">=</span> <span class="n">prediction</span><span class="p">[</span><span class="n">image</span><span class="p">,</span> <span class="n">anchor</span><span class="p">,</span> <span class="n">grid_y</span><span class="p">,</span> <span class="n">grid_x</span><span class="p">]</span>  <span class="c1"># prediction subset corresponding to targets</span>

                <span class="c1"># Boxes loss</span>
                <span class="n">pxy</span> <span class="o">=</span> <span class="n">ps</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">()</span> <span class="o">*</span> <span class="mf">2.</span> <span class="o">-</span> <span class="mf">0.5</span>
                <span class="n">pwh</span> <span class="o">=</span> <span class="p">(</span><span class="n">ps</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:</span><span class="mi">4</span><span class="p">]</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">()</span> <span class="o">*</span> <span class="mi">2</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">anchors</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
                <span class="n">pbox</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">pxy</span><span class="p">,</span> <span class="n">pwh</span><span class="p">),</span> <span class="mi">1</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>  <span class="c1"># predicted box</span>
                <span class="n">iou</span> <span class="o">=</span> <span class="n">calculate_bbox_iou_elementwise</span><span class="p">(</span><span class="n">pbox</span><span class="o">.</span><span class="n">T</span><span class="p">,</span> <span class="n">target_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="n">x1y1x2y2</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">CIoU</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                <span class="n">loss_boxes</span> <span class="o">+=</span> <span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">iou</span><span class="p">)</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>  <span class="c1"># iou loss</span>

                <span class="c1"># Objectness loss target</span>
                <span class="n">target_obj</span><span class="p">[</span><span class="n">image</span><span class="p">,</span> <span class="n">anchor</span><span class="p">,</span> <span class="n">grid_y</span><span class="p">,</span> <span class="n">grid_x</span><span class="p">]</span> <span class="o">=</span> \
                    <span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">giou_loss_ratio</span><span class="p">)</span> <span class="o">+</span> <span class="n">giou_loss_ratio</span> <span class="o">*</span> <span class="n">iou</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">target_obj</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
                <span class="c1"># Weights for weighted objectness</span>
                <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_obj_weights</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="c1"># NOTE: for grid cells that have a few ground truths with different classes assigned to them</span>
                    <span class="c1"># objectness weight will be picked randomly from one of these classes</span>
                    <span class="n">weight_obj</span><span class="p">[</span><span class="n">image</span><span class="p">,</span> <span class="n">anchor</span><span class="p">,</span> <span class="n">grid_y</span><span class="p">,</span> <span class="n">grid_x</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_obj_weights</span><span class="p">[</span><span class="n">target_classes</span><span class="p">[</span><span class="n">i</span><span class="p">]]</span>

                <span class="c1"># Classification loss</span>
                <span class="k">if</span> <span class="n">ps</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">6</span><span class="p">:</span>   <span class="c1"># cls loss (only if multiple classes)</span>
                    <span class="n">t</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">full_like</span><span class="p">(</span><span class="n">ps</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">:],</span> <span class="mi">0</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="c1"># targets</span>
                    <span class="n">t</span><span class="p">[</span><span class="nb">range</span><span class="p">(</span><span class="n">n</span><span class="p">),</span> <span class="n">target_classes</span><span class="p">[</span><span class="n">i</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">1</span>
                    <span class="n">loss_classification</span> <span class="o">+=</span> <span class="n">BCEcls</span><span class="p">(</span><span class="n">ps</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">:],</span> <span class="n">t</span><span class="p">)</span>  <span class="c1"># BCE</span>

            <span class="c1"># Objectness loss</span>
            <span class="n">loss_obj_cur_head</span> <span class="o">=</span> <span class="n">BCEobj</span><span class="p">(</span><span class="n">prediction</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span> <span class="n">target_obj</span><span class="p">)</span>
            <span class="n">loss_obj_cur_head</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">loss_obj_cur_head</span> <span class="o">*</span> <span class="n">weight_obj</span> <span class="o">/</span> <span class="n">torch</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">weight_obj</span><span class="p">))</span>
            <span class="n">loss_objectivness</span> <span class="o">+=</span> <span class="n">loss_obj_cur_head</span> <span class="o">*</span> <span class="n">balance</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>  <span class="c1"># obj loss</span>

        <span class="n">batch_size</span> <span class="o">=</span> <span class="n">prediction</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>  <span class="c1"># batch size</span>

        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_boxes</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">box_loss_gain</span> <span class="o">+</span> <span class="n">loss_objectivness</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">obj_loss_gain</span> <span class="o">+</span> <span class="n">loss_classification</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_loss_gain</span>
        <span class="c1"># IMPORTANT: box, obj and cls loss are logged scaled by gain in ultralytics</span>
        <span class="c1"># and are logged unscaled in our codebase</span>
        <span class="k">return</span> <span class="n">loss</span> <span class="o">*</span> <span class="n">batch_size</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">loss_boxes</span><span class="p">,</span> <span class="n">loss_objectivness</span><span class="p">,</span> <span class="n">loss_classification</span><span class="p">,</span> <span class="n">loss</span><span class="p">))</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.losses.yolox_loss &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.losses.yolox_loss &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -95,10 +97,13 @@
 
                             <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span>
                
 
                             <span class="kn">import</span> <span class="nn">torch</span>
                
 
                            +<span class="kn">import</span> <span class="nn">torch.distributed</span> <span class="k">as</span> <span class="nn">dist</span>
                
 
                             <span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>
                
 
                             <span class="kn">from</span> <span class="nn">torch.nn.modules.loss</span> <span class="kn">import</span> <span class="n">_Loss</span>
                
 
                             <span class="kn">import</span> <span class="nn">torch.nn.functional</span> <span class="k">as</span> <span class="nn">F</span>
                
 
                            +
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">torch_version_is_greater_or_equal</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">calculate_bbox_iou_matrix</span>
                
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
@@ -127,10 +132,9 @@
 
                                     <span class="n">supported_losses</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;iou&quot;</span><span class="p">,</span> <span class="s2">&quot;giou&quot;</span><span class="p">]</span>
                
 
                                     <span class="n">supported_reductions</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;mean&quot;</span><span class="p">,</span> <span class="s2">&quot;sum&quot;</span><span class="p">,</span> <span class="s2">&quot;none&quot;</span><span class="p">]</span>
                
 
                                     <span class="k">if</span> <span class="n">loss_type</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">supported_losses</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Illegal loss_type value: &quot;</span> <span class="o">+</span> <span class="n">loss_type</span> <span class="o">+</span> <span class="s1">&#39;, expected one of: &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">supported_losses</span><span class="p">))</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Illegal loss_type value: &quot;</span> <span class="o">+</span> <span class="n">loss_type</span> <span class="o">+</span> <span class="s2">&quot;, expected one of: &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">supported_losses</span><span class="p">))</span>
                
 
                                     <span class="k">if</span> <span class="n">reduction</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">supported_reductions</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                
 
                            -                <span class="s2">&quot;Illegal reduction value: &quot;</span> <span class="o">+</span> <span class="n">reduction</span> <span class="o">+</span> <span class="s1">&#39;, expected one of: &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">supported_reductions</span><span class="p">))</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Illegal reduction value: &quot;</span> <span class="o">+</span> <span class="n">reduction</span> <span class="o">+</span> <span class="s2">&quot;, expected one of: &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">supported_reductions</span><span class="p">))</span>
                
 
                                 <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pred</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
                
 
                                     <span class="k">assert</span> <span class="n">pred</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">target</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                
@@ -149,7 +153,7 @@
 
                                     <span class="n">iou</span> <span class="o">=</span> <span class="p">(</span><span class="n">area_i</span><span class="p">)</span> <span class="o">/</span> <span class="p">(</span><span class="n">area_u</span> <span class="o">+</span> <span class="mf">1e-16</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_type</span> <span class="o">==</span> <span class="s2">&quot;iou&quot;</span><span class="p">:</span>
                
 
                            -            <span class="n">loss</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">iou</span> <span class="o">**</span> <span class="mi">2</span>
                
 
                            +            <span class="n">loss</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">iou</span><span class="o">**</span><span class="mi">2</span>
                
 
                                     <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_type</span> <span class="o">==</span> <span class="s2">&quot;giou&quot;</span><span class="p">:</span>
                
 
                                         <span class="n">c_tl</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">((</span><span class="n">pred</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">pred</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">),</spa
                
 
                                         <span class="n">c_br</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">((</span><span class="n">pred</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">pred</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">),</spa
                
@@ -165,7 +169,7 @@
 
                                     <span class="k">return</span> <span class="n">loss</span>
                
 
                            -<div class="viewcode-block" id="YoloXDetectionLoss"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoloXDetectionLoss">[docs]</a><span class="k">class</span> <span class="nc">YoloXDetectionLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="YoloXDetectionLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXDetectionLoss">[docs]</a><span class="k">class</span> <span class="nc">YoloXDetectionLoss</span><span class="p">(</span><span class="n">_Loss</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Calculate YOLOX loss:</span>
                
 
                             <span class="sd">    L = L_objectivness + L_iou + L_classification + 1[use_l1]*L_l1</span>
                
@@ -200,8 +204,7 @@
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">strides</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">use_l1</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><sp
                
 
                            -                 <span class="n">iou_type</span><span class="o">=</span><span class="s1">&#39;iou&#39;</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">strides</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">use_l1</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><sp
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">grids</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">)]</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">strides</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">strides</span> <span class="o">=</span> <span class="n">strides</span>
                
@@ -213,7 +216,16 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">bcewithlog_loss</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">BCEWithLogitsLoss</span><span class="p">(</span><span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">iou_loss</span> <span class="o">=</span> <span class="n">IOUloss</span><span class="p">(</span><span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">,</span> <span class="n">loss_type</span><span class="o">=</span><span class="n">iou_type</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="YoloXDetectionLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoloXDetectionLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_output</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span clas
                
 
                            +    <span class="nd">@property</span>
                
 
                            +    <span class="k">def</span> <span class="nf">component_names</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Component names for logging during training.</span>
                
 
                            +<span class="sd">        These correspond to 2nd item in the tuple returned in self.forward(...).</span>
                
 
                            +<span class="sd">        See super_gradients.Trainer.train() docs for more info.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="k">return</span> <span class="p">[</span><span class="s2">&quot;iou&quot;</span><span class="p">,</span> <span class="s2">&quot;obj&quot;</span><span class="p">,</span> <span class="s2">&quot;cls&quot;</span><span class="p">,</span> <span class="s2">&quot;l1&quot;</span><span class="p">,</span> <span class="s2">&quot;num_fg&quot;</span><span class="p">,</span> <span class="s2">&quot;Loss&quot;</span><span class="p">]</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="YoloXDetectionLoss.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXDetectionLoss.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_output</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span class="p">,
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        :param model_output: Union[list, Tuple[torch.Tensor, List]]:</span>
                
 
                             <span class="sd">             When list-</span>
                
@@ -241,11 +253,14 @@
 
                             <span class="sd">        :param ny: int: cells along the y axis (default=20)</span>
                
 
                             <span class="sd">        :return: torch.tensor of xy coordinates of size (1,1,nx,ny,2)</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">yv</span><span class="p">,</span> <span class="n">xv</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">([</span><span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">ny</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">nx</s
                
 
                            +        <span class="k">if</span> <span class="n">torch_version_is_greater_or_equal</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">10</span><span class="p">):</span>
                
 
                            +            <span class="c1"># https://github.com/pytorch/pytorch/issues/50276</span>
                
 
                            +            <span class="n">yv</span><span class="p">,</span> <span class="n">xv</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">([</span><span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">ny</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">n
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="n">yv</span><span class="p">,</span> <span class="n">xv</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">([</span><span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">ny</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">n
                
 
                                     <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">((</span><span class="n">xv</span><span class="p">,</span> <span class="n">yv</span><span class="p">),</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">view</span><span class="p">((</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">ny</span><span
                
 
                            -    <span class="k">def</span> <span class="nf">_compute_loss</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><
                
 
                            -            <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
                
 
                            +    <span class="k">def</span> <span class="nf">_compute_loss</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        :param predictions:     output from all Yolo levels, each of shape</span>
                
 
                             <span class="sd">                                [Batch x 1 x GridSizeY x GridSizeX x (4 + 1 + Num_classes)]</span>
                
@@ -267,7 +282,7 @@
 
                                     <span class="n">obj_targets</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                                     <span class="n">fg_masks</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                            -        <span class="n">num_fg</span><span class="p">,</span> <span class="n">num_gts</span> <span class="o">=</span> <span class="mf">0.</span><span class="p">,</span> <span class="mf">0.</span>
                
 
                            +        <span class="n">num_fg</span><span class="p">,</span> <span class="n">num_gts</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.0</span>
                
 
                                     <span class="k">for</span> <span class="n">image_idx</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">transformed_outputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]):</span>
                
 
                                         <span class="n">labels_im</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="n">targets</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">image_idx</span><span class="p">]</span>
                
@@ -287,21 +302,42 @@
 
                                             <span class="k">try</span><span class="p">:</span>
                
 
                                                 <span class="c1"># assign cells to ground truths, at most one GT per cell</span>
                
 
                            -                    <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">fg_mask</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span><span class="p">,</span> <span class="n">num_fg_img</span> <span class="o">=</span> \
                
 
                            -                        <span class="bp">self</span><span class="o">.</span><span class="n">get_assignments</span><span class="p">(</span><span class="n">image_idx</span><span class="p">,</span> <span class="n">num_gt</span><span class="p">,</span> <span class="n">total_num_anchors</span><span class="p">,</span> <span class="n">gt_bboxes_per_image</span><span class="p">,</span>
                
 
                            -                                             <span class="n">gt_classes</span><span class="p">,</span> <span class="n">bboxes_preds_per_image</span><span class="p">,</span>
                
 
                            -                                             <span class="n">expanded_strides</span><span class="p">,</span> <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span class="n">cls_preds</span><span class="p">,</span> <span class="n">obj_preds</span><span class="p">)</span>
                
 
                            +                    <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">fg_mask</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span><span class="p">,</span> <span class="n">num_fg_img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_assignments</span><span class="p">(</span>
                
 
                            +                        <span class="n">image_idx</span><span class="p">,</span>
                
 
                            +                        <span class="n">num_gt</span><span class="p">,</span>
                
 
                            +                        <span class="n">total_num_anchors</span><span class="p">,</span>
                
 
                            +                        <span class="n">gt_bboxes_per_image</span><span class="p">,</span>
                
 
                            +                        <span class="n">gt_classes</span><span class="p">,</span>
                
 
                            +                        <span class="n">bboxes_preds_per_image</span><span class="p">,</span>
                
 
                            +                        <span class="n">expanded_strides</span><span class="p">,</span>
                
 
                            +                        <span class="n">x_shifts</span><span class="p">,</span>
                
 
                            +                        <span class="n">y_shifts</span><span class="p">,</span>
                
 
                            +                        <span class="n">cls_preds</span><span class="p">,</span>
                
 
                            +                        <span class="n">obj_preds</span><span class="p">,</span>
                
 
                            +                    <span class="p">)</span>
                
 
                                             <span class="c1"># TODO: CHECK IF ERROR IS CUDA OUT OF MEMORY</span>
                
 
                                             <span class="k">except</span> <span class="ne">RuntimeError</span><span class="p">:</span>
                
 
                            -                    <span class="n">logging</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="s2">&quot;OOM RuntimeError is raised due to the huge memory cost during label assignment. </span><span class="se">\</span>
                
 
                            +                    <span class="n">logging</span><span class="o">.</span><span class="n">error</span><span class="p">(</span>
                
 
                            +                        <span class="s2">&quot;OOM RuntimeError is raised due to the huge memory cost during label assignment. </span><span class="se">\</span>
                
 
                             <span class="s2">                                   CPU mode is applied in this batch. If you want to avoid this issue, </span><span class="se">\</span>
                
 
                            -<span class="s2">                                   try to reduce the batch size or image size.&quot;</span><span class="p">)</span>
                
 
                            +<span class="s2">                                   try to reduce the batch size or image size.&quot;</span>
                
 
                            +                    <span class="p">)</span>
                
 
                                                 <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">empty_cache</span><span class="p">()</span>
                
 
                            -                    <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">fg_mask</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span><span class="p">,</span> <span class="n">num_fg_img</span> <span class="o">=</span> \
                
 
                            -                        <span class="bp">self</span><span class="o">.</span><span class="n">get_assignments</span><span class="p">(</span><span class="n">image_idx</span><span class="p">,</span> <span class="n">num_gt</span><span class="p">,</span> <span class="n">total_num_anchors</span><span class="p">,</span> <span class="n">gt_bboxes_per_image</span><span class="p">,</span>
                
 
                            -                                             <span class="n">gt_classes</span><span class="p">,</span> <span class="n">bboxes_preds_per_image</span><span class="p">,</span>
                
 
                            -                                             <span class="n">expanded_strides</span><span class="p">,</span> <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span class="n">cls_preds</span><span class="p">,</span> <span class="n">obj_preds</span><span class="p">,</span> <span class="s1">&#39;cpu&#39;</span><span class="p">)</span>
                
 
                            +                    <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">fg_mask</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span><span class="p">,</span> <span class="n">num_fg_img</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_assignments</span><span class="p">(</span>
                
 
                            +                        <span class="n">image_idx</span><span class="p">,</span>
                
 
                            +                        <span class="n">num_gt</span><span class="p">,</span>
                
 
                            +                        <span class="n">total_num_anchors</span><span class="p">,</span>
                
 
                            +                        <span class="n">gt_bboxes_per_image</span><span class="p">,</span>
                
 
                            +                        <span class="n">gt_classes</span><span class="p">,</span>
                
 
                            +                        <span class="n">bboxes_preds_per_image</span><span class="p">,</span>
                
 
                            +                        <span class="n">expanded_strides</span><span class="p">,</span>
                
 
                            +                        <span class="n">x_shifts</span><span class="p">,</span>
                
 
                            +                        <span class="n">y_shifts</span><span class="p">,</span>
                
 
                            +                        <span class="n">cls_preds</span><span class="p">,</span>
                
 
                            +                        <span class="n">obj_preds</span><span class="p">,</span>
                
 
                            +                        <span class="s2">&quot;cpu&quot;</span><span class="p">,</span>
                
 
                            +                    <span class="p">)</span>
                
 
                                             <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">empty_cache</span><span class="p">()</span>
                
 
                                             <span class="n">num_fg</span> <span class="o">+=</span> <span class="n">num_fg_img</span>
                
@@ -310,9 +346,13 @@
 
                                             <span class="n">obj_target</span> <span class="o">=</span> <span class="n">fg_mask</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
                
 
                                             <span class="n">reg_target</span> <span class="o">=</span> <span class="n">gt_bboxes_per_image</span><span class="p">[</span><span class="n">matched_gt_inds</span><span class="p">]</span>
                
 
                                             <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_l1</span><span class="p">:</span>
                
 
                            -                    <span class="n">l1_target</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_l1_target</span><span class="p">(</span><span class="n">transformed_outputs</span><span class="o">.</span><span class="n">new_zeros</span><span class="p">((</span><span class="n">num_fg_img</span><span class="p">,</span> <span class="mi">4</span><span class="p">)),</span>
                
 
                            -                                                   <span class="n">gt_bboxes_per_image</span><span class="p">[</span><span class="n">matched_gt_inds</span><span class="p">],</span> <span class="n">expanded_strides</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">fg_mask</span><span class="p">],</span>
                
 
                            -                                                   <span class="n">x_shifts</span><span class="o">=</span><span class="n">x_shifts</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">fg_mask</span><span class="p">],</span> <span class="n">y_shifts</span><span class="o">=</span><span class="n">y_shifts</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">fg_mask</span><span class="p">])</span>
                
 
                            +                    <span class="n">l1_target</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_l1_target</span><span class="p">(</span>
                
 
                            +                        <span class="n">transformed_outputs</span><span class="o">.</span><span class="n">new_zeros</span><span class="p">((</span><span class="n">num_fg_img</span><span class="p">,</span> <span class="mi">4</span><span class="p">)),</span>
                
 
                            +                        <span class="n">gt_bboxes_per_image</span><span class="p">[</span><span class="n">matched_gt_inds</span><span class="p">],</span>
                
 
                            +                        <span class="n">expanded_strides</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">fg_mask</span><span class="p">],</span>
                
 
                            +                        <span class="n">x_shifts</span><span class="o">=</span><span class="n">x_shifts</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">fg_mask</span><span class="p">],</span>
                
 
                            +                        <span class="n">y_shifts</span><span class="o">=</span><span class="n">y_shifts</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">fg_mask</span><span class="p">],</span>
                
 
                            +                    <span class="p">)</span>
                
 
                                         <span class="c1"># collect targets for all loss terms over the whole batch</span>
                
 
                                         <span class="n">cls_targets</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">cls_target</span><span class="p">)</span>
                
@@ -343,13 +383,21 @@
 
                                     <span class="n">reg_weight</span> <span class="o">=</span> <span class="mf">5.0</span>
                
 
                                     <span class="n">loss</span> <span class="o">=</span> <span class="n">reg_weight</span> <span class="o">*</span> <span class="n">loss_iou</span> <span class="o">+</span> <span class="n">loss_obj</span> <span class="o">+</span> <span class="n">loss_cls</span> <span class="o">+</span> <span class="n">loss_l1</span>
                
 
                            -        <span class="k">return</span> <span class="n">loss</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">loss_iou</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span> <span class="n">loss_obj</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span cl
                
 
                            -                                <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">loss_l1</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">loss</span><span class="o">.</span><span class="n">device</span><span class="p">),</span>
                
 
                            -                                <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">num_fg</span> <span class="o">/</span> <span class="nb">max</span><span class="p">(</span><span class="n">num_gts</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</spa
                
 
                            -                                <span class="n">loss</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)))</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="YoloXDetectionLoss.prepare_predictions"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoloXDetectionLoss.prepare_predictions">[docs]</a>    <span class="k">def</span> <span class="nf">prepare_predictions</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><spa
                
 
                            -            <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class
                
 
                            +        <span class="k">return</span> <span class="p">(</span>
                
 
                            +            <span class="n">loss</span><span class="p">,</span>
                
 
                            +            <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
                
 
                            +                <span class="p">(</span>
                
 
                            +                    <span class="n">loss_iou</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                    <span class="n">loss_obj</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                    <span class="n">loss_cls</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                    <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">loss_l1</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">loss</span><span class="o">.</span><span class="n">device</span><span class="p">),</span>
                
 
                            +                    <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">num_fg</span> <span class="o">/</span> <span class="nb">max</span><span class="p">(</span><span class="n">num_gts</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span clas
                
 
                            +                    <span class="n">loss</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                <span class="p">)</span>
                
 
                            +            <span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">(),</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="YoloXDetectionLoss.prepare_predictions"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXDetectionLoss.prepare_predictions">[docs]</a>    <span class="k">def</span> <span class="nf">prepare_predictions</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Convert raw outputs of the network into a format that merges outputs from all levels</span>
                
 
                             <span class="sd">        :param predictions:     output from all Yolo levels, each of shape</span>
                
@@ -414,7 +462,7 @@
 
                                     <span class="k">return</span> <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">,</span> <span class="n">transformed_outputs</span><span class="p">,</span> <span class="n">raw_outputs</span></div>
                
 
                            -<div class="viewcode-block" id="YoloXDetectionLoss.get_l1_target"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoloXDetectionLoss.get_l1_target">[docs]</a>    <span class="k">def</span> <span class="nf">get_l1_target</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">l1_target</span><span class="p">,</span> <span class="n">gt</span><span class="p">,</span> <span class="n">stride</s
                
 
                            +<div class="viewcode-block" id="YoloXDetectionLoss.get_l1_target"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXDetectionLoss.get_l1_target">[docs]</a>    <span class="k">def</span> <span class="nf">get_l1_target</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">l1_target</span><span class="p">,</span> <span class="n">gt</span><span class="p">,</span> <span class="n">stride</span><sp
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        :param l1_target:   tensor of zeros of shape [Num_cell_gt_pairs x 4]</span>
                
 
                             <span class="sd">        :param gt:          targets in coordinates [Num_cell_gt_pairs x (4 + 1 + num_classes)]</span>
                
@@ -427,10 +475,24 @@
 
                                     <span class="n">l1_target</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">gt</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="n">stride</span> <span class="o">+</span> <span class="n">eps</span><span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="n">l1_target</span></div>
                
 
                            -<div class="viewcode-block" id="YoloXDetectionLoss.get_assignments"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoloXDetectionLoss.get_assignments">[docs]</a>    <span class="nd">@torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">()</span>
                
 
                            -    <span class="k">def</span> <span class="nf">get_assignments</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">image_idx</span><span class="p">,</span> <span class="n">num_gt</span><span class="p">,</span> <span class="n">total_num_anchors</span><span class="p">,</span> <span class="n">gt_bboxes_per_image</span><span class="p">,</span> <span class="n">gt_classes</span><span class="p">,</span>
                
 
                            -                        <span class="n">bboxes_preds_per_image</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">,</span> <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span class="n">cls_preds</span><span class="p">,</span>
                
 
                            -                        <span class="n">obj_preds</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;gpu&quot;</span><span class="p">,</span> <span class="n">ious_loss_cost_coeff</span><span class="o">=</span><span class="mf">3.0</span><span class="p">,</span> <span class="n">outside_boxes_and_center_cost_coeff</span><span class="o">=</span><span class="mf">100000.0</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="YoloXDetectionLoss.get_assignments"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXDetectionLoss.get_assignments">[docs]</a>    <span class="nd">@torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">()</span>
                
 
                            +    <span class="k">def</span> <span class="nf">get_assignments</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">image_idx</span><span class="p">,</span>
                
 
                            +        <span class="n">num_gt</span><span class="p">,</span>
                
 
                            +        <span class="n">total_num_anchors</span><span class="p">,</span>
                
 
                            +        <span class="n">gt_bboxes_per_image</span><span class="p">,</span>
                
 
                            +        <span class="n">gt_classes</span><span class="p">,</span>
                
 
                            +        <span class="n">bboxes_preds_per_image</span><span class="p">,</span>
                
 
                            +        <span class="n">expanded_strides</span><span class="p">,</span>
                
 
                            +        <span class="n">x_shifts</span><span class="p">,</span>
                
 
                            +        <span class="n">y_shifts</span><span class="p">,</span>
                
 
                            +        <span class="n">cls_preds</span><span class="p">,</span>
                
 
                            +        <span class="n">obj_preds</span><span class="p">,</span>
                
 
                            +        <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;gpu&quot;</span><span class="p">,</span>
                
 
                            +        <span class="n">ious_loss_cost_coeff</span><span class="o">=</span><span class="mf">3.0</span><span class="p">,</span>
                
 
                            +        <span class="n">outside_boxes_and_center_cost_coeff</span><span class="o">=</span><span class="mf">100000.0</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Match cells to ground truth:</span>
                
 
                             <span class="sd">            * at most 1 GT per cell</span>
                
@@ -464,8 +526,7 @@
 
                                         <span class="n">y_shifts</span> <span class="o">=</span> <span class="n">y_shifts</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span>
                
 
                                     <span class="c1"># create a mask for foreground cells</span>
                
 
                            -        <span class="n">fg_mask</span><span class="p">,</span> <span class="n">is_in_boxes_and_center</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_in_boxes_info</span><span class="p">(</span><span class="n">gt_bboxes_per_image</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">,</span>
                
 
                            -                                                                 <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span class="n">total_num_anchors</span><span class="p">,</span> <span class="n">num_gt</span><span class="p">)</span>
                
 
                            +        <span class="n">fg_mask</span><span class="p">,</span> <span class="n">is_in_boxes_and_center</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_in_boxes_info</span><span class="p">(</span><span class="n">gt_bboxes_per_image</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">,</span> <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span 
                
 
                                     <span class="n">bboxes_preds_per_image</span> <span class="o">=</span> <span class="n">bboxes_preds_per_image</span><span class="p">[</span><span class="n">fg_mask</span><span class="p">]</span>
                
 
                                     <span class="n">cls_preds_</span> <span class="o">=</span> <span class="n">cls_preds</span><span class="p">[</span><span class="n">image_idx</span><span class="p">][</span><span class="n">fg_mask</span><span class="p">]</span>
                
@@ -490,12 +551,10 @@
 
                                         <span class="n">pair_wise_cls_loss</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">binary_cross_entropy</span><span class="p">(</span><span class="n">cls_preds_</span><span class="o">.</span><span class="n">sqrt_</span><span class="p">(),</span> <span class="n">gt_cls_per_image</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">)</span><span cla
                
 
                                     <span class="k">del</span> <span class="n">cls_preds_</span>
                
 
                            -        <span class="n">cost</span> <span class="o">=</span> <span class="n">pair_wise_cls_loss</span> <span class="o">+</span> <span class="n">ious_loss_cost_coeff</span> <span class="o">*</span> <span class="n">pair_wise_ious_loss</span> <span class="o">+</span> <span class="n">outside_boxes_and_center_cost_coeff</span> <span class="o">*</span> <span class="p">(</span>
                
 
                            -            <span class="o">~</span><span class="n">is_in_boxes_and_center</span><span class="p">)</span>
                
 
                            +        <span class="n">cost</span> <span class="o">=</span> <span class="n">pair_wise_cls_loss</span> <span class="o">+</span> <span class="n">ious_loss_cost_coeff</span> <span class="o">*</span> <span class="n">pair_wise_ious_loss</span> <span class="o">+</span> <span class="n">outside_boxes_and_center_cost_coeff</span> <span class="o">*</span> <span class="p">(</span><span class="o">~</span><span class="n">is_in_boxes_and_center</span><span class="p">)</span>
                
 
                                     <span class="c1"># further filter foregrounds: create pairs between cells and ground truth, based on cost and IoUs</span>
                
 
                            -        <span class="n">num_fg</span><span class="p">,</span> <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span> <span class="o">=</span> \
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_k_matching</span><span class="p">(</span><span class="n">cost</span><span class="p">,</span> <span class="n">pair_wise_ious</span><span class="p">,</span> <span class="n">gt_classes</span><span class="p">,</span> <span class="n">num_gt</span><span class="p">,</span> <span class="n">fg_mask</span><span class="p">)</span>
                
 
                            +        <span class="n">num_fg</span><span class="p">,</span> <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_k_matching</span><span class="p">(</span><span class="n">cost</span><span class="p">,</span> <span class="n">pair_wise_ious</span><span class="p">,</span> <span
                
 
                                     <span class="c1"># discard tensors related to cost</span>
                
 
                                     <span class="k">del</span> <span class="n">pair_wise_cls_loss</span><span class="p">,</span> <span class="n">cost</span><span class="p">,</span> <span class="n">pair_wise_ious</span><span class="p">,</span> <span class="n">pair_wise_ious_loss</span>
                
@@ -507,7 +566,7 @@
 
                                     <span class="k">return</span> <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">fg_mask</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span><span class="p">,</span> <span class="n">num_fg</span></div>
                
 
                            -<div class="viewcode-block" id="YoloXDetectionLoss.get_in_boxes_info"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoloXDetectionLoss.get_in_boxes_info">[docs]</a>    <span class="k">def</span> <span class="nf">get_in_boxes_info</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">gt_bboxes_per_image</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p
                
 
                            +<div class="viewcode-block" id="YoloXDetectionLoss.get_in_boxes_info"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXDetectionLoss.get_in_boxes_info">[docs]</a>    <span class="k">def</span> <span class="nf">get_in_boxes_info</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">gt_bboxes_per_image</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">,</sp
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Create a mask for all cells, mask in only foreground: cells that have a center located:</span>
                
 
                             <span class="sd">            * withing a GT box;</span>
                
@@ -561,14 +620,18 @@
 
                                     <span class="c1"># FIND CELL CENTERS THAT ARE WITHIN +- self.center_sampling_radius CELLS FROM GROUND TRUTH BOXES CENTERS</span>
                
 
                                     <span class="c1"># define fake boxes: instead of ground truth boxes step +- self.center_sampling_radius from their centers</span>
                
 
                            -        <span class="n">gt_bboxes_per_image_l</span> <span class="o">=</span> <span class="p">((</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">total
                
 
                            -                                 <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
                
 
                            -        <span class="n">gt_bboxes_per_image_r</span> <span class="o">=</span> <span class="p">((</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">total
                
 
                            -                                 <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
                
 
                            -        <span class="n">gt_bboxes_per_image_t</span> <span class="o">=</span> <span class="p">((</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">total
                
 
                            -                                 <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
                
 
                            -        <span class="n">gt_bboxes_per_image_b</span> <span class="o">=</span> <span class="p">((</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">total
                
 
                            -                                 <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">))</span>
                
 
                            +        <span class="n">gt_bboxes_per_image_l</span> <span class="o">=</span> <span class="p">(</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span>
                
 
                            +            <span class="mi">1</span><span class="p">,</span> <span class="n">total_num_anchors</span>
                
 
                            +        <span class="p">)</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
                
 
                            +        <span class="n">gt_bboxes_per_image_r</span> <span class="o">=</span> <span class="p">(</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span>
                
 
                            +            <span class="mi">1</span><span class="p">,</span> <span class="n">total_num_anchors</span>
                
 
                            +        <span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
                
 
                            +        <span class="n">gt_bboxes_per_image_t</span> <span class="o">=</span> <span class="p">(</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span>
                
 
                            +            <span class="mi">1</span><span class="p">,</span> <span class="n">total_num_anchors</span>
                
 
                            +        <span class="p">)</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
                
 
                            +        <span class="n">gt_bboxes_per_image_b</span> <span class="o">=</span> <span class="p">(</span><span class="n">gt_bboxes_per_image</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span>
                
 
                            +            <span class="mi">1</span><span class="p">,</span> <span class="n">total_num_anchors</span>
                
 
                            +        <span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">center_sampling_radius</span> <span class="o">*</span> <span class="n">expanded_strides_per_image</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
                
 
                                     <span class="n">c_l</span> <span class="o">=</span> <span class="n">x_centers_per_image</span> <span class="o">-</span> <span class="n">gt_bboxes_per_image_l</span>
                
 
                                     <span class="n">c_r</span> <span class="o">=</span> <span class="n">gt_bboxes_per_image_r</span> <span class="o">-</span> <span class="n">x_centers_per_image</span>
                
@@ -582,10 +645,10 @@
 
                                     <span class="n">is_in_boxes_anchor</span> <span class="o">=</span> <span class="n">is_in_boxes_all</span> <span class="o">|</span> <span class="n">is_in_centers_all</span>
                
 
                                     <span class="c1"># in boxes AND in centers, preserving a shape [num_GTs x num_FGs]</span>
                
 
                            -        <span class="n">is_in_boxes_and_center</span> <span class="o">=</span> <span class="p">(</span><span class="n">is_in_boxes</span><span class="p">[:,</span> <span class="n">is_in_boxes_anchor</span><span class="p">]</span> <span class="o">&amp;</span> <span class="n">is_in_centers</span><span class="p">[:,</span> <span class="n">is_in_boxes_anchor</span><span class="p">])</span>
                
 
                            +        <span class="n">is_in_boxes_and_center</span> <span class="o">=</span> <span class="n">is_in_boxes</span><span class="p">[:,</span> <span class="n">is_in_boxes_anchor</span><span class="p">]</span> <span class="o">&amp;</span> <span class="n">is_in_centers</span><span class="p">[:,</span> <span class="n">is_in_boxes_anchor</span><span class="p">]</span>
                
 
                                     <span class="k">return</span> <span class="n">is_in_boxes_anchor</span><span class="p">,</span> <span class="n">is_in_boxes_and_center</span></div>
                
 
                            -<div class="viewcode-block" id="YoloXDetectionLoss.dynamic_k_matching"><a class="viewcode-back" href="../../../../super_gradients.training.losses.html#super_gradients.training.losses.YoloXDetectionLoss.dynamic_k_matching">[docs]</a>    <span class="k">def</span> <span class="nf">dynamic_k_matching</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">cost</span><span class="p">,</span> <span class="n">pair_wise_ious</span><span class="p">,</span> <sp
                
 
                            +<div class="viewcode-block" id="YoloXDetectionLoss.dynamic_k_matching"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXDetectionLoss.dynamic_k_matching">[docs]</a>    <span class="k">def</span> <span class="nf">dynamic_k_matching</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">cost</span><span class="p">,</span> <span class="n">pair_wise_ious</span><span class="p">,</span> <span clas
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        :param cost:            pairwise cost, [num_FGs x num_GTs]</span>
                
 
                             <span class="sd">        :param pair_wise_ious:  pairwise IoUs, [num_FGs x num_GTs]</span>
                
@@ -632,6 +695,399 @@
 
                                     <span class="n">pred_ious_this_matching</span> <span class="o">=</span> <span class="p">(</span><span class="n">matching_matrix</span> <span class="o">*</span> <span class="n">pair_wise_ious</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="mi">0</span><span class="p">)[</span><span class="n">fg_mask_inboxes</span><span class="p">]</span>
                
 
                                     <span class="k">return</span> <span class="n">num_fg</span><span class="p">,</span> <span class="n">gt_matched_classes</span><span class="p">,</span> <span class="n">pred_ious_this_matching</span><span class="p">,</span> <span class="n">matched_gt_inds</span></div></div>
                
 
                            +
                
 
                            +
                
 
                            +<div class="viewcode-block" id="YoloXFastDetectionLoss"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.losses.YoloXFastDetectionLoss">[docs]</a><span class="k">class</span> <span class="nc">YoloXFastDetectionLoss</span><span class="p">(</span><span class="n">YoloXDetectionLoss</span><span class="p">):</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">    A completely new implementation of YOLOX loss.</span>
                
 
                            +<span class="sd">    This is NOT an equivalent implementation to the regular yolox loss.</span>
                
 
                            +
                
 
                            +<span class="sd">    * Completely avoids using loops compared to the nested loops in the original implementation.</span>
                
 
                            +<span class="sd">        As a result runs much faster (speedup depends on the type of GPUs, their count, the batch size, etc.).</span>
                
 
                            +<span class="sd">    * Tensors format is very different the original implementation.</span>
                
 
                            +<span class="sd">        Tensors contain image ids, ground truth ids and anchor ids as values to support variable length data.</span>
                
 
                            +<span class="sd">    * There are differences in terms of the algorithm itself:</span>
                
 
                            +<span class="sd">    1. When computing a dynamic k for a ground truth,</span>
                
 
                            +<span class="sd">        in the original implementation they consider the sum of top 10 predictions sorted by ious among the initial</span>
                
 
                            +<span class="sd">        foregrounds of any ground truth in the image,</span>
                
 
                            +<span class="sd">        while in our implementation we consider only the initial foreground of that particular ground truth.</span>
                
 
                            +<span class="sd">        To compensate for that difference we introduce the dynamic_ks_bias hyperparamter which makes the dynamic ks larger.</span>
                
 
                            +<span class="sd">    2. When computing the k matched detections for a ground truth,</span>
                
 
                            +<span class="sd">        in the original implementation they consider the initial foregrounds of any ground truth in the image as candidates,</span>
                
 
                            +<span class="sd">        while in our implementation we consider only the initial foreground of that particular ground truth as candidates.</span>
                
 
                            +<span class="sd">        We believe that this difference is minor.</span>
                
 
                            +
                
 
                            +<span class="sd">    :param dynamic_ks_bias: hyperparameter to compensate for the discrepancies between the regular loss and this loss.</span>
                
 
                            +<span class="sd">    :param sync_num_fgs:    sync num of fgs.</span>
                
 
                            +<span class="sd">                            Can be used for DDP training.</span>
                
 
                            +<span class="sd">    :param obj_loss_fix:    devide by total of num anchors instead num of matching fgs.</span>
                
 
                            +<span class="sd">                            Can be used for objectness loss.</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span> <span class="n">strides</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">,</span> <span class="n">use_l1</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">center_sampling_radius</span><span class="o">=</span><span class="mf">2.5</span><span class="p">,</span> <span class="n">iou_type</span><span class="o">=</span><span class="s2">&quot;iou&quot;</span><s
                
 
                            +    <span class="p">):</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">strides</span><span class="o">=</span><span class="n">strides</span><span class="p">,</span> <span class="n">num_classes</span><span class="o">=</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">use_l1</span><span class="o">=</span><span class="n">use_l1</span><span class="p">,</span> <span class="n">center_s
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_ks_bias</span> <span class="o">=</span> <span class="n">dynamic_ks_bias</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">sync_num_fgs</span> <span class="o">=</span> <span class="n">sync_num_fgs</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">obj_loss_fix</span> <span class="o">=</span> <span class="n">obj_loss_fix</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_compute_loss</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        L = L_objectness + L_iou + L_classification + 1[no_aug_epoch]*L_l1</span>
                
 
                            +<span class="sd">        where:</span>
                
 
                            +<span class="sd">            * L_iou, L_classification and L_l1 are calculated only between cells and targets that suit them;</span>
                
 
                            +<span class="sd">            * L_objectness is calculated for all cells.</span>
                
 
                            +
                
 
                            +<span class="sd">        L_classification:</span>
                
 
                            +<span class="sd">            for cells that have suitable ground truths in their grid locations add BCEs</span>
                
 
                            +<span class="sd">            to force a prediction of IoU with a GT in a multi-label way</span>
                
 
                            +<span class="sd">            Coef: 1.</span>
                
 
                            +<span class="sd">        L_iou:</span>
                
 
                            +<span class="sd">            for cells that have suitable ground truths in their grid locations</span>
                
 
                            +<span class="sd">            add (1 - IoU^2), IoU between a predicted box and each GT box, force maximum IoU</span>
                
 
                            +<span class="sd">            Coef: 1.</span>
                
 
                            +<span class="sd">        L_l1:</span>
                
 
                            +<span class="sd">            for cells that have suitable ground truths in their grid locations</span>
                
 
                            +<span class="sd">            l1 distance between the logits and GTs in “logits” format (the inverse of “logits to predictions” ops)</span>
                
 
                            +<span class="sd">            Coef: 1[no_aug_epoch]</span>
                
 
                            +<span class="sd">        L_objectness:</span>
                
 
                            +<span class="sd">            for each cell add BCE with a label of 1 if there is GT assigned to the cell</span>
                
 
                            +<span class="sd">            Coef: 5</span>
                
 
                            +
                
 
                            +<span class="sd">        :param predictions:     output from all Yolo levels, each of shape</span>
                
 
                            +<span class="sd">                                [Batch x Num_Anchors x GridSizeY x GridSizeX x (4 + 1 + Num_classes)]</span>
                
 
                            +<span class="sd">        :param targets:         [Num_targets x (4 + 2)], values on dim 1 are: image id in a batch, class, box x y w h</span>
                
 
                            +
                
 
                            +<span class="sd">        :return:                loss, all losses separately in a detached tensor</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">,</span> <span class="n">transformed_outputs</span><span class="p">,</span> <span class="n">raw_outputs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">prepare_predictions</span><span class="p">(</span><span class="n">predictions</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">bbox_preds</span> <span class="o">=</span> <span class="n">transformed_outputs</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span>  <span class="c1"># [batch, n_anchors_all, 4]</span>
                
 
                            +        <span class="n">obj_preds</span> <span class="o">=</span> <span class="n">transformed_outputs</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">4</span><span class="p">:</span><span class="mi">5</span><span class="p">]</span>  <span class="c1"># [batch, n_anchors_all, 1]</span>
                
 
                            +        <span class="n">cls_preds</span> <span class="o">=</span> <span class="n">transformed_outputs</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">5</span><span class="p">:]</span>  <span class="c1"># [batch, n_anchors_all, n_cls]</span>
                
 
                            +
                
 
                            +        <span class="c1"># assign cells to ground truths, at most one GT per cell</span>
                
 
                            +        <span class="n">matched_fg_ids</span><span class="p">,</span> <span class="n">matched_gt_classes</span><span class="p">,</span> <span class="n">matched_gt_ids</span><span class="p">,</span> <span class="n">matched_img_ids</span><span class="p">,</span> <span class="n">matched_ious</span><span class="p">,</span> <span class="n">flattened_gts</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_matching</span><span class="p">(</span>
                
 
                            +            <span class="n">bbox_preds</span><span class="p">,</span> <span class="n">cls_preds</span><span class="p">,</span> <span class="n">obj_preds</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">,</span> <span class="n">x_shifts</span><span class="p">,</span> <span class="n">y_shifts</span><span class="p">,</span> <span class="n">targets</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">num_gts</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">flattened_gts</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="mi">1</span><span class="p">)</span>
                
 
                            +        <span class="n">num_fg</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">matched_gt_ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="mi">1</span><span class="p">)</span>
                
 
                            +        <span class="n">total_num_anchors</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">transformed_outputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">transformed_outputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class=
                
 
                            +
                
 
                            +        <span class="n">cls_targets</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">one_hot</span><span class="p">(</span><span class="n">matched_gt_classes</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">int64</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_classes</span><span class="p">)</span
                
 
                            +        <span class="n">obj_targets</span> <span class="o">=</span> <span class="n">transformed_outputs</span><span class="o">.</span><span class="n">new_zeros</span><span class="p">((</span><span class="n">transformed_outputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">transformed_outputs</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</s
                
 
                            +        <span class="n">obj_targets</span><span class="p">[</span><span class="n">matched_img_ids</span><span class="p">,</span> <span class="n">matched_fg_ids</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
                
 
                            +        <span class="n">reg_targets</span> <span class="o">=</span> <span class="n">flattened_gts</span><span class="p">[</span><span class="n">matched_gt_ids</span><span class="p">][:,</span> <span class="mi">1</span><span class="p">:]</span>
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_l1</span><span class="p">:</span>
                
 
                            +            <span class="n">l1_targets</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_l1_target</span><span class="p">(</span>
                
 
                            +                <span class="n">transformed_outputs</span><span class="o">.</span><span class="n">new_zeros</span><span class="p">((</span><span class="n">num_fg</span><span class="p">,</span> <span class="mi">4</span><span class="p">)),</span>
                
 
                            +                <span class="n">flattened_gts</span><span class="p">[</span><span class="n">matched_gt_ids</span><span class="p">][:,</span> <span class="mi">1</span><span class="p">:],</span>
                
 
                            +                <span class="n">expanded_strides</span><span class="o">.</span><span class="n">squeeze</span><span class="p">()[</span><span class="n">matched_fg_ids</span><span class="p">],</span>
                
 
                            +                <span class="n">x_shifts</span><span class="o">=</span><span class="n">x_shifts</span><span class="o">.</span><span class="n">squeeze</span><span class="p">()[</span><span class="n">matched_fg_ids</span><span class="p">],</span>
                
 
                            +                <span class="n">y_shifts</span><span class="o">=</span><span class="n">y_shifts</span><span class="o">.</span><span class="n">squeeze</span><span class="p">()[</span><span class="n">matched_fg_ids</span><span class="p">],</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">sync_num_fgs</span> <span class="ow">and</span> <span class="n">dist</span><span class="o">.</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="n">num_fg</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">scalar_tensor</span><span class="p">(</span><span class="n">num_fg</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">matched_gt_ids</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
                
 
                            +            <span class="n">dist</span><span class="o">.</span><span class="n">all_reduce</span><span class="p">(</span><span class="n">num_fg</span><span class="p">,</span> <span class="n">op</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_distributed_c10d</span><span class="o">.</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">AVG</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">loss_iou</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">iou_loss</span><span class="p">(</span><span class="n">bbox_preds</span><span class="p">[</span><span class="n">matched_img_ids</span><span class="p">,</span> <span class="n">matched_fg_ids</span><span class="p">],</span> <span class="n">reg_targets</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span clas
                
 
                            +        <span class="n">loss_obj</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">bcewithlog_loss</span><span class="p">(</span><span class="n">obj_preds</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">obj_targets</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <
                
 
                            +        <span class="n">loss_cls</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">bcewithlog_loss</span><span class="p">(</span><span class="n">cls_preds</span><span class="p">[</span><span class="n">matched_img_ids</span><span class="p">,</span> <span class="n">matched_fg_ids</span><span class="p">],</span> <span class="n">cls_targets</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <spa
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">use_l1</span><span class="p">:</span>
                
 
                            +            <span class="n">loss_l1</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">l1_loss</span><span class="p">(</span><span class="n">raw_outputs</span><span class="p">[</span><span class="n">matched_img_ids</span><span class="p">,</span> <span class="n">matched_fg_ids</span><span class="p">],</span> <span class="n">l1_targets</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span cl
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="n">loss_l1</span> <span class="o">=</span> <span class="mf">0.0</span>
                
 
                            +
                
 
                            +        <span class="n">reg_weight</span> <span class="o">=</span> <span class="mf">5.0</span>
                
 
                            +        <span class="n">loss</span> <span class="o">=</span> <span class="n">reg_weight</span> <span class="o">*</span> <span class="n">loss_iou</span> <span class="o">+</span> <span class="n">loss_obj</span> <span class="o">+</span> <span class="n">loss_cls</span> <span class="o">+</span> <span class="n">loss_l1</span>
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="p">(</span>
                
 
                            +            <span class="n">loss</span><span class="p">,</span>
                
 
                            +            <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
                
 
                            +                <span class="p">(</span>
                
 
                            +                    <span class="n">loss_iou</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                    <span class="n">loss_obj</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                    <span class="n">loss_cls</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                    <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">loss_l1</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">transformed_outputs</span><span class="o">.</span><span class="n">device</span><span class="p">),</spa
                
 
                            +                    <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">num_fg</span> <span class="o">/</span> <span class="nb">max</span><span class="p">(</span><span class="n">num_gts</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span clas
                
 
                            +                    <span class="n">loss</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span>
                
 
                            +                <span class="p">)</span>
                
 
                            +            <span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">(),</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_get_initial_matching</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span> <span class="n">gt_bboxes</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">expanded_strides</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">x_shifts</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</sp
                
 
                            +    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Get candidates using a mask for all cells.</span>
                
 
                            +<span class="sd">        Mask in only foreground cells that have a center located:</span>
                
 
                            +<span class="sd">            * withing a GT box (param: is_in_boxes);</span>
                
 
                            +<span class="sd">            OR</span>
                
 
                            +<span class="sd">            * within a fixed radius around a GT box (center sampling) (param: is_in_centers);</span>
                
 
                            +
                
 
                            +<span class="sd">        return:</span>
                
 
                            +<span class="sd">            initial_matching: get a list of candidates pairs of (gt box id, anchor box id) based on cell = is_in_boxes | is_in_centers.</span>
                
 
                            +<span class="sd">                              shape: [num_candidates, 2]</span>
                
 
                            +<span class="sd">            strong candidate mask: get a list whether a candidate is a strong one or not.</span>
                
 
                            +<span class="sd">                                   strong candidate is a cell from is_in_boxes &amp; is_in_centers.</span>
                
 
                            +<span class="sd">                                   shape: [num_candidates].</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">cell_x_centers</span> <span class="o">=</span> <span class="p">(</span><span class="n">x_shifts</span> <span class="o">+</span> <span class="mf">0.5</span><span class="p">)</span> <span class="o">*</span> <span class="n">expanded_strides</span>
                
 
                            +        <span class="n">cell_y_centers</span> <span class="o">=</span> <span class="p">(</span><span class="n">y_shifts</span> <span class="o">+</span> <span class="mf">0.5</span><span class="p">)</span> <span class="o">*</span> <span class="n">expanded_strides</span>
                
 
                            +
                
 
                            +        <span class="n">gt_bboxes_x_centers</span> <span class="o">=</span> <span class="n">gt_bboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
                
 
                            +        <span class="n">gt_bboxes_y_centers</span> <span class="o">=</span> <span class="n">gt_bboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">gt_bboxes_half_w</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">gt_bboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
                
 
                            +        <span class="n">gt_bboxes_half_h</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">gt_bboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">])</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">is_in_boxes</span> <span class="o">=</span> <span class="p">(</span>
                
 
                            +            <span class="p">(</span><span class="n">cell_x_centers</span> <span class="o">&gt;</span> <span class="n">gt_bboxes_x_centers</span> <span class="o">-</span> <span class="n">gt_bboxes_half_w</span><span class="p">)</span>
                
 
                            +            <span class="o">&amp;</span> <span class="p">(</span><span class="n">gt_bboxes_x_centers</span> <span class="o">+</span> <span class="n">gt_bboxes_half_w</span> <span class="o">&gt;</span> <span class="n">cell_x_centers</span><span class="p">)</span>
                
 
                            +            <span class="o">&amp;</span> <span class="p">(</span><span class="n">cell_y_centers</span> <span class="o">&gt;</span> <span class="n">gt_bboxes_y_centers</span> <span class="o">-</span> <span class="n">gt_bboxes_half_h</span><span class="p">)</span>
                
 
                            +            <span class="o">&amp;</span> <span class="p">(</span><span class="n">gt_bboxes_y_centers</span> <span class="o">+</span> <span class="n">gt_bboxes_half_h</span> <span class="o">&gt;</span> <span class="n">cell_y_centers</span><span class="p">)</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">radius_shifts</span> <span class="o">=</span> <span class="mf">2.5</span> <span class="o">*</span> <span class="n">expanded_strides</span>
                
 
                            +
                
 
                            +        <span class="n">is_in_centers</span> <span class="o">=</span> <span class="p">(</span>
                
 
                            +            <span class="p">(</span><span class="n">cell_x_centers</span> <span class="o">+</span> <span class="n">radius_shifts</span> <span class="o">&gt;</span> <span class="n">gt_bboxes_x_centers</span><span class="p">)</span>
                
 
                            +            <span class="o">&amp;</span> <span class="p">(</span><span class="n">gt_bboxes_x_centers</span> <span class="o">&gt;</span> <span class="n">cell_x_centers</span> <span class="o">-</span> <span class="n">radius_shifts</span><span class="p">)</span>
                
 
                            +            <span class="o">&amp;</span> <span class="p">(</span><span class="n">cell_y_centers</span> <span class="o">+</span> <span class="n">radius_shifts</span> <span class="o">&gt;</span> <span class="n">gt_bboxes_y_centers</span><span class="p">)</span>
                
 
                            +            <span class="o">&amp;</span> <span class="p">(</span><span class="n">gt_bboxes_y_centers</span> <span class="o">&gt;</span> <span class="n">cell_y_centers</span> <span class="o">-</span> <span class="n">radius_shifts</span><span class="p">)</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">initial_mask</span> <span class="o">=</span> <span class="n">is_in_boxes</span> <span class="o">|</span> <span class="n">is_in_centers</span>
                
 
                            +        <span class="n">initial_matching</span> <span class="o">=</span> <span class="n">initial_mask</span><span class="o">.</span><span class="n">nonzero</span><span class="p">()</span>
                
 
                            +        <span class="n">strong_candidate_mask</span> <span class="o">=</span> <span class="p">(</span><span class="n">is_in_boxes</span> <span class="o">&amp;</span> <span class="n">is_in_centers</span><span class="p">)[</span><span class="n">initial_mask</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="n">initial_matching</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">initial_matching</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">strong_candidate_mask</span>
                
 
                            +
                
 
                            +    <span class="nd">@torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">()</span>
                
 
                            +    <span class="k">def</span> <span class="nf">_compute_matching</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">bbox_preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">cls_preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">obj_preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">expanded_strides</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">x_shifts</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">y_shifts</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">labels</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">ious_loss_cost_coeff</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">3.0</span><span class="p">,</span>
                
 
                            +        <span class="n">outside_boxes_and_center_cost_coeff</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">100000.0</span><span class="p">,</span>
                
 
                            +    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Ten
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Match cells to ground truth:</span>
                
 
                            +<span class="sd">            * at most 1 GT per cell</span>
                
 
                            +<span class="sd">            * dynamic number of cells per GT</span>
                
 
                            +
                
 
                            +<span class="sd">        :param bbox_preds: predictions of bounding boxes. shape [batch, n_anchors_all, 4]</span>
                
 
                            +<span class="sd">        :param cls_preds:  predictions of class.          shape [batch, n_anchors_all, n_cls]</span>
                
 
                            +<span class="sd">        :param obj_preds:  predictions for objectness.    shape [batch, n_anchors_all, 1]</span>
                
 
                            +<span class="sd">        :param expanded_strides:  stride of the output grid the prediction is coming from. shape [1, n_anchors_all]</span>
                
 
                            +<span class="sd">        :param x_shifts: x coordinate on the grid cell the prediction is coming from.      shape [1, n_anchors_all]</span>
                
 
                            +<span class="sd">        :param y_shifts: y coordinate on the grid cell the prediction is coming from.      shape [1, n_anchors_all]</span>
                
 
                            +<span class="sd">        :param labels:   labels for each grid cell.  shape [n_anchors_all, (4 + 2)]</span>
                
 
                            +<span class="sd">        :return: candidate_fg_ids       shape [num_fg]</span>
                
 
                            +<span class="sd">                 candidate_gt_classes   shape [num_fg]</span>
                
 
                            +<span class="sd">                 candidate_gt_ids       shape [num_fg]</span>
                
 
                            +<span class="sd">                 candidate_img_ids      shape [num_fg]</span>
                
 
                            +<span class="sd">                 candidate_ious         shape [num_fg]</span>
                
 
                            +<span class="sd">                 flattened_gts          shape [num_gts, 5]</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +
                
 
                            +        <span class="n">flattened_gts</span><span class="p">,</span> <span class="n">gt_id_to_img_id</span> <span class="o">=</span> <span class="n">labels</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">:],</span> <span class="n">labels</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">type</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">int64</spa
                
 
                            +
                
 
                            +        <span class="c1"># COMPUTE CANDIDATES</span>
                
 
                            +        <span class="n">candidate_gt_ids</span><span class="p">,</span> <span class="n">candidate_fg_ids</span><span class="p">,</span> <span class="n">strong_candidate_mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_initial_matching</span><span class="p">(</span><span class="n">flattened_gts</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">:],</span> <span class="n">expanded_strides</span><span class="p">,</
                
 
                            +        <span class="n">candidate_img_ids</span> <span class="o">=</span> <span class="n">gt_id_to_img_id</span><span class="p">[</span><span class="n">candidate_gt_ids</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_gts_bbox</span> <span class="o">=</span> <span class="n">flattened_gts</span><span class="p">[</span><span class="n">candidate_gt_ids</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span>
                
 
                            +        <span class="n">candidate_det_bbox</span> <span class="o">=</span> <span class="n">bbox_preds</span><span class="p">[</span><span class="n">candidate_img_ids</span><span class="p">,</span> <span class="n">candidate_fg_ids</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="c1"># COMPUTE DYNAMIC KS</span>
                
 
                            +        <span class="n">candidate_ious</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_calculate_pairwise_bbox_iou</span><span class="p">(</span><span class="n">candidate_gts_bbox</span><span class="p">,</span> <span class="n">candidate_det_bbox</span><span class="p">,</span> <span class="n">xyxy</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
                
 
                            +        <span class="n">dynamic_ks</span><span class="p">,</span> <span class="n">matching_index_to_dynamic_k_index</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_dynamic_ks</span><span class="p">(</span><span class="n">candidate_gt_ids</span><span class="p">,</span> <span class="n">candidate_ious</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_ks_bias</span><span class="p">)
                
 
                            +        <span class="k">del</span> <span class="n">candidate_gts_bbox</span><span class="p">,</span> <span class="n">candidate_det_bbox</span>
                
 
                            +
                
 
                            +        <span class="c1"># ORDER CANDIDATES BY COST</span>
                
 
                            +        <span class="n">candidate_gt_classes</span> <span class="o">=</span> <span class="n">flattened_gts</span><span class="p">[</span><span class="n">candidate_gt_ids</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
                
 
                            +        <span class="n">cost_order</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_cost_order</span><span class="p">(</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">num_classes</span><span class="p">,</span>
                
 
                            +            <span class="n">candidate_img_ids</span><span class="p">,</span>
                
 
                            +            <span class="n">candidate_gt_classes</span><span class="p">,</span>
                
 
                            +            <span class="n">candidate_fg_ids</span><span class="p">,</span>
                
 
                            +            <span class="n">candidate_ious</span><span class="p">,</span>
                
 
                            +            <span class="n">cls_preds</span><span class="p">,</span>
                
 
                            +            <span class="n">obj_preds</span><span class="p">,</span>
                
 
                            +            <span class="n">strong_candidate_mask</span><span class="p">,</span>
                
 
                            +            <span class="n">ious_loss_cost_coeff</span><span class="p">,</span>
                
 
                            +            <span class="n">outside_boxes_and_center_cost_coeff</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">candidate_gt_ids</span> <span class="o">=</span> <span class="n">candidate_gt_ids</span><span class="p">[</span><span class="n">cost_order</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_gt_classes</span> <span class="o">=</span> <span class="n">candidate_gt_classes</span><span class="p">[</span><span class="n">cost_order</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_img_ids</span> <span class="o">=</span> <span class="n">candidate_img_ids</span><span class="p">[</span><span class="n">cost_order</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_fg_ids</span> <span class="o">=</span> <span class="n">candidate_fg_ids</span><span class="p">[</span><span class="n">cost_order</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_ious</span> <span class="o">=</span> <span class="n">candidate_ious</span><span class="p">[</span><span class="n">cost_order</span><span class="p">]</span>
                
 
                            +        <span class="n">matching_index_to_dynamic_k_index</span> <span class="o">=</span> <span class="n">matching_index_to_dynamic_k_index</span><span class="p">[</span><span class="n">cost_order</span><span class="p">]</span>
                
 
                            +        <span class="k">del</span> <span class="n">cost_order</span>
                
 
                            +
                
 
                            +        <span class="c1"># FILTER MATCHING TO LOWEST K COST MATCHES PER GT</span>
                
 
                            +        <span class="n">ranks</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_ranks</span><span class="p">(</span><span class="n">candidate_gt_ids</span><span class="p">)</span>
                
 
                            +        <span class="n">corresponding_dynamic_ks</span> <span class="o">=</span> <span class="n">dynamic_ks</span><span class="p">[</span><span class="n">matching_index_to_dynamic_k_index</span><span class="p">]</span>
                
 
                            +        <span class="n">topk_mask</span> <span class="o">=</span> <span class="n">ranks</span> <span class="o">&lt;</span> <span class="n">corresponding_dynamic_ks</span>
                
 
                            +
                
 
                            +        <span class="n">candidate_gt_ids</span> <span class="o">=</span> <span class="n">candidate_gt_ids</span><span class="p">[</span><span class="n">topk_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_gt_classes</span> <span class="o">=</span> <span class="n">candidate_gt_classes</span><span class="p">[</span><span class="n">topk_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_img_ids</span> <span class="o">=</span> <span class="n">candidate_img_ids</span><span class="p">[</span><span class="n">topk_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_fg_ids</span> <span class="o">=</span> <span class="n">candidate_fg_ids</span><span class="p">[</span><span class="n">topk_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_ious</span> <span class="o">=</span> <span class="n">candidate_ious</span><span class="p">[</span><span class="n">topk_mask</span><span class="p">]</span>
                
 
                            +        <span class="k">del</span> <span class="n">ranks</span><span class="p">,</span> <span class="n">topk_mask</span><span class="p">,</span> <span class="n">dynamic_ks</span><span class="p">,</span> <span class="n">matching_index_to_dynamic_k_index</span><span class="p">,</span> <span class="n">corresponding_dynamic_ks</span>
                
 
                            +
                
 
                            +        <span class="c1"># FILTER MATCHING TO AT MOST 1 MATCH FOR DET BY TAKING THE LOWEST COST MATCH</span>
                
 
                            +        <span class="n">candidate_img_and_fg_ids_combined</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_combine_candidates_img_id_fg_id</span><span class="p">(</span><span class="n">candidate_img_ids</span><span class="p">,</span> <span class="n">candidate_fg_ids</span><span class="p">)</span>
                
 
                            +        <span class="n">top1_mask</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_is_first_mask</span><span class="p">(</span><span class="n">candidate_img_and_fg_ids_combined</span><span class="p">)</span>
                
 
                            +        <span class="n">candidate_gt_ids</span> <span class="o">=</span> <span class="n">candidate_gt_ids</span><span class="p">[</span><span class="n">top1_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_gt_classes</span> <span class="o">=</span> <span class="n">candidate_gt_classes</span><span class="p">[</span><span class="n">top1_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_fg_ids</span> <span class="o">=</span> <span class="n">candidate_fg_ids</span><span class="p">[</span><span class="n">top1_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_img_ids</span> <span class="o">=</span> <span class="n">candidate_img_ids</span><span class="p">[</span><span class="n">top1_mask</span><span class="p">]</span>
                
 
                            +        <span class="n">candidate_ious</span> <span class="o">=</span> <span class="n">candidate_ious</span><span class="p">[</span><span class="n">top1_mask</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="n">candidate_fg_ids</span><span class="p">,</span> <span class="n">candidate_gt_classes</span><span class="p">,</span> <span class="n">candidate_gt_ids</span><span class="p">,</span> <span class="n">candidate_img_ids</span><span class="p">,</span> <span class="n">candidate_ious</span><span class="p">,</span> <span class="n">flattened_gts</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_combine_candidates_img_id_fg_id</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">candidate_img_ids</span><span class="p">,</span> <span class="n">candidate_anchor_ids</span><span class="p">):</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Create one dim tensor with unique pairs of img_id and fg_id.</span>
                
 
                            +<span class="sd">        e.g: candidate_img_ids = [0,1,0,0]</span>
                
 
                            +<span class="sd">             candidate_fg_ids = [0,0,0,1]</span>
                
 
                            +<span class="sd">             result = [0,1,0,2]</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">candidate_img_and_fg_ids_combined</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">((</span><span class="n">candidate_img_ids</span><span class="p">,</span> <span class="n">candidate_anchor_ids</span><span class="p">),</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">unique</span><span class="p">(<
                
 
                            +        <span class="k">return</span> <span class="n">candidate_img_and_fg_ids_combined</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_compute_dynamic_ks</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ids</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">ious</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">dynamic_ks_bias</
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        :param ids:                 ids of GTs, shape: [num_candidates]</span>
                
 
                            +<span class="sd">        :param ious:                pairwise IoUs, shape: [num_candidates]</span>
                
 
                            +<span class="sd">        :param dynamic_ks_bias:     multiply the resulted k to compensate the regular loss</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">ids</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;ids must be of shape [num_candidates]&quot;</span>
                
 
                            +        <span class="k">assert</span> <span class="nb">len</span><span class="p">(</span><span class="n">ious</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;ious must be of shape [num_candidates]&quot;</span>
                
 
                            +        <span class="k">assert</span> <span class="n">ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">ious</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="s2">&quot;num of ids.shape[0] must be the same as num of ious.shape[0]&quot;</span>
                
 
                            +        <span class="c1"># sort ious and ids by ious</span>
                
 
                            +        <span class="n">ious</span><span class="p">,</span> <span class="n">ious_argsort</span> <span class="o">=</span> <span class="n">ious</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">descending</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +        <span class="n">ids</span> <span class="o">=</span> <span class="n">ids</span><span class="p">[</span><span class="n">ious_argsort</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="c1"># stable sort indices, so that ious are first sorted by id and second by value</span>
                
 
                            +        <span class="n">ids</span><span class="p">,</span> <span class="n">ids_argsort</span> <span class="o">=</span> <span class="n">ids</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">stable</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +        <span class="n">ious</span> <span class="o">=</span> <span class="n">ious</span><span class="p">[</span><span class="n">ids_argsort</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="n">unique_ids</span><span class="p">,</span> <span class="n">ids_index_to_unique_ids_index</span> <span class="o">=</span> <span class="n">ids</span><span class="o">.</span><span class="n">unique_consecutive</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +        <span class="n">num_unique_ids</span> <span class="o">=</span> <span class="n">unique_ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">10</span><span class="p">:</span>
                
 
                            +            <span class="n">is_in_top_10</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">10</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">bool</span><span class="p">,</span> <span clas
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="n">is_in_top_10</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">ids</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">bool</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">dynamic_ks</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">num_unique_ids</span><span class="p">,),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">ious</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">ious</span><span class="o">.</span><sp
                
 
                            +        <span class="n">dynamic_ks</span><span class="o">.</span><span class="n">index_put_</span><span class="p">((</span><span class="n">ids_index_to_unique_ids_index</span><span class="p">,),</span> <span class="n">is_in_top_10</span> <span class="o">*</span> <span class="n">ious</span><span class="p">,</span> <span class="n">accumulate</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +        <span class="k">if</span> <span class="n">dynamic_ks_bias</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="n">dynamic_ks</span> <span class="o">*=</span> <span class="n">dynamic_ks_bias</span>
                
 
                            +        <span class="n">dynamic_ks</span> <span class="o">=</span> <span class="n">dynamic_ks</span><span class="o">.</span><span class="n">long</span><span class="p">()</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="nb">min</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">all_argsort</span> <span class="o">=</span> <span class="n">ious_argsort</span><span class="p">[</span><span class="n">ids_argsort</span><span class="p">]</span>
                
 
                            +        <span class="n">inverse_all_argsort</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">ious_argsort</span><span class="p">)</span>
                
 
                            +        <span class="n">inverse_all_argsort</span><span class="p">[</span><span class="n">all_argsort</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">all_argsort</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">all_arg
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="n">dynamic_ks</span><span class="p">,</span> <span class="n">ids_index_to_unique_ids_index</span><span class="p">[</span><span class="n">inverse_all_argsort</span><span class="p">]</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_compute_cost_order</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">num_classes</span><span class="p">,</span>
                
 
                            +        <span class="n">candidate_gt_img_ids</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">candidate_gt_classes</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">candidate_anchor_ids</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">candidate_ious</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">cls_preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">obj_preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">strong_candidate_mask</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                
 
                            +        <span class="n">ious_loss_cost_coeff</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                
 
                            +        <span class="n">outside_boxes_and_center_cost_coeff</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                
 
                            +    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
                
 
                            +        <span class="n">gt_cls_per_image</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">one_hot</span><span class="p">(</span><span class="n">candidate_gt_classes</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">int64</span><span class="p">),</span> <span class="n">num_classes</span><span class="p">)</span><span class="o">.</span><span class="n">floa
                
 
                            +        <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">amp</span><span class="o">.</span><span class="n">autocast</span><span class="p">(</span><span class="n">enabled</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
                
 
                            +            <span class="n">cls_preds_</span> <span class="o">=</span> <span class="p">(</span>
                
 
                            +                <span class="n">cls_preds</span><span class="p">[</span><span class="n">candidate_gt_img_ids</span><span class="p">,</span> <span class="n">candidate_anchor_ids</span><span class="p">]</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">sigmoid_</span><span class="p">()</span>
                
 
                            +                <span class="o">*</span> <span class="n">obj_preds</span><span class="p">[</span><span class="n">candidate_gt_img_ids</span><span class="p">,</span> <span class="n">candidate_anchor_ids</span><span class="p">]</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">sigmoid_</span><span class="p">()</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +            <span class="n">pair_wise_cls_cost</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">binary_cross_entropy</span><span class="p">(</span><span class="n">cls_preds_</span><span class="o">.</span><span class="n">sqrt_</span><span class="p">(),</span> <span class="n">gt_cls_per_image</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">)</span><span cla
                
 
                            +
                
 
                            +        <span class="n">ious_cost</span> <span class="o">=</span> <span class="o">-</span><span class="n">torch</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">candidate_ious</span> <span class="o">+</span> <span class="mf">1e-8</span><span class="p">)</span>
                
 
                            +        <span class="n">cost</span> <span class="o">=</span> <span class="n">pair_wise_cls_cost</span> <span class="o">+</span> <span class="n">ious_loss_cost_coeff</span> <span class="o">*</span> <span class="n">ious_cost</span> <span class="o">+</span> <span class="n">outside_boxes_and_center_cost_coeff</span> <span class="o">*</span> <span class="n">strong_candidate_mask</span><span class="o">.</span><span class="n">logical_not</span><span class="p">()</span>
                
 
                            +        <span class="k">return</span> <span class="n">cost</span><span class="o">.</span><span class="n">argsort</span><span class="p">()</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_calculate_pairwise_bbox_iou</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">bboxes_a</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">bboxes_b</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n"
                
 
                            +        <span class="k">if</span> <span class="n">bboxes_a</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">4</span> <span class="ow">or</span> <span class="n">bboxes_b</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">4</span><span class="p">:</sp
                
 
                            +            <span class="k">raise</span> <span class="ne">IndexError</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">xyxy</span><span class="p">:</span>
                
 
                            +            <span class="n">tl</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">bboxes_a</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">],</span> <span class="n">bboxes_b</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">])</span>
                
 
                            +            <span class="n">br</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">bboxes_a</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:],</span> <span class="n">bboxes_b</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:])</span>
                
 
                            +            <span class="n">area_a</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">bboxes_a</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">-</span> <span class="n">bboxes_a</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">],</span> <span class="mi">1</span><span class="p">)</span>
                
 
                            +            <span class="n">area_b</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">bboxes_b</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">-</span> <span class="n">bboxes_b</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">],</span> <span class="mi">1</span><span class="p">)</span>
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="n">tl</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span>
                
 
                            +                <span class="p">(</span><span class="n">bboxes_a</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">bboxes_a</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">),</span>
                
 
                            +                <span class="p">(</span><span class="n">bboxes_b</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">bboxes_b</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">),</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +            <span class="n">br</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span>
                
 
                            +                <span class="p">(</span><span class="n">bboxes_a</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">bboxes_a</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">),</span>
                
 
                            +                <span class="p">(</span><span class="n">bboxes_b</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">bboxes_b</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">),</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +
                
 
                            +            <span class="n">area_a</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">bboxes_a</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:],</span> <span class="mi">1</span><span class="p">)</span>
                
 
                            +            <span class="n">area_b</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">bboxes_b</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:],</span> <span class="mi">1</span><span class="p">)</span>
                
 
                            +        <span class="n">en</span> <span class="o">=</span> <span class="p">(</span><span class="n">tl</span> <span class="o">&lt;</span> <span class="n">br</span><span class="p">)</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
                
 
                            +        <span class="n">area_i</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="n">br</span> <span class="o">-</span> <span class="n">tl</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="o">*</span> <span class="n">en</span>
                
 
                            +        <span class="k">return</span> <span class="n">area_i</span> <span class="o">/</span> <span class="p">(</span><span class="n">area_a</span> <span class="o">+</span> <span class="n">area_b</span> <span class="o">-</span> <span class="n">area_i</span><span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_compute_ranks</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ids</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
                
 
                            +        <span class="n">ids</span><span class="p">,</span> <span class="n">ids_argsort</span> <span class="o">=</span> <span class="n">ids</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">stable</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
                
 
                            +            <span class="n">is_not_first</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="mi">1</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">bool</span><span class="p">,</span> <span clas
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="n">is_not_first</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">ids</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">bool</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">subtract</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">ids_argsort</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="
                
 
                            +        <span class="n">subtract</span><span class="p">[</span><span class="n">is_not_first</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
                
 
                            +        <span class="n">subtract</span> <span class="o">=</span> <span class="n">subtract</span><span class="o">.</span><span class="n">cummax</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
                
 
                            +        <span class="n">rank</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">ids_argsort</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">d
                
 
                            +
                
 
                            +        <span class="n">inverse_argsort</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">ids_argsort</span><span class="p">)</span>
                
 
                            +        <span class="n">inverse_argsort</span><span class="p">[</span><span class="n">ids_argsort</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">ids_argsort</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">ids_argsort
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="n">rank</span><span class="p">[</span><span class="n">inverse_argsort</span><span class="p">]</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_compute_is_first_mask</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ids</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Filter fg that matches two gts.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">ids</span><span class="p">,</span> <span class="n">ids_argsort</span> <span class="o">=</span> <span class="n">ids</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">stable</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">ids</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
                
 
                            +            <span class="n">is_first</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">1</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">bool</span><span class="p">,</span> <span class="n"
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="n">is_first</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">ids</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">bool</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">inverse_argsort</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">ids_argsort</span><span class="p">)</span>
                
 
                            +        <span class="n">inverse_argsort</span><span class="p">[</span><span class="n">ids_argsort</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">ids_argsort</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">ids_argsort
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="n">is_first</span><span class="p">[</span><span class="n">inverse_argsort</span><span class="p">]</span></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -661,4 +1117,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.metrics.classification_metrics &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.metrics.classification_metrics &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -92,7 +94,7 @@
 
                             <span class="kn">from</span> <span class="nn">torchmetrics</span> <span class="kn">import</span> <span class="n">Metric</span>
                
 
                            -<div class="viewcode-block" id="accuracy"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.accuracy">[docs]</a><span class="k">def</span> <span class="nf">accuracy</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">topk</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,)):</span>
                
 
                            +<div class="viewcode-block" id="accuracy"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.accuracy">[docs]</a><span class="k">def</span> <span class="nf">accuracy</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">topk</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,)):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Computes the precision@k for the specified values of k</span>
                
 
                             <span class="sd">    :param output: Tensor / Numpy / List</span>
                
 
                             <span class="sd">        The prediction</span>
                
@@ -122,24 +124,26 @@
 
                                 <span class="k">return</span> <span class="n">res</span></div>
                
 
                            -<div class="viewcode-block" id="Accuracy"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Accuracy">[docs]</a><span class="k">class</span> <span class="nc">Accuracy</span><span class="p">(</span><span class="n">torchmetrics</span><span class="o">.</span><span class="n">Accuracy</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="Accuracy"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Accuracy">[docs]</a><span class="k">class</span> <span class="nc">Accuracy</span><span class="p">(</span><span class="n">torchmetrics</span><span class="o">.</span><span class="n">Accuracy</span><span class="p">):</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                            -<div class="viewcode-block" id="Accuracy.update"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Accuracy.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class=
                
 
                            +<div class="viewcode-block" id="Accuracy.update"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Accuracy.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">targ
                
 
                                     <span class="k">if</span> <span class="n">target</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="n">preds</span><span class="o">.</span><span class="n">shape</span><span class="p">:</span>
                
 
                                         <span class="n">target</span> <span class="o">=</span> <span class="n">target</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>  <span class="c1"># supports smooth labels</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">preds</span><span class="o">=</span><span class="n">preds</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="n">target</span><span class="p">)</span></div></div>
                
 
                            -<div class="viewcode-block" id="Top5"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Top5">[docs]</a><span class="k">class</span> <span class="nc">Top5</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="Top5"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Top5">[docs]</a><span class="k">class</span> <span class="nc">Top5</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;correct&quot;</span><span class="p">,</span> <span class="n">default</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="mf">0.</span><span class="p">),</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="s2">&quot;sum&quot;</span><spa
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;correct&quot;</span><span class="p">,</span> <span class="n">default</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="mf">0.0</span><span class="p">),</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="s2">&quot;sum&quot;</span><sp
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;total&quot;</span><span class="p">,</span> <span class="n">default</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="mi">0</span><span class="p">),</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="s2">&quot;sum&quot;</span><span c
                
 
                            -<div class="viewcode-block" id="Top5.update"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Top5.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">targ
                
 
                            +<div class="viewcode-block" id="Top5.update"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Top5.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">target</span
                
 
                                     <span class="k">if</span> <span class="n">target</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="n">preds</span><span class="o">.</span><span class="n">shape</span><span class="p">:</span>
                
 
                                         <span class="n">target</span> <span class="o">=</span> <span class="n">target</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>  <span class="c1"># supports smooth labels</span>
                
@@ -155,21 +159,22 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">correct</span> <span class="o">+=</span> <span class="n">correct5</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">total</span> <span class="o">+=</span> <span class="n">batch_size</span></div>
                
 
                            -<div class="viewcode-block" id="Top5.compute"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Top5.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="Top5.compute"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Top5.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">correct</span><span class="o">.</span><span class="n">float</span><span class="p">()</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">total</span></div></div>
                
 
                            -<div class="viewcode-block" id="ToyTestClassificationMetric"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.ToyTestClassificationMetric">[docs]</a><span class="k">class</span> <span class="nc">ToyTestClassificationMetric</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="ToyTestClassificationMetric"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.ToyTestClassificationMetric">[docs]</a><span class="k">class</span> <span class="nc">ToyTestClassificationMetric</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Dummy classification Mettric object returning 0 always (for testing).</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="ToyTestClassificationMetric.update"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.ToyTestClassificationMetric.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span
                
 
                            +<div class="viewcode-block" id="ToyTestClassificationMetric.update"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.ToyTestClassificationMetric.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span c
                
 
                                     <span class="k">pass</span></div>
                
 
                            -<div class="viewcode-block" id="ToyTestClassificationMetric.compute"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.ToyTestClassificationMetric.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="ToyTestClassificationMetric.compute"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.ToyTestClassificationMetric.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="k">return</span> <span class="mi">0</span></div></div>
                
 
                             </pre></div>
                
@@ -200,4 +205,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.metrics.detection_metrics &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.metrics.detection_metrics &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -94,10 +96,11 @@
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">compute_detection_matching</span><span class="p">,</span> <span class="n">compute_detection_metrics</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">DetectionPostPredictionCallback</span><span class="p">,</span> <span class="n">IouThreshold</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                            +
                
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="DetectionMetrics"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.DetectionMetrics">[docs]</a><span class="k">class</span> <span class="nc">DetectionMetrics</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="DetectionMetrics"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.DetectionMetrics">[docs]</a><span class="k">class</span> <span class="nc">DetectionMetrics</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    DetectionMetrics</span>
                
@@ -121,37 +124,52 @@
 
                             <span class="sd">        accumulate_on_cpu:     Run on CPU regardless of device used in other parts.</span>
                
 
                             <span class="sd">                            This is to avoid &quot;CUDA out of memory&quot; that might happen on GPU (default False)</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">num_cls</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            -                 <span class="n">post_prediction_callback</span><span class="p">:</span> <span class="n">DetectionPostPredictionCallback</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">normalize_targets</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -                 <span class="n">iou_thres</span><span class="p">:</span> <span class="n">IouThreshold</span> <span class="o">=</span> <span class="n">IouThreshold</span><span class="o">.</span><span class="n">MAP_05_TO_095</span><span class="p">,</span>
                
 
                            -                 <span class="n">recall_thres</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">score_thres</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                
 
                            -                 <span class="n">top_k_predictions</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span>
                
 
                            -                 <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -                 <span class="n">accumulate_on_cpu</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">):</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">num_cls</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            +        <span class="n">post_prediction_callback</span><span class="p">:</span> <span class="n">DetectionPostPredictionCallback</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">normalize_targets</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">iou_thres</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">IouThreshold</span><span class="p">,</span> <span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="n">IouThreshold</span><span class="o">.</span><span class="n">MAP_05_TO_095</span><span class="p">,</span>
                
 
                            +        <span class="n">recall_thres</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">score_thres</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                
 
                            +        <span class="n">top_k_predictions</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">accumulate_on_cpu</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">num_cls</span> <span class="o">=</span> <span class="n">num_cls</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">iou_thres</span> <span class="o">=</span> <span class="n">iou_thres</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">map_str</span> <span class="o">=</span> <span class="s1">&#39;mAP@</span><span class="si">%.1f</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">iou_thres</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">iou_thres</span><span class="o">.</span><span class="n">is_range</span><span class="p">()</
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">component_names</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;Precision&quot;</span><span class="p">,</span> <span class="s2">&quot;Recall&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">map_str</span><span class="p">,</span> <span class="s2">&quot;F1&quot;</span><span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">iou_thres</span><span class="p">,</span> <span class="n">IouThreshold</span><span class="p">):</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">iou_thresholds</span> <span class="o">=</span> <span class="n">iou_thres</span><span class="o">.</span><span class="n">to_tensor</span><span class="p">()</span>
                
 
                            +        <span class="k">else</span><span class="p">:</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">iou_thresholds</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">iou_thres</span><span class="p">])</span>
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">map_str</span> <span class="o">=</span> <span class="s2">&quot;mAP&quot;</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_component_is_better</span> <span class="o">=</span> <span class="p">{</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;Precision</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;Recall</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;mAP</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;F1</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +        <span class="p">}</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">component_names</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">greater_component_is_better</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">components</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">component_names</span><span class="p">)</span>
                
 
                            +
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">post_prediction_callback</span> <span class="o">=</span> <span class="n">post_prediction_callback</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">is_distributed</span> <span class="o">=</span> <span class="n">super_gradients</span><span class="o">.</span><span class="n">is_distributed</span><span class="p">()</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">denormalize_targets</span> <span class="o">=</span> <span class="ow">not</span> <span class="n">normalize_targets</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">world_size</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">rank</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;matching_info&quot;</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="p">[],</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;matching_info</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="p">[],</span> <span class="n"
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">iou_thresholds</span> <span class="o">=</span> <span class="n">iou_thres</span><span class="o">.</span><span class="n">to_tensor</span><span class="p">()</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">recall_thresholds</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">101</span><span class="p">)</span> <span class="k">if</span> <span class="n">recall_thres</span> <span class="ow">is</span> <span class="kc">None</sp
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">score_threshold</span> <span class="o">=</span> <span class="n">score_thres</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">top_k_predictions</span> <span class="o">=</span> <span class="n">top_k_predictions</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">accumulate_on_cpu</span> <span class="o">=</span> <span class="n">accumulate_on_cpu</span>
                
 
                            -<div class="viewcode-block" id="DetectionMetrics.update"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.DetectionMetrics.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</s
                
 
                            -               <span class="n">inputs</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">,</span> <span class="n">crowd_targets</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="DetectionMetrics.update"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.DetectionMetrics.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><spa
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Apply NMS and match all the predictions and targets of a given batch, and update the metric state accordingly.</span>
                
@@ -173,27 +191,38 @@
 
                                     <span class="n">preds</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">post_prediction_callback</span><span class="p">(</span><span class="n">preds</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">new_matching_info</span> <span class="o">=</span> <span class="n">compute_detection_matching</span><span class="p">(</span>
                
 
                            -            <span class="n">preds</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">height</span><span class="p">,</span> <span class="n">width</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">iou_thresholds</span><span class="p">,</span> <span class="n">crowd_targets</span><span class="o">=</span><span class="n">crowd_targets</span><span class="p">,</span>
                
 
                            -            <span class="n">top_k</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">top_k_predictions</span><span class="p">,</span> <span class="n">denormalize_targets</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">denormalize_targets</span><span class="p">,</span>
                
 
                            -            <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">return_on_cpu</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">accumulate_on_cpu</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="n">accumulated_matching_info</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;matching_info&quot;</span><span class="p">)</span>
                
 
                            -        <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;matching_info&quot;</span><span class="p">,</span> <span class="n">accumulated_matching_info</span> <span class="o">+</span> <span class="n">new_matching_info</span><span class="p">)</span></div>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="DetectionMetrics.compute"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.DetectionMetrics.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Unio
                
 
                            +            <span class="n">preds</span><span class="p">,</span>
                
 
                            +            <span class="n">targets</span><span class="p">,</span>
                
 
                            +            <span class="n">height</span><span class="p">,</span>
                
 
                            +            <span class="n">width</span><span class="p">,</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">iou_thresholds</span><span class="p">,</span>
                
 
                            +            <span class="n">crowd_targets</span><span class="o">=</span><span class="n">crowd_targets</span><span class="p">,</span>
                
 
                            +            <span class="n">top_k</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">top_k_predictions</span><span class="p">,</span>
                
 
                            +            <span class="n">denormalize_targets</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">denormalize_targets</span><span class="p">,</span>
                
 
                            +            <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span>
                
 
                            +            <span class="n">return_on_cpu</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">accumulate_on_cpu</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">accumulated_matching_info</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;matching_info</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
                
 
                            +        <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;matching_info</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">accumulated_matching_info</span> <span class="o">+</span> <span class="n">new_ma
                
 
                            +
                
 
                            +<div class="viewcode-block" id="DetectionMetrics.compute"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.DetectionMetrics.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Union</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Compute the metrics for all the accumulated results.</span>
                
 
                            -<span class="sd">            :return: Metrics of interest</span>
                
 
                            +<span class="sd">        :return: Metrics of interest</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="n">mean_ap</span><span class="p">,</span> <span class="n">mean_precision</span><span class="p">,</span> <span class="n">mean_recall</span><span class="p">,</span> <span class="n">mean_f1</span> <span class="o">=</span> <span class="o">-</span><span class="mf">1.</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.</span><span class="p">,</span> <span class="o">-</span><span clas
                
 
                            -        <span class="n">accumulated_matching_info</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;matching_info&quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">mean_ap</span><span class="p">,</span> <span class="n">mean_precision</span><span class="p">,</span> <span class="n">mean_recall</span><span class="p">,</span> <span class="n">mean_f1</span> <span class="o">=</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">,</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">,</span> <span class="o">-</span><span c
                
 
                            +        <span class="n">accumulated_matching_info</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;matching_info</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">accumulated_matching_info</span><span class="p">):</span>
                
 
                                         <span class="n">matching_info_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span clas
                
 
                                         <span class="c1"># shape (n_class, nb_iou_thresh)</span>
                
 
                                         <span class="n">ap</span><span class="p">,</span> <span class="n">precision</span><span class="p">,</span> <span class="n">recall</span><span class="p">,</span> <span class="n">f1</span><span class="p">,</span> <span class="n">unique_classes</span> <span class="o">=</span> <span class="n">compute_detection_metrics</span><span class="p">(</span>
                
 
                            -                <span class="o">*</span><span class="n">matching_info_tensors</span><span class="p">,</span> <span class="n">recall_thresholds</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">recall_thresholds</span><span class="p">,</span> <span class="n">score_threshold</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">score_threshold</span><span class="p">,</span>
                
 
                            -                <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">accumulate_on_cpu</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
                
 
                            +                <span class="o">*</span><span class="n">matching_info_tensors</span><span class="p">,</span>
                
 
                            +                <span class="n">recall_thresholds</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">recall_thresholds</span><span class="p">,</span>
                
 
                            +                <span class="n">score_threshold</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">score_threshold</span><span class="p">,</span>
                
 
                            +                <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">accumulate_on_cpu</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span>
                
 
                            +            <span class="p">)</span>
                
 
                                         <span class="c1"># Precision, recall and f1 are computed for smallest IoU threshold (usually 0.5), averaged over classes</span>
                
 
                                         <span class="n">mean_precision</span><span class="p">,</span> <span class="n">mean_recall</span><span class="p">,</span> <span class="n">mean_f1</span> <span class="o">=</span> <span class="n">precision</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">mean</span><span class="p">(),</span> <span class="n">recall</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</
                
@@ -201,7 +230,12 @@
 
                                         <span class="c1"># MaP is averaged over IoU thresholds and over classes</span>
                
 
                                         <span class="n">mean_ap</span> <span class="o">=</span> <span class="n">ap</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>
                
 
                            -        <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;Precision&quot;</span><span class="p">:</span> <span class="n">mean_precision</span><span class="p">,</span> <span class="s2">&quot;Recall&quot;</span><span class="p">:</span> <span class="n">mean_recall</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">map_str</span><span class="p">:</span> <span class="n">mean_ap</span><span class="p">,</span> <span class="s2">
                
 
                            +        <span class="k">return</span> <span class="p">{</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;Precision</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="n">mean_precision</span><span class="p">,</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;Recall</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="n">mean_recall</span><span class="p">,</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;mAP</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="n">mean_ap</span><span class="p">,</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;F1</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">:</span> <span class="n">mean_f1</span><span class="p">,</span>
                
 
                            +        <span class="p">}</span></div>
                
 
                                 <span class="k">def</span> <span class="nf">_sync_dist</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dist_sync_fn</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">process_group</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
@@ -223,10 +257,83 @@
 
                                         <span class="n">torch</span><span class="o">.</span><span class="n">distributed</span><span class="o">.</span><span class="n">all_gather_object</span><span class="p">(</span><span class="n">gathered_state_dicts</span><span class="p">,</span> <span class="n">local_state_dict</span><span class="p">)</span>
                
 
                                         <span class="n">matching_info</span> <span class="o">=</span> <span class="p">[]</span>
                
 
                                         <span class="k">for</span> <span class="n">state_dict</span> <span class="ow">in</span> <span class="n">gathered_state_dicts</span><span class="p">:</span>
                
 
                            -                <span class="n">matching_info</span> <span class="o">+=</span> <span class="n">state_dict</span><span class="p">[</span><span class="s2">&quot;matching_info&quot;</span><span class="p">]</span>
                
 
                            +                <span class="n">matching_info</span> <span class="o">+=</span> <span class="n">state_dict</span><span class="p">[</span><span class="sa">f</span><span class="s2">&quot;matching_info</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">]</span>
                
 
                                         <span class="n">matching_info</span> <span class="o">=</span> <span class="n">tensor_container_to_device</span><span class="p">(</span><span class="n">matching_info</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">accumulate_on_cpu</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span
                
 
                            -            <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;matching_info&quot;</span><span class="p">,</span> <span class="n">matching_info</span><span class="p">)</span></div>
                
 
                            +            <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;matching_info</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_range_str</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">matching_info</span><span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_get_range_str</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="k">return</span> <span class="s2">&quot;@</span><span class="si">%.2f</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">iou_thresholds</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">if</span> <span class="ow">not</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">iou_thresh
                
 
                            +
                
 
                            +
                
 
                            +<div class="viewcode-block" id="DetectionMetrics_050"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.DetectionMetrics_050">[docs]</a><span class="k">class</span> <span class="nc">DetectionMetrics_050</span><span class="p">(</span><span class="n">DetectionMetrics</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">num_cls</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            +        <span class="n">post_prediction_callback</span><span class="p">:</span> <span class="n">DetectionPostPredictionCallback</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">normalize_targets</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">recall_thres</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">score_thres</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                
 
                            +        <span class="n">top_k_predictions</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">accumulate_on_cpu</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                            +
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
                
 
                            +            <span class="n">num_cls</span><span class="p">,</span>
                
 
                            +            <span class="n">post_prediction_callback</span><span class="p">,</span>
                
 
                            +            <span class="n">normalize_targets</span><span class="p">,</span>
                
 
                            +            <span class="n">IouThreshold</span><span class="o">.</span><span class="n">MAP_05</span><span class="p">,</span>
                
 
                            +            <span class="n">recall_thres</span><span class="p">,</span>
                
 
                            +            <span class="n">score_thres</span><span class="p">,</span>
                
 
                            +            <span class="n">top_k_predictions</span><span class="p">,</span>
                
 
                            +            <span class="n">dist_sync_on_step</span><span class="p">,</span>
                
 
                            +            <span class="n">accumulate_on_cpu</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span></div>
                
 
                            +
                
 
                            +
                
 
                            +<div class="viewcode-block" id="DetectionMetrics_075"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.DetectionMetrics_075">[docs]</a><span class="k">class</span> <span class="nc">DetectionMetrics_075</span><span class="p">(</span><span class="n">DetectionMetrics</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">num_cls</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            +        <span class="n">post_prediction_callback</span><span class="p">:</span> <span class="n">DetectionPostPredictionCallback</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">normalize_targets</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">recall_thres</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">score_thres</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                
 
                            +        <span class="n">top_k_predictions</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">accumulate_on_cpu</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                            +
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
                
 
                            +            <span class="n">num_cls</span><span class="p">,</span> <span class="n">post_prediction_callback</span><span class="p">,</span> <span class="n">normalize_targets</span><span class="p">,</span> <span class="mf">0.75</span><span class="p">,</span> <span class="n">recall_thres</span><span class="p">,</span> <span class="n">score_thres</span><span class="p">,</span> <span class="n">top_k_predictions</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="p">,<
                
 
                            +        <span class="p">)</span></div>
                
 
                            +
                
 
                            +
                
 
                            +<div class="viewcode-block" id="DetectionMetrics_050_095"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.DetectionMetrics_050_095">[docs]</a><span class="k">class</span> <span class="nc">DetectionMetrics_050_095</span><span class="p">(</span><span class="n">DetectionMetrics</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">num_cls</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            +        <span class="n">post_prediction_callback</span><span class="p">:</span> <span class="n">DetectionPostPredictionCallback</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">normalize_targets</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">recall_thres</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">score_thres</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                
 
                            +        <span class="n">top_k_predictions</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">accumulate_on_cpu</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                            +
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
                
 
                            +            <span class="n">num_cls</span><span class="p">,</span>
                
 
                            +            <span class="n">post_prediction_callback</span><span class="p">,</span>
                
 
                            +            <span class="n">normalize_targets</span><span class="p">,</span>
                
 
                            +            <span class="n">IouThreshold</span><span class="o">.</span><span class="n">MAP_05_TO_095</span><span class="p">,</span>
                
 
                            +            <span class="n">recall_thres</span><span class="p">,</span>
                
 
                            +            <span class="n">score_thres</span><span class="p">,</span>
                
 
                            +            <span class="n">top_k_predictions</span><span class="p">,</span>
                
 
                            +            <span class="n">dist_sync_on_step</span><span class="p">,</span>
                
 
                            +            <span class="n">accumulate_on_cpu</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span></div>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -256,4 +363,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.metrics.metric_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.metrics.metric_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.metrics.metric_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torchmetrics</span> <span class="kn">import</span> <span class="n">MetricCollection</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">AverageMeter</span>


<div class="viewcode-block" id="get_logging_values"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.metric_utils.get_logging_values">[docs]</a><span class="k">def</span> <span class="nf">get_logging_values</span><span class="p">(</span><span class="n">loss_loggings</span><span class="p">:</span> <span class="n">AverageMeter</span><span class="p">,</span> <span class="n">metrics</span><span class="p">:</span> <span class="n">MetricCollection</span><span class="p">,</span> <span class="n">criterion</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    @param loss_loggings: AverageMeter running average for the loss items</span>
<span class="sd">    @param metrics: MetricCollection object for running user specified metrics</span>
<span class="sd">    @param criterion the object loss_loggings average meter is monitoring, when set to None- only the metrics values are</span>
<span class="sd">    computed and returned.</span>

<span class="sd">    @return: tuple of the computed values</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">criterion</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">loss_loggingg_avg</span> <span class="o">=</span> <span class="n">loss_loggings</span><span class="o">.</span><span class="n">average</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">loss_loggingg_avg</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="n">loss_loggingg_avg</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">([</span><span class="n">loss_loggingg_avg</span><span class="p">])</span>
        <span class="n">logging_vals</span> <span class="o">=</span> <span class="n">loss_loggingg_avg</span> <span class="o">+</span> <span class="n">get_metrics_results_tuple</span><span class="p">(</span><span class="n">metrics</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">logging_vals</span> <span class="o">=</span> <span class="n">get_metrics_results_tuple</span><span class="p">(</span><span class="n">metrics</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">logging_vals</span></div>


<div class="viewcode-block" id="get_metrics_titles"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.metric_utils.get_metrics_titles">[docs]</a><span class="k">def</span> <span class="nf">get_metrics_titles</span><span class="p">(</span><span class="n">metrics_collection</span><span class="p">:</span> <span class="n">MetricCollection</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">    @param metrics_collection: MetricCollection object for running user specified metrics</span>
<span class="sd">    @return: list of all the names of the computed values list(str)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">titles</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">metric_name</span><span class="p">,</span> <span class="n">metric</span> <span class="ow">in</span> <span class="n">metrics_collection</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">metric_name</span> <span class="o">==</span> <span class="s2">&quot;additional_items&quot;</span><span class="p">:</span>
            <span class="k">continue</span>
        <span class="k">elif</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">metric</span><span class="p">,</span> <span class="s2">&quot;component_names&quot;</span><span class="p">):</span>
            <span class="n">titles</span> <span class="o">+=</span> <span class="n">metric</span><span class="o">.</span><span class="n">component_names</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">titles</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">metric_name</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">titles</span></div>


<div class="viewcode-block" id="get_metrics_results_tuple"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.metric_utils.get_metrics_results_tuple">[docs]</a><span class="k">def</span> <span class="nf">get_metrics_results_tuple</span><span class="p">(</span><span class="n">metrics_collection</span><span class="p">:</span> <span class="n">MetricCollection</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">    @param metrics_collection: metrics collection of the user specified metrics</span>
<span class="sd">    @type metrics_collection</span>
<span class="sd">    @return: tuple of metrics values</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">metrics_collection</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">results_tuple</span> <span class="o">=</span> <span class="p">()</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">results_tuple</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">flatten_metrics_dict</span><span class="p">(</span><span class="n">metrics_collection</span><span class="o">.</span><span class="n">compute</span><span class="p">())</span><span class="o">.</span><span class="n">values</span><span class="p">())</span>
    <span class="k">return</span> <span class="n">results_tuple</span></div>


<div class="viewcode-block" id="flatten_metrics_dict"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.metric_utils.flatten_metrics_dict">[docs]</a><span class="k">def</span> <span class="nf">flatten_metrics_dict</span><span class="p">(</span><span class="n">metrics_dict</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    :param metrics_dict - dictionary of metric values where values can also be dictionaries containing subvalues</span>
<span class="sd">    (in the case of compound metrics)</span>

<span class="sd">    @return: flattened dict of metric values i.e {metric1_name: metric1_value...}</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">flattened</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">metric_name</span><span class="p">,</span> <span class="n">metric_val</span> <span class="ow">in</span> <span class="n">metrics_dict</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">metric_name</span> <span class="o">==</span> <span class="s2">&quot;additional_items&quot;</span><span class="p">:</span>
            <span class="k">continue</span>
        <span class="c1"># COLLECT ALL OF THE COMPONENTS IN THE CASE OF COMPOUND METRICS</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">metric_val</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">sub_metric_name</span><span class="p">,</span> <span class="n">sub_metric_val</span> <span class="ow">in</span> <span class="n">metric_val</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                <span class="n">flattened</span><span class="p">[</span><span class="n">sub_metric_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">sub_metric_val</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">flattened</span><span class="p">[</span><span class="n">metric_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">metric_val</span>

    <span class="k">return</span> <span class="n">flattened</span></div>


<div class="viewcode-block" id="get_metrics_dict"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.metric_utils.get_metrics_dict">[docs]</a><span class="k">def</span> <span class="nf">get_metrics_dict</span><span class="p">(</span><span class="n">metrics_tuple</span><span class="p">,</span> <span class="n">metrics_collection</span><span class="p">,</span> <span class="n">loss_logging_item_names</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a dictionary with the epoch results as values and their names as keys.</span>
<span class="sd">    @param metrics_tuple: the result tuple</span>
<span class="sd">    @param metrics_collection: MetricsCollection</span>
<span class="sd">    @param loss_logging_item_names: loss component&#39;s names.</span>
<span class="sd">    @return: dict</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">keys</span> <span class="o">=</span> <span class="n">loss_logging_item_names</span> <span class="o">+</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="n">metrics_collection</span><span class="p">)</span>
    <span class="n">metrics_dict</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="nb">list</span><span class="p">(</span><span class="n">metrics_tuple</span><span class="p">)))</span>
    <span class="k">return</span> <span class="n">metrics_dict</span></div>


<div class="viewcode-block" id="get_train_loop_description_dict"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.metric_utils.get_train_loop_description_dict">[docs]</a><span class="k">def</span> <span class="nf">get_train_loop_description_dict</span><span class="p">(</span><span class="n">metrics_tuple</span><span class="p">,</span> <span class="n">metrics_collection</span><span class="p">,</span> <span class="n">loss_logging_item_names</span><span class="p">,</span> <span class="o">**</span><span class="n">log_items</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a dictionary with the epoch&#39;s logging items as values and their names as keys, with the purpose of</span>
<span class="sd">     passing it as a description to tqdm&#39;s progress bar.</span>

<span class="sd">    @param metrics_tuple: the result tuple</span>
<span class="sd">    @param metrics_collection: MetricsCollection</span>
<span class="sd">    @param loss_logging_item_names: loss component&#39;s names.</span>
<span class="sd">    @param log_items additional logging items to be rendered.</span>
<span class="sd">    @return: dict</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">log_items</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">get_metrics_dict</span><span class="p">(</span><span class="n">metrics_tuple</span><span class="p">,</span> <span class="n">metrics_collection</span><span class="p">,</span> <span class="n">loss_logging_item_names</span><span class="p">))</span>
    <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">log_items</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
            <span class="n">log_items</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">value</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">log_items</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.metrics.segmentation_metrics &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.metrics.segmentation_metrics &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -95,7 +97,7 @@
 
                             <span class="kn">from</span> <span class="nn">abc</span> <span class="kn">import</span> <span class="n">ABC</span><span class="p">,</span> <span class="n">abstractmethod</span>
                
 
                            -<div class="viewcode-block" id="batch_pix_accuracy"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.batch_pix_accuracy">[docs]</a><span class="k">def</span> <span class="nf">batch_pix_accuracy</span><span class="p">(</span><span class="n">predict</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">batch_pix_accuracy</span><span class="p">(</span><span class="n">predict</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Batch Pixel Accuracy</span>
                
 
                             <span class="sd">    Args:</span>
                
 
                             <span class="sd">        predict: input 4D tensor</span>
                
@@ -106,12 +108,11 @@
 
                                 <span class="n">target</span> <span class="o">=</span> <span class="n">target</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> <span class="o">+</span> <span class="mi">1</span>
                
 
                                 <span class="n">pixel_labeled</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">target</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">)</span>
                
 
                                 <span class="n">pixel_correct</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">((</span><span class="n">predict</span> <span class="o">==</span> <span class="n">target</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="n">target</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">))</span>
                
 
                            -    <span class="k">assert</span> <span class="n">pixel_correct</span> <span class="o">&lt;=</span> <span class="n">pixel_labeled</span><span class="p">,</span> \
                
 
                            -        <span class="s2">&quot;Correct area should be smaller than Labeled&quot;</span>
                
 
                            -    <span class="k">return</span> <span class="n">pixel_correct</span><span class="p">,</span> <span class="n">pixel_labeled</span></div>
                
 
                            +    <span class="k">assert</span> <span class="n">pixel_correct</span> <span class="o">&lt;=</span> <span class="n">pixel_labeled</span><span class="p">,</span> <span class="s2">&quot;Correct area should be smaller than Labeled&quot;</span>
                
 
                            +    <span class="k">return</span> <span class="n">pixel_correct</span><span class="p">,</span> <span class="n">pixel_labeled</span>
                
 
                            -<div class="viewcode-block" id="batch_intersection_union"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.batch_intersection_union">[docs]</a><span class="k">def</span> <span class="nf">batch_intersection_union</span><span class="p">(</span><span class="n">predict</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">nclass</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">batch_intersection_union</span><span class="p">(</span><span class="n">predict</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">nclass</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Batch Intersection of Union</span>
                
 
                             <span class="sd">    Args:</span>
                
 
                             <span class="sd">        predict: input 4D tensor</span>
                
@@ -132,13 +133,12 @@
 
                                 <span class="n">area_pred</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">predict</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">nbins</span><span class="p">,</span> <span class="nb">range</span><span class="o">=</span><span class="p">(</span><span class="n">mini</span><span class="p">,</s
                
 
                                 <span class="n">area_lab</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">nbins</span><span class="p">,</span> <span class="nb">range</span><span class="o">=</span><span class="p">(</span><span class="n">mini</span><span class="p">,</spa
                
 
                                 <span class="n">area_union</span> <span class="o">=</span> <span class="n">area_pred</span> <span class="o">+</span> <span class="n">area_lab</span> <span class="o">-</span> <span class="n">area_inter</span>
                
 
                            -    <span class="k">assert</span> <span class="p">(</span><span class="n">area_inter</span> <span class="o">&lt;=</span> <span class="n">area_union</span><span class="p">)</span><span class="o">.</span><span class="n">all</span><span class="p">(),</span> \
                
 
                            -        <span class="s2">&quot;Intersection area should be smaller than Union area&quot;</span>
                
 
                            -    <span class="k">return</span> <span class="n">area_inter</span><span class="p">,</span> <span class="n">area_union</span></div>
                
 
                            +    <span class="k">assert</span> <span class="p">(</span><span class="n">area_inter</span> <span class="o">&lt;=</span> <span class="n">area_union</span><span class="p">)</span><span class="o">.</span><span class="n">all</span><span class="p">(),</span> <span class="s2">&quot;Intersection area should be smaller than Union area&quot;</span>
                
 
                            +    <span class="k">return</span> <span class="n">area_inter</span><span class="p">,</span> <span class="n">area_union</span>
                
 
                             <span class="c1"># ref https://github.com/CSAILVision/sceneparsing/blob/master/evaluationCode/utils_eval.py</span>
                
 
                            -<div class="viewcode-block" id="pixel_accuracy"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.pixel_accuracy">[docs]</a><span class="k">def</span> <span class="nf">pixel_accuracy</span><span class="p">(</span><span class="n">im_pred</span><span class="p">,</span> <span class="n">im_lab</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">pixel_accuracy</span><span class="p">(</span><span class="n">im_pred</span><span class="p">,</span> <span class="n">im_lab</span><span class="p">):</span>
                
 
                                 <span class="n">im_pred</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">im_pred</span><span class="p">)</span>
                
 
                                 <span class="n">im_lab</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">im_lab</span><span class="p">)</span>
                
@@ -147,7 +147,7 @@
 
                                 <span class="n">pixel_labeled</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">im_lab</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">)</span>
                
 
                                 <span class="n">pixel_correct</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">((</span><span class="n">im_pred</span> <span class="o">==</span> <span class="n">im_lab</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="n">im_lab</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">))</span>
                
 
                                 <span class="c1"># pixel_accuracy = 1.0 * pixel_correct / pixel_labeled</span>
                
 
                            -    <span class="k">return</span> <span class="n">pixel_correct</span><span class="p">,</span> <span class="n">pixel_labeled</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">pixel_correct</span><span class="p">,</span> <span class="n">pixel_labeled</span>
                
 
                             <span class="k">def</span> <span class="nf">_dice_from_confmat</span><span class="p">(</span>
                
@@ -188,51 +188,48 @@
 
                                     <span class="n">scores</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span>
                
 
                                         <span class="p">[</span>
                
 
                                             <span class="n">scores</span><span class="p">[:</span><span class="n">ignore_index</span><span class="p">],</span>
                
 
                            -                <span class="n">scores</span><span class="p">[</span><span class="n">ignore_index</span> <span class="o">+</span> <span class="mi">1</span><span class="p">:],</span>
                
 
                            +                <span class="n">scores</span><span class="p">[</span><span class="n">ignore_index</span> <span class="o">+</span> <span class="mi">1</span> <span class="p">:],</span>
                
 
                                         <span class="p">]</span>
                
 
                                     <span class="p">)</span>
                
 
                                 <span class="k">return</span> <span class="n">reduce</span><span class="p">(</span><span class="n">scores</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="intersection_and_union"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.intersection_and_union">[docs]</a><span class="k">def</span> <span class="nf">intersection_and_union</span><span class="p">(</span><span class="n">im_pred</span><span class="p">,</span> <span class="n">im_lab</span><span class="p">,</span> <span class="n">num_class</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">intersection_and_union</span><span class="p">(</span><span class="n">im_pred</span><span class="p">,</span> <span class="n">im_lab</span><span class="p">,</span> <span class="n">num_class</span><span class="p">):</span>
                
 
                                 <span class="n">im_pred</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">im_pred</span><span class="p">)</span>
                
 
                                 <span class="n">im_lab</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">im_lab</span><span class="p">)</span>
                
 
                                 <span class="c1"># Remove classes from unlabeled pixels in gt image.</span>
                
 
                                 <span class="n">im_pred</span> <span class="o">=</span> <span class="n">im_pred</span> <span class="o">*</span> <span class="p">(</span><span class="n">im_lab</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">)</span>
                
 
                                 <span class="c1"># Compute area intersection:</span>
                
 
                                 <span class="n">intersection</span> <span class="o">=</span> <span class="n">im_pred</span> <span class="o">*</span> <span class="p">(</span><span class="n">im_pred</span> <span class="o">==</span> <span class="n">im_lab</span><span class="p">)</span>
                
 
                            -    <span class="n">area_inter</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">intersection</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span>
                
 
                            -                                 <span class="nb">range</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">))</span>
                
 
                            +    <span class="n">area_inter</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">intersection</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="nb">range</span><span class="o">=</span><span clas
                
 
                                 <span class="c1"># Compute area union:</span>
                
 
                            -    <span class="n">area_pred</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">im_pred</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span>
                
 
                            -                                <span class="nb">range</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">))</span>
                
 
                            -    <span class="n">area_lab</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">im_lab</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span>
                
 
                            -                               <span class="nb">range</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">))</span>
                
 
                            +    <span class="n">area_pred</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">im_pred</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="nb">range</span><span class="o">=</span><span class="p">
                
 
                            +    <span class="n">area_lab</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">histogram</span><span class="p">(</span><span class="n">im_lab</span><span class="p">,</span> <span class="n">bins</span><span class="o">=</span><span class="n">num_class</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="nb">range</span><span class="o">=</span><span class="p">(<
                
 
                                 <span class="n">area_union</span> <span class="o">=</span> <span class="n">area_pred</span> <span class="o">+</span> <span class="n">area_lab</span> <span class="o">-</span> <span class="n">area_inter</span>
                
 
                            -    <span class="k">return</span> <span class="n">area_inter</span><span class="p">,</span> <span class="n">area_union</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">area_inter</span><span class="p">,</span> <span class="n">area_union</span>
                
 
                            -<div class="viewcode-block" id="AbstractMetricsArgsPrepFn"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.AbstractMetricsArgsPrepFn">[docs]</a><span class="k">class</span> <span class="nc">AbstractMetricsArgsPrepFn</span><span class="p">(</span><span class="n">ABC</span><span class="p">):</span>
                
 
                            +<span class="k">class</span> <span class="nc">AbstractMetricsArgsPrepFn</span><span class="p">(</span><span class="n">ABC</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Abstract preprocess metrics arguments class.</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                                 <span class="nd">@abstractmethod</span>
                
 
                                 <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        All base classes must implement this function and return a tuple of torch tensors (predictions, target).</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">()</span></div>
                
 
                            +        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">()</span>
                
 
                            -<div class="viewcode-block" id="PreprocessSegmentationMetricsArgs"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.PreprocessSegmentationMetricsArgs">[docs]</a><span class="k">class</span> <span class="nc">PreprocessSegmentationMetricsArgs</span><span class="p">(</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="PreprocessSegmentationMetricsArgs"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.PreprocessSegmentationMetricsArgs">[docs]</a><span class="k">class</span> <span class="nc">PreprocessSegmentationMetricsArgs</span><span class="p">(</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Default segmentation inputs preprocess function before updating segmentation metrics, handles multiple inputs and</span>
                
 
                             <span class="sd">    apply normalizations.</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                            -                 <span class="n">apply_arg_max</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -                 <span class="n">apply_sigmoid</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">apply_arg_max</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">apply_sigmoid</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        :param apply_arg_max: Whether to apply argmax on predictions tensor.</span>
                
 
                             <span class="sd">        :param apply_sigmoid:  Whether to apply sigmoid on predictions tensor.</span>
                
@@ -253,18 +250,16 @@
 
                                     <span class="k">return</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span></div>
                
 
                            -<div class="viewcode-block" id="PixelAccuracy"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.PixelAccuracy">[docs]</a><span class="k">class</span> <span class="nc">PixelAccuracy</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                            -                 <span class="n">ignore_label</span><span class="o">=-</span><span class="mi">100</span><span class="p">,</span>
                
 
                            -                 <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                
 
                            -                 <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="PixelAccuracy"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.PixelAccuracy">[docs]</a><span class="k">class</span> <span class="nc">PixelAccuracy</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ignore_label</span><span class="o">=-</span><span class="mi">100</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span 
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">ignore_label</span> <span class="o">=</span> <span class="n">ignore_label</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;total_correct&quot;</span><span class="p">,</span> <span class="n">default</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="mf">0.</span><span class="p">),</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="s2">&quot;sum&quot;</spa
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;total_label&quot;</span><span class="p">,</span> <span class="n">default</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="mf">0.</span><span class="p">),</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="s2">&quot;sum&quot;</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;total_correct&quot;</span><span class="p">,</span> <span class="n">default</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="mf">0.0</span><span class="p">),</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="s2">&quot;sum&quot;</sp
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">add_state</span><span class="p">(</span><span class="s2">&quot;total_label&quot;</span><span class="p">,</span> <span class="n">default</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="mf">0.0</span><span class="p">),</span> <span class="n">dist_reduce_fx</span><span class="o">=</span><span class="s2">&quot;sum&quot;</span
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">metrics_args_prep_fn</span> <span class="o">=</span> <span class="n">metrics_args_prep_fn</span> <span class="ow">or</span> <span class="n">PreprocessSegmentationMetricsArgs</span><span class="p">(</span><span class="n">apply_arg_max</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="PixelAccuracy.update"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.PixelAccuracy.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <s
                
 
                            +<div class="viewcode-block" id="PixelAccuracy.update"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.PixelAccuracy.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span clas
                
 
                                     <span class="n">predict</span><span class="p">,</span> <span class="n">target</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">metrics_args_prep_fn</span><span class="p">(</span><span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
                
 
                                     <span class="n">labeled_mask</span> <span class="o">=</span> <span class="n">target</span><span class="o">.</span><span class="n">ne</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ignore_label</span><span class="p">)</span>
                
@@ -273,81 +268,117 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">total_correct</span> <span class="o">+=</span> <span class="n">pixel_correct</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">total_label</span> <span class="o">+=</span> <span class="n">pixel_labeled</span></div>
                
 
                            -<div class="viewcode-block" id="PixelAccuracy.compute"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.PixelAccuracy.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            -        <span class="n">_total_correct</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">total_correct</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s1">&#39;int64&#39;<
                
 
                            -        <span class="n">_total_label</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">total_label</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s1">&#39;int64&#39;</spa
                
 
                            +<div class="viewcode-block" id="PixelAccuracy.compute"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.PixelAccuracy.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="n">_total_correct</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">total_correct</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s2">&quot;int64&quot
                
 
                            +        <span class="n">_total_label</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">total_label</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s2">&quot;int64&quot;</s
                
 
                                     <span class="n">pix_acc</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float64</span><span class="p">(</span><span class="mf">1.0</span><span class="p">)</span> <span class="o">*</span> <span class="n">_total_correct</span> <span class="o">/</span> <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">spacing</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n"
                
 
                                     <span class="k">return</span> <span class="n">pix_acc</span></div></div>
                
 
                            -<div class="viewcode-block" id="IoU"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.IoU">[docs]</a><span class="k">class</span> <span class="nc">IoU</span><span class="p">(</span><span class="n">torchmetrics</span><span class="o">.</span><span class="n">JaccardIndex</span><span class="p">):</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                            -                 <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            -                 <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -                 <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;elementwise_mean&quot;</span><span class="p">,</span>
                
 
                            -                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            -                 <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_classes</span><span class="o">=</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,
                
 
                            -                         <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span> <span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">)</span>
                
 
                            +<div class="viewcode-block" id="IoU"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.IoU">[docs]</a><span class="k">class</span> <span class="nc">IoU</span><span class="p">(</span><span class="n">torchmetrics</span><span class="o">.</span><span class="n">JaccardIndex</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;elementwise_mean&quot;</span><span class="p">,</span>
                
 
                            +        <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            +        <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">num_classes</span> <span class="o">&lt;=</span> <span class="mi">1</span><span class="p">:</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;IoU class only for multi-class usage! For binary usage, please call </span><span class="si">{</span><span class="n">BinaryIOU</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_classes</span><span class="o">=</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">metrics_args_prep_fn</span> <span class="o">=</span> <span class="n">metrics_args_prep_fn</span> <span class="ow">or</span> <span class="n">PreprocessSegmentationMetricsArgs</span><span class="p">(</span><span class="n">apply_arg_max</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                            -<div class="viewcode-block" id="IoU.update"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.IoU.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor
                
 
                            +<div class="viewcode-block" id="IoU.update"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.IoU.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><
                
 
                                     <span class="n">preds</span><span class="p">,</span> <span class="n">target</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">metrics_args_prep_fn</span><span class="p">(</span><span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">preds</span><span class="o">=</span><span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="o">=</span><span class="n">target</span><span class="p">)</span></div></div>
                
 
                            -<div class="viewcode-block" id="Dice"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Dice">[docs]</a><span class="k">class</span> <span class="nc">Dice</span><span class="p">(</span><span class="n">torchmetrics</span><span class="o">.</span><span class="n">JaccardIndex</span><span class="p">):</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                            -                 <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            -                 <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -                 <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;elementwise_mean&quot;</span><span class="p">,</span>
                
 
                            -                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            -                 <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_classes</span><span class="o">=</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,
                
 
                            -                         <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span> <span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">)</span>
                
 
                            +<div class="viewcode-block" id="Dice"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Dice">[docs]</a><span class="k">class</span> <span class="nc">Dice</span><span class="p">(</span><span class="n">torchmetrics</span><span class="o">.</span><span class="n">JaccardIndex</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +        <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;elementwise_mean&quot;</span><span class="p">,</span>
                
 
                            +        <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            +        <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="n">num_classes</span> <span class="o">&lt;=</span> <span class="mi">1</span><span class="p">:</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Dice class only for multi-class usage! For binary usage, please call </span><span class="si">{</span><span class="n">BinaryDice</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_classes</span><span class="o">=</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">metrics_args_prep_fn</span> <span class="o">=</span> <span class="n">metrics_args_prep_fn</span> <span class="ow">or</span> <span class="n">PreprocessSegmentationMetricsArgs</span><span class="p">(</span><span class="n">apply_arg_max</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                            -<div class="viewcode-block" id="Dice.update"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Dice.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tens
                
 
                            +<div class="viewcode-block" id="Dice.update"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Dice.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span
                
 
                                     <span class="n">preds</span><span class="p">,</span> <span class="n">target</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">metrics_args_prep_fn</span><span class="p">(</span><span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
                
 
                                     <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">preds</span><span class="o">=</span><span class="n">preds</span><span class="p">,</span> <span class="n">target</span><span class="o">=</span><span class="n">target</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="Dice.compute"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.Dice.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="Dice.compute"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.Dice.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;Computes Dice coefficient&quot;&quot;&quot;</span>
                
 
                            -        <span class="k">return</span> <span class="n">_dice_from_confmat</span><span class="p">(</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">confmat</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_classes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">ignore_index</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">absent_score</span><span class="p">,</span> <span class="bp">self</span><span class="o">.
                
 
                            -        <span class="p">)</span></div></div>
                
 
                            +        <span class="k">return</span> <span class="n">_dice_from_confmat</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">confmat</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">num_classes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">ignore_index</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n
                
 
                            -<div class="viewcode-block" id="BinaryIOU"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.BinaryIOU">[docs]</a><span class="k">class</span> <span class="nc">BinaryIOU</span><span class="p">(</span><span class="n">IoU</span><span class="p">):</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                            -                 <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                
 
                            -                 <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            -                 <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="BinaryIOU"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.BinaryIOU">[docs]</a><span class="k">class</span> <span class="nc">BinaryIOU</span><span class="p">(</span><span class="n">IoU</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                
 
                            +        <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            +        <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                                     <span class="n">metrics_args_prep_fn</span> <span class="o">=</span> <span class="n">metrics_args_prep_fn</span> <span class="ow">or</span> <span class="n">PreprocessSegmentationMetricsArgs</span><span class="p">(</span><span class="n">apply_sigmoid</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_classes</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,</span>
                
 
                            -                         <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">,</span> <span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span> <span class="n">metrics_args_prep_fn</span><span class="o">=</span><span class="n">metrics_args_prep_fn</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">component_names</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;target_IOU&quot;</span><span class="p">,</span> <span class="s2">&quot;background_IOU&quot;</span><span class="p">,</span> <span class="s2">&quot;mean_IOU&quot;</span><span class="p">]</span>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="BinaryIOU.compute"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.BinaryIOU.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
                
 
                            +            <span class="n">num_classes</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
                
 
                            +            <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span>
                
 
                            +            <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,</span>
                
 
                            +            <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">,</span>
                
 
                            +            <span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span>
                
 
                            +            <span class="n">metrics_args_prep_fn</span><span class="o">=</span><span class="n">metrics_args_prep_fn</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_component_is_better</span> <span class="o">=</span> <span class="p">{</span>
                
 
                            +            <span class="s2">&quot;target_IOU&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;background_IOU&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;mean_IOU&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +        <span class="p">}</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">component_names</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">greater_component_is_better</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="BinaryIOU.compute"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.BinaryIOU.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="n">ious</span> <span class="o">=</span> <span class="nb">super</span><span class="p">(</span><span class="n">BinaryIOU</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">compute</span><span class="p">()</span>
                
 
                                     <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;target_IOU&quot;</span><span class="p">:</span> <span class="n">ious</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="s2">&quot;background_IOU&quot;</span><span class="p">:</span> <span class="n">ious</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="s2">&quot;mean_IOU&quot;</span><span class="p">:</span> <span class="n">io
                
 
                            -<div class="viewcode-block" id="BinaryDice"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.BinaryDice">[docs]</a><span class="k">class</span> <span class="nc">BinaryDice</span><span class="p">(</span><span class="n">Dice</span><span class="p">):</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                
 
                            -                 <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                
 
                            -                 <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            -                 <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="BinaryDice"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.BinaryDice">[docs]</a><span class="k">class</span> <span class="nc">BinaryDice</span><span class="p">(</span><span class="n">Dice</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                
 
                            +        <span class="n">ignore_index</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span>
                
 
                            +        <span class="n">metrics_args_prep_fn</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AbstractMetricsArgsPrepFn</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>
                
 
                                     <span class="n">metrics_args_prep_fn</span> <span class="o">=</span> <span class="n">metrics_args_prep_fn</span> <span class="ow">or</span> <span class="n">PreprocessSegmentationMetricsArgs</span><span class="p">(</span><span class="n">apply_sigmoid</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">num_classes</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,</span>
                
 
                            -                         <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">,</span> <span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span> <span class="n">metrics_args_prep_fn</span><span class="o">=</span><span class="n">metrics_args_prep_fn</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">component_names</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;target_Dice&quot;</span><span class="p">,</span> <span class="s2">&quot;background_Dice&quot;</span><span class="p">,</span> <span class="s2">&quot;mean_Dice&quot;</span><span class="p">]</span>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="BinaryDice.compute"><a class="viewcode-back" href="../../../../super_gradients.training.metrics.html#super_gradients.training.metrics.BinaryDice.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span>
                
 
                            +            <span class="n">num_classes</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
                
 
                            +            <span class="n">dist_sync_on_step</span><span class="o">=</span><span class="n">dist_sync_on_step</span><span class="p">,</span>
                
 
                            +            <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,</span>
                
 
                            +            <span class="n">reduction</span><span class="o">=</span><span class="s2">&quot;none&quot;</span><span class="p">,</span>
                
 
                            +            <span class="n">threshold</span><span class="o">=</span><span class="n">threshold</span><span class="p">,</span>
                
 
                            +            <span class="n">metrics_args_prep_fn</span><span class="o">=</span><span class="n">metrics_args_prep_fn</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_component_is_better</span> <span class="o">=</span> <span class="p">{</span>
                
 
                            +            <span class="s2">&quot;target_Dice&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;background_Dice&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +            <span class="s2">&quot;mean_Dice&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                
 
                            +        <span class="p">}</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">component_names</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">greater_component_is_better</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="BinaryDice.compute"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.metrics.BinaryDice.compute">[docs]</a>    <span class="k">def</span> <span class="nf">compute</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="n">dices</span> <span class="o">=</span> <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">compute</span><span class="p">()</span>
                
 
                                     <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;target_Dice&quot;</span><span class="p">:</span> <span class="n">dices</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="s2">&quot;background_Dice&quot;</span><span class="p">:</span> <span class="n">dices</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="s2">&quot;mean_Dice&quot;</span><span class="p">:</span> <span class="
                
 
                             </pre></div>
                
@@ -379,4 +410,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.models.sg_module &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.models.sg_module</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.models.sg_module</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>

<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">HpmStruct</span>


<div class="viewcode-block" id="SgModule"><a class="viewcode-back" href="../../../../super_gradients.training.models.html#super_gradients.training.models.sg_module.SgModule">[docs]</a><span class="k">class</span> <span class="nc">SgModule</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
<div class="viewcode-block" id="SgModule.initialize_param_groups"><a class="viewcode-back" href="../../../../super_gradients.training.models.html#super_gradients.training.models.sg_module.SgModule.initialize_param_groups">[docs]</a>    <span class="k">def</span> <span class="nf">initialize_param_groups</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">lr</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">training_params</span><span class="p">:</span> <span class="n">HpmStruct</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">        :return: list of dictionaries containing the key &#39;named_params&#39; with a list of named params</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="p">[{</span><span class="s2">&quot;named_params&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">named_parameters</span><span class="p">()}]</span></div>

<div class="viewcode-block" id="SgModule.update_param_groups"><a class="viewcode-back" href="../../../../super_gradients.training.models.html#super_gradients.training.models.sg_module.SgModule.update_param_groups">[docs]</a>    <span class="k">def</span> <span class="nf">update_param_groups</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span> <span class="n">param_groups</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">lr</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">epoch</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="nb">iter</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">training_params</span><span class="p">:</span> <span class="n">HpmStruct</span><span class="p">,</span> <span class="n">total_batch</span><span class="p">:</span> <span class="nb">int</span>
    <span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">        :param param_groups: list of dictionaries containing the params</span>
<span class="sd">        :return: list of dictionaries containing the params</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">for</span> <span class="n">param_group</span> <span class="ow">in</span> <span class="n">param_groups</span><span class="p">:</span>
            <span class="n">param_group</span><span class="p">[</span><span class="s2">&quot;lr&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">lr</span>
        <span class="k">return</span> <span class="n">param_groups</span></div>

<div class="viewcode-block" id="SgModule.get_include_attributes"><a class="viewcode-back" href="../../../../super_gradients.training.models.html#super_gradients.training.models.sg_module.SgModule.get_include_attributes">[docs]</a>    <span class="k">def</span> <span class="nf">get_include_attributes</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This function is used by the EMA. When updating the EMA model, some attributes of the main model (used in training)</span>
<span class="sd">        are updated to the EMA model along with the model weights.</span>
<span class="sd">        By default, all attributes are updated except for private attributes (starting with &#39;_&#39;)</span>
<span class="sd">        You can either set include_attributes or exclude_attributes. By returning a non empty list from this function,</span>
<span class="sd">        you override the default behaviour and only attributes named in this list will be updated.</span>
<span class="sd">        Note: This will also override the get_exclude_attributes list.</span>
<span class="sd">            :return: list of attributes to update from main model to EMA model</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="p">[]</span></div>

<div class="viewcode-block" id="SgModule.get_exclude_attributes"><a class="viewcode-back" href="../../../../super_gradients.training.models.html#super_gradients.training.models.sg_module.SgModule.get_exclude_attributes">[docs]</a>    <span class="k">def</span> <span class="nf">get_exclude_attributes</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This function is used by the EMA. When updating the EMA model, some attributes of the main model (used in training)</span>
<span class="sd">        are updated to the EMA model along with the model weights.</span>
<span class="sd">        By default, all attributes are updated except for private attributes (starting with &#39;_&#39;)</span>
<span class="sd">        You can either set include_attributes or exclude_attributes. By returning a non empty list from this function,</span>
<span class="sd">        you override the default behaviour and attributes named in this list will also be excluded from update.</span>
<span class="sd">        Note: if get_include_attributes is not empty, it will override this list.</span>
<span class="sd">            :return: list of attributes to not update from main model to EMA mode</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="p">[]</span></div>

<div class="viewcode-block" id="SgModule.prep_model_for_conversion"><a class="viewcode-back" href="../../../../super_gradients.training.models.html#super_gradients.training.models.sg_module.SgModule.prep_model_for_conversion">[docs]</a>    <span class="k">def</span> <span class="nf">prep_model_for_conversion</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Prepare the model to be converted to ONNX or other frameworks.</span>
<span class="sd">        Typically, this function will freeze the size of layers which is otherwise flexible, replace some modules</span>
<span class="sd">        with convertible substitutes and remove all auxiliary or training related parts.</span>
<span class="sd">        :param input_size: [H,W]</span>
<span class="sd">        &quot;&quot;&quot;</span></div>

<div class="viewcode-block" id="SgModule.replace_head"><a class="viewcode-back" href="../../../../super_gradients.training.models.html#super_gradients.training.models.sg_module.SgModule.replace_head">[docs]</a>    <span class="k">def</span> <span class="nf">replace_head</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Replace final layer for pretrained models. Since this varies between architectures, we leave it to the inheriting</span>
<span class="sd">        class to implement.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.params &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../" id="documentation_options" src="../../../_static/documentation_options.js"></script>
        <script src="../../../_static/jquery.js"></script>
        <script src="../../../_static/underscore.js"></script>
        <script src="../../../_static/doctools.js"></script>
    <script src="../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../genindex.html" />
    <link rel="search" title="Search" href="../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../welcome.html">SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../super_gradients.common.html">Common</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../super_gradients.training.html">Training</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.params</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.params</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">HpmStruct</span>

<span class="n">DEFAULT_TRAINING_PARAMS</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;lr_warmup_epochs&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span>
                           <span class="s2">&quot;cosine_final_lr_ratio&quot;</span><span class="p">:</span> <span class="mf">0.01</span><span class="p">,</span>
                           <span class="s2">&quot;optimizer&quot;</span><span class="p">:</span> <span class="s2">&quot;SGD&quot;</span><span class="p">,</span>
                           <span class="s2">&quot;criterion_params&quot;</span><span class="p">:</span> <span class="p">{},</span>
                           <span class="s2">&quot;ema&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                           <span class="s2">&quot;batch_accumulate&quot;</span><span class="p">:</span> <span class="mi">1</span><span class="p">,</span>  <span class="c1"># number of batches to accumulate before every backward pass</span>
                           <span class="s2">&quot;ema_params&quot;</span><span class="p">:</span> <span class="p">{},</span>
                           <span class="s2">&quot;zero_weight_decay_on_bias_and_bn&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                           <span class="s2">&quot;load_opt_params&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                           <span class="s2">&quot;run_validation_freq&quot;</span><span class="p">:</span> <span class="mi">1</span><span class="p">,</span>
                           <span class="s2">&quot;save_model&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                           <span class="s2">&quot;metric_to_watch&quot;</span><span class="p">:</span> <span class="s2">&quot;Accuracy&quot;</span><span class="p">,</span>
                           <span class="s2">&quot;launch_tensorboard&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                           <span class="s2">&quot;tb_files_user_prompt&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>  <span class="c1"># Asks User for Tensorboard Deletion Prompt</span>
                           <span class="s2">&quot;silent_mode&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>  <span class="c1"># Silents the Print outs</span>
                           <span class="s2">&quot;mixed_precision&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                           <span class="s2">&quot;tensorboard_port&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
                           <span class="s2">&quot;save_ckpt_epoch_list&quot;</span><span class="p">:</span> <span class="p">[],</span>  <span class="c1"># indices where the ckpt will save automatically</span>
                           <span class="s2">&quot;average_best_models&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                           <span class="s2">&quot;dataset_statistics&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>  <span class="c1"># add a dataset statistical analysis and sample images to tensorboard</span>
                           <span class="s2">&quot;save_tensorboard_to_s3&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                           <span class="s2">&quot;lr_schedule_function&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
                           <span class="s2">&quot;train_metrics_list&quot;</span><span class="p">:</span> <span class="p">[],</span>
                           <span class="s2">&quot;valid_metrics_list&quot;</span><span class="p">:</span> <span class="p">[],</span>
                           <span class="s2">&quot;loss_logging_items_names&quot;</span><span class="p">:</span> <span class="p">[</span><span class="s2">&quot;Loss&quot;</span><span class="p">],</span>
                           <span class="s2">&quot;greater_metric_to_watch_is_better&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                           <span class="s2">&quot;precise_bn&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                           <span class="s2">&quot;precise_bn_batch_size&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
                           <span class="s2">&quot;seed&quot;</span><span class="p">:</span> <span class="mi">42</span><span class="p">,</span>
                           <span class="s2">&quot;lr_mode&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
                           <span class="s2">&quot;phase_callbacks&quot;</span><span class="p">:</span> <span class="p">[],</span>
                           <span class="s2">&quot;log_installed_packages&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
                           <span class="s2">&quot;save_full_train_log&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                           <span class="s2">&quot;sg_logger&quot;</span><span class="p">:</span> <span class="s2">&quot;base_sg_logger&quot;</span><span class="p">,</span>
                           <span class="s2">&quot;sg_logger_params&quot;</span><span class="p">:</span>
                               <span class="p">{</span><span class="s2">&quot;tb_files_user_prompt&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>  <span class="c1"># Asks User for Tensorboard Deletion Prompt</span>
                                <span class="s2">&quot;project_name&quot;</span><span class="p">:</span> <span class="s2">&quot;&quot;</span><span class="p">,</span>
                                <span class="s2">&quot;launch_tensorboard&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
                                <span class="s2">&quot;tensorboard_port&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
                                <span class="s2">&quot;save_checkpoints_remote&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>  <span class="c1"># upload checkpoint files to s3</span>
                                <span class="s2">&quot;save_tensorboard_remote&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>  <span class="c1"># upload tensorboard files to s3</span>
                                <span class="s2">&quot;save_logs_remote&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">},</span>  <span class="c1"># upload log files to s3</span>
                           <span class="s2">&quot;warmup_mode&quot;</span><span class="p">:</span> <span class="s2">&quot;linear_step&quot;</span>
                           <span class="p">}</span>

<span class="n">DEFAULT_OPTIMIZER_PARAMS_SGD</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;weight_decay&quot;</span><span class="p">:</span> <span class="mf">1e-4</span><span class="p">,</span> <span class="s2">&quot;momentum&quot;</span><span class="p">:</span> <span class="mf">0.9</span><span class="p">}</span>

<span class="n">DEFAULT_OPTIMIZER_PARAMS_ADAM</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;weight_decay&quot;</span><span class="p">:</span> <span class="mf">1e-4</span><span class="p">}</span>

<span class="n">DEFAULT_OPTIMIZER_PARAMS_RMSPROP</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;weight_decay&quot;</span><span class="p">:</span> <span class="mf">1e-4</span><span class="p">,</span> <span class="s2">&quot;momentum&quot;</span><span class="p">:</span> <span class="mf">0.9</span><span class="p">}</span>

<span class="n">DEFAULT_OPTIMIZER_PARAMS_RMSPROPTF</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;weight_decay&quot;</span><span class="p">:</span> <span class="mf">1e-4</span><span class="p">,</span> <span class="s2">&quot;momentum&quot;</span><span class="p">:</span> <span class="mf">0.9</span><span class="p">}</span>

<span class="n">TRAINING_PARAM_SCHEMA</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;type&quot;</span><span class="p">:</span> <span class="s2">&quot;object&quot;</span><span class="p">,</span>
                         <span class="s2">&quot;properties&quot;</span><span class="p">:</span> <span class="p">{</span>
                             <span class="s2">&quot;max_epochs&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;type&quot;</span><span class="p">:</span> <span class="s2">&quot;number&quot;</span><span class="p">,</span> <span class="s2">&quot;minimum&quot;</span><span class="p">:</span> <span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;maximum&quot;</span><span class="p">:</span> <span class="mi">800</span><span class="p">},</span>

                             <span class="c1"># FIXME: CHECK THE IMPORTANCE OF THE COMMENTED SCHEMA- AS IT CAUSES HYDRA USE TO CRASH</span>

                             <span class="c1"># &quot;lr_updates&quot;: {&quot;type&quot;: &quot;array&quot;, &quot;minItems&quot;: 1},</span>
                             <span class="s2">&quot;lr_decay_factor&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;type&quot;</span><span class="p">:</span> <span class="s2">&quot;number&quot;</span><span class="p">,</span> <span class="s2">&quot;minimum&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;maximum&quot;</span><span class="p">:</span> <span class="mi">1</span><span class="p">},</span>
                             <span class="s2">&quot;lr_warmup_epochs&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;type&quot;</span><span class="p">:</span> <span class="s2">&quot;number&quot;</span><span class="p">,</span> <span class="s2">&quot;minimum&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;maximum&quot;</span><span class="p">:</span> <span class="mi">10</span><span class="p">},</span>
                             <span class="s2">&quot;initial_lr&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;type&quot;</span><span class="p">:</span> <span class="s2">&quot;number&quot;</span><span class="p">,</span> <span class="s2">&quot;exclusiveMinimum&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span> <span class="s2">&quot;maximum&quot;</span><span class="p">:</span> <span class="mi">10</span><span class="p">}</span>
                         <span class="p">},</span>
                         <span class="s2">&quot;if&quot;</span><span class="p">:</span> <span class="p">{</span>
                             <span class="s2">&quot;properties&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;lr_mode&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;const&quot;</span><span class="p">:</span> <span class="s2">&quot;step&quot;</span><span class="p">}}</span>
                         <span class="p">},</span>
                         <span class="s2">&quot;then&quot;</span><span class="p">:</span> <span class="p">{</span>
                             <span class="s2">&quot;required&quot;</span><span class="p">:</span> <span class="p">[</span><span class="s2">&quot;lr_updates&quot;</span><span class="p">,</span> <span class="s2">&quot;lr_decay_factor&quot;</span><span class="p">]</span>
                         <span class="p">},</span>
                         <span class="s2">&quot;required&quot;</span><span class="p">:</span> <span class="p">[</span><span class="s2">&quot;max_epochs&quot;</span><span class="p">,</span> <span class="s2">&quot;lr_mode&quot;</span><span class="p">,</span> <span class="s2">&quot;initial_lr&quot;</span><span class="p">,</span> <span class="s2">&quot;loss&quot;</span><span class="p">]</span>
                         <span class="p">}</span>


<div class="viewcode-block" id="TrainingParams"><a class="viewcode-back" href="../../../super_gradients.training.html#super_gradients.training.params.TrainingParams">[docs]</a><span class="k">class</span> <span class="nc">TrainingParams</span><span class="p">(</span><span class="n">HpmStruct</span><span class="p">):</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">entries</span><span class="p">):</span>
        <span class="c1"># WE initialize by the default training params, overridden by the provided params</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="o">**</span><span class="n">DEFAULT_TRAINING_PARAMS</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">set_schema</span><span class="p">(</span><span class="n">TRAINING_PARAM_SCHEMA</span><span class="p">)</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">entries</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">override</span><span class="p">(</span><span class="o">**</span><span class="n">entries</span><span class="p">)</span>

<div class="viewcode-block" id="TrainingParams.override"><a class="viewcode-back" href="../../../super_gradients.training.html#super_gradients.training.params.TrainingParams.override">[docs]</a>    <span class="k">def</span> <span class="nf">override</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">entries</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">override</span><span class="p">(</span><span class="o">**</span><span class="n">entries</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">validate</span><span class="p">()</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.sg_model.sg_model &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.sg_trainer.sg_trainer &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -76,7 +78,7 @@
 
                               <ul class="wy-breadcrumbs">
                
 
                                   <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
                
 
                                       <li><a href="../../../index.html">Module code</a> &raquo;</li>
                
 
                            -      <li>super_gradients.training.sg_model.sg_model</li>
                
 
                            +      <li>super_gradients.training.sg_trainer.sg_trainer</li>
                
 
                                   <li class="wy-breadcrumbs-aside">
                
 
                                   </li>
                
 
                               </ul>
                
@@ -85,74 +87,99 @@
 
                                       <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
                
 
                                        <div itemprop="articleBody">
                
 
                            -  <h1>Source code for super_gradients.training.sg_model.sg_model</h1><div class="highlight"><pre>
                
 
                            +  <h1>Source code for super_gradients.training.sg_trainer.sg_trainer</h1><div class="highlight"><pre>
                
 
                             <span></span><span class="kn">import</span> <span class="nn">inspect</span>
                
 
                             <span class="kn">import</span> <span class="nn">os</span>
                
 
                            -<span class="kn">import</span> <span class="nn">sys</span>
                
 
                             <span class="kn">from</span> <span class="nn">copy</span> <span class="kn">import</span> <span class="n">deepcopy</span>
                
 
                            -<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Mapping</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Any</span>
                
 
                            +<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Mapping</span><span class="p">,</span> <span class="n">Dict</span>
                
 
                            +<span class="kn">from</span> <span class="nn">pathlib</span> <span class="kn">import</span> <span class="n">Path</span>
                
 
                             <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
                
 
                            -<span class="kn">import</span> <span class="nn">pkg_resources</span>
                
 
                             <span class="kn">import</span> <span class="nn">torch</span>
                
 
                            -<span class="kn">import</span> <span class="nn">torchvision.transforms</span> <span class="k">as</span> <span class="nn">transforms</span>
                
 
                            -<span class="kn">from</span> <span class="nn">deprecated</span> <span class="kn">import</span> <span class="n">deprecated</span>
                
 
                            +<span class="kn">import</span> <span class="nn">hydra</span>
                
 
                            +<span class="kn">from</span> <span class="nn">omegaconf</span> <span class="kn">import</span> <span class="n">DictConfig</span>
                
 
                             <span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>
                
 
                            -<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">DataLoader</span><span class="p">,</span> <span class="n">DistributedSampler</span>
                
 
                            +<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">DataLoader</span><span class="p">,</span> <span class="n">SequentialSampler</span>
                
 
                             <span class="kn">from</span> <span class="nn">torch.cuda.amp</span> <span class="kn">import</span> <span class="n">GradScaler</span><span class="p">,</span> <span class="n">autocast</span>
                
 
                             <span class="kn">from</span> <span class="nn">torchmetrics</span> <span class="kn">import</span> <span class="n">MetricCollection</span>
                
 
                             <span class="kn">from</span> <span class="nn">tqdm</span> <span class="kn">import</span> <span class="n">tqdm</span>
                
 
                             <span class="kn">from</span> <span class="nn">piptools.scripts.sync</span> <span class="kn">import</span> <span class="n">_get_installed_distributions</span>
                
 
                            +<span class="kn">from</span> <span class="nn">torch.utils.data.distributed</span> <span class="kn">import</span> <span class="n">DistributedSampler</span>
                
 
                            +
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.factories.type_factory</span> <span class="kn">import</span> <span class="n">TypeFactory</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.samplers</span> <span class="kn">import</span> <span class="n">InfiniteSampler</span><span class="p">,</span> <span class="n">RepeatAugSampler</span>
                
 
                            +
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.factories.callbacks_factory</span> <span class="kn">import</span> <span class="n">CallbacksFactory</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.data_types.enum</span> <span class="kn">import</span> <span class="n">MultiGPUMode</span><span class="p">,</span> <span class="n">StrictLoad</span><span class="p">,</span> <span class="n">EvaluationType</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.models.all_architectures</span> <span class="kn">import</span> <span class="n">ARCHITECTURES</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.decorators.factory_decorator</span> <span class="kn">import</span> <span class="n">resolve_param</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.environment</span> <span class="kn">import</span> <span class="n">env_helpers</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.common.factories.datasets_factory</span> <span class="kn">import</span> <span class="n">DatasetsFactory</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span><span class="p">,</span> <span class="n">mute_current_process</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.factories.list_factory</span> <span class="kn">import</span> <span class="n">ListFactory</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.factories.losses_factory</span> <span class="kn">import</span> <span class="n">LossesFactory</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.factories.metrics_factory</span> <span class="kn">import</span> <span class="n">MetricsFactory</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.sg_loggers</span> <span class="kn">import</span> <span class="n">SG_LOGGERS</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.sg_loggers.abstract_sg_logger</span> <span class="kn">import</span> <span class="n">AbstractSGLogger</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.sg_loggers.base_sg_logger</span> <span class="kn">import</span> <span class="n">BaseSGLogger</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span><span class="p">,</span> <span class="n">models</span><span class="p">,</span> <span class="n">dataloaders</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.models</span> <span class="kn">import</span> <span class="n">SgModule</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.pretrained_models</span> <span class="kn">import</span> <span class="n">PRETRAINED_NUM_CLASSES</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">sg_model_utils</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils.quantization_utils</span> <span class="kn">import</span> <span class="n">QATCallback</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils.sg_model_utils</span> <span class="kn">import</span> <span class="n">MonitoredValue</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">metrics</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.sg_model_exceptions</span> <span class="kn">import</span> <span class="n">UnsupportedOptimizerFormat</span><span class="p">,</span> \
                
 
                            -    <span class="n">IllegalDataloaderInitialization</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.datasets</span> <span class="kn">import</span> <span class="n">DatasetInterface</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">sg_trainer_utils</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils.sg_trainer_utils</span> <span class="kn">import</span> <span class="n">MonitoredValue</span><span class="p">,</span> <span class="n">parse_args</span><span class="p">,</span> <span class="n">log_main_training_params</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.sg_trainer_exceptions</span> <span class="kn">import</span> <span class="n">UnsupportedOptimizerFormat</span><span class="p">,</span> <span class="n">GPUModeNotSetupError</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.losses</span> <span class="kn">import</span> <span class="n">LOSSES</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.metrics.metric_utils</span> <span class="kn">import</span> <span class="n">get_metrics_titles</span><span class="p">,</span> <span class="n">get_metrics_results_tuple</span><span class="p">,</span> \
                
 
                            -    <span class="n">get_logging_values</span><span class="p">,</span> \
                
 
                            -    <span class="n">get_metrics_dict</span><span class="p">,</span> <span class="n">get_train_loop_description_dict</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.metrics.metric_utils</span> <span class="kn">import</span> <span class="p">(</span>
                
 
                            +    <span class="n">get_metrics_titles</span><span class="p">,</span>
                
 
                            +    <span class="n">get_metrics_results_tuple</span><span class="p">,</span>
                
 
                            +    <span class="n">get_logging_values</span><span class="p">,</span>
                
 
                            +    <span class="n">get_metrics_dict</span><span class="p">,</span>
                
 
                            +    <span class="n">get_train_loop_description_dict</span><span class="p">,</span>
                
 
                            +<span class="p">)</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.params</span> <span class="kn">import</span> <span class="n">TrainingParams</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">DetectionPostPredictionCallback</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils.distributed_training_utils</span> <span class="kn">import</span> <span class="n">MultiGPUModeAutocastWrapper</span><span class="p">,</span> \
                
 
                            -    <span class="n">reduce_results_tuple_for_ddp</span><span class="p">,</span> <span class="n">compute_precise_bn_stats</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils.distributed_training_utils</span> <span class="kn">import</span> <span class="p">(</span>
                
 
                            +    <span class="n">MultiGPUModeAutocastWrapper</span><span class="p">,</span>
                
 
                            +    <span class="n">reduce_results_tuple_for_ddp</span><span class="p">,</span>
                
 
                            +    <span class="n">compute_precise_bn_stats</span><span class="p">,</span>
                
 
                            +    <span class="n">setup_device</span><span class="p">,</span>
                
 
                            +    <span class="n">require_gpu_setup</span><span class="p">,</span>
                
 
                            +    <span class="n">get_gpu_mem_utilization</span><span class="p">,</span>
                
 
                            +    <span class="n">get_world_size</span><span class="p">,</span>
                
 
                            +    <span class="n">get_local_rank</span><span class="p">,</span>
                
 
                            +    <span class="n">wait_for_the_master</span><span class="p">,</span>
                
 
                            +<span class="p">)</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.ema</span> <span class="kn">import</span> <span class="n">ModelEMA</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.optimizer_utils</span> <span class="kn">import</span> <span class="n">build_optimizer</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils.weight_averaging_utils</span> <span class="kn">import</span> <span class="n">ModelWeightAveraging</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.metrics</span> <span class="kn">import</span> <span class="n">Accuracy</span><span class="p">,</span> <span class="n">Top5</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">random_seed</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils.checkpoint_utils</span> <span class="kn">import</span> <span class="n">get_ckpt_local_path</span><span class="p">,</span> <span class="n">read_ckpt_state_dict</span><span class="p">,</span> \
                
 
                            -    <span class="n">load_checkpoint_to_model</span><span class="p">,</span> <span class="n">load_pretrained_weights</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils.checkpoint_utils</span> <span class="kn">import</span> <span class="p">(</span>
                
 
                            +    <span class="n">get_ckpt_local_path</span><span class="p">,</span>
                
 
                            +    <span class="n">read_ckpt_state_dict</span><span class="p">,</span>
                
 
                            +    <span class="n">load_checkpoint_to_model</span><span class="p">,</span>
                
 
                            +    <span class="n">load_pretrained_weights</span><span class="p">,</span>
                
 
                            +    <span class="n">get_checkpoints_dir_path</span><span class="p">,</span>
                
 
                            +<span class="p">)</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.datasets.datasets_utils</span> <span class="kn">import</span> <span class="n">DatasetStatisticsTensorboardLogger</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.utils.callbacks</span> <span class="kn">import</span> <span class="n">CallbackHandler</span><span class="p">,</span> <span class="n">Phase</span><span class="p">,</span> <span class="n">LR_SCHEDULERS_CLS_DICT</span><span class="p">,</span> <span class="n">PhaseContext</span><span class="p">,</span> \
                
 
                            -    <span class="n">MetricsUpdateCallback</span><span class="p">,</span> <span class="n">LR_WARMUP_CLS_DICT</span><span class="p">,</span> <span class="n">ContextSgMethods</span><span class="p">,</span> <span class="n">LRCallbackBase</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils.callbacks</span> <span class="kn">import</span> <span class="p">(</span>
                
 
                            +    <span class="n">CallbackHandler</span><span class="p">,</span>
                
 
                            +    <span class="n">Phase</span><span class="p">,</span>
                
 
                            +    <span class="n">LR_SCHEDULERS_CLS_DICT</span><span class="p">,</span>
                
 
                            +    <span class="n">PhaseContext</span><span class="p">,</span>
                
 
                            +    <span class="n">MetricsUpdateCallback</span><span class="p">,</span>
                
 
                            +    <span class="n">LR_WARMUP_CLS_DICT</span><span class="p">,</span>
                
 
                            +    <span class="n">ContextSgMethods</span><span class="p">,</span>
                
 
                            +    <span class="n">LRCallbackBase</span><span class="p">,</span>
                
 
                            +<span class="p">)</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common.environment</span> <span class="kn">import</span> <span class="n">environment_config</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">HpmStruct</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.training.datasets.samplers.infinite_sampler</span> <span class="kn">import</span> <span class="n">InfiniteSampler</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.training.utils.hydra_utils</span> <span class="kn">import</span> <span class="n">load_experiment_cfg</span><span class="p">,</span> <span class="n">add_params_to_cfg</span>
                
 
                            +<span class="kn">from</span> <span class="nn">omegaconf</span> <span class="kn">import</span> <span class="n">OmegaConf</span>
                
 
                            -<span class="kn">from</span> <span class="nn">super_gradients.common</span> <span class="kn">import</span> <span class="n">StrictLoad</span><span class="p">,</span> <span class="n">MultiGPUMode</span><span class="p">,</span> <span class="n">EvaluationType</span>
                
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="SgModel"><a class="viewcode-back" href="../../../../super_gradients.training.sg_model.html#super_gradients.training.SgModel">[docs]</a><span class="k">class</span> <span class="nc">SgModel</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="Trainer"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.Trainer">[docs]</a><span class="k">class</span> <span class="nc">Trainer</span><span class="p">:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    SuperGradient Model - Base Class for Sg Models</span>
                
@@ -168,30 +195,17 @@
 
                             <span class="sd">        returns the test loss, accuracy and runtime</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">device</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">multi_gpu</span><span class="p">:</span> <span class="n">Union</span
                
 
                            -                 <span class="n">model_checkpoints_location</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;local&#39;</span><span class="p">,</span>
                
 
                            -                 <span class="n">overwrite_local_checkpoint</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> <span class="n">ckpt_name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;ckpt_latest.pth&#39;</span><span class="p">,</span>
                
 
                            -                 <span class="n">post_prediction_callback</span><span class="p">:</span> <span class="n">DetectionPostPredictionCallback</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                 <span class="n">train_loader</span><span class="p">:</span> <span class="n">DataLoader</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">valid_loader</span><span class="p">:</span> <span class="n">DataLoader</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">test_loader</span><span class="p">:</span> <span class="n">DataLoader</span> <span class="o">=</span> <span class="kc">Non
                
 
                            -                 <span class="n">classes</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">device</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">multi_gpu</span><span class="p">:</span> <span class="n">Union</span
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        :param experiment_name:                      Used for logging and loading purposes</span>
                
 
                             <span class="sd">        :param device:                          If equal to &#39;cpu&#39; runs on the CPU otherwise on GPU</span>
                
 
                             <span class="sd">        :param multi_gpu:                       If True, runs on all available devices</span>
                
 
                            -<span class="sd">        :param model_checkpoints_location:      If set to &#39;s3&#39; saves the Checkpoints in AWS S3</span>
                
 
                             <span class="sd">                                                otherwise saves the Checkpoints Locally</span>
                
 
                            -<span class="sd">        :param overwrite_local_checkpoint:      If set to False keeps the current local checkpoint when importing</span>
                
 
                             <span class="sd">                                                checkpoint from cloud service, otherwise overwrites the local checkpoints file</span>
                
 
                            -<span class="sd">        :param ckpt_name:                       The Checkpoint to Load</span>
                
 
                             <span class="sd">        :param ckpt_root_dir:                   Local root directory path where all experiment logging directories will</span>
                
 
                             <span class="sd">                                                reside. When none is give, it is assumed that</span>
                
 
                             <span class="sd">                                                pkg_resources.resource_filename(&#39;checkpoints&#39;, &quot;&quot;) exists and will be used.</span>
                
 
                            -<span class="sd">        :param train_loader:                    Training set Dataloader instead of using DatasetInterface, must pass &quot;valid_loader&quot;</span>
                
 
                            -<span class="sd">                                                and &quot;classes&quot; along with it</span>
                
 
                            -<span class="sd">        :param valid_loader:                    Validation set Dataloader</span>
                
 
                            -<span class="sd">        :param test_loader:                     Test set Dataloader</span>
                
 
                            -<span class="sd">        :param classes:                         List of class labels</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="c1"># SET THE EMPTY PROPERTIES</span>
                
@@ -201,7 +215,6 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">ema_model</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">update_param_groups</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">post_prediction_callback</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">criterion</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">scaler</span> <span class="o">=</span> <span class="kc">None</span>
                
@@ -217,157 +230,189 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">ddp_silent_mode</span> <span class="o">=</span> <span class="kc">False</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">source_ckpt_folder_name</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">model_weight_averaging</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">average_model_checkpoint_filename</span> <span class="o">=</span> <span class="s1">&#39;average_model.pth&#39;</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">average_model_checkpoint_filename</span> <span class="o">=</span> <span class="s2">&quot;average_model.pth&quot;</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">start_epoch</span> <span class="o">=</span> <span class="mi">0</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">best_metric</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">inf</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">external_checkpoint_path</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">strict_load</span> <span class="o">=</span> <span class="n">StrictLoad</span><span class="o">.</span><span class="n">ON</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">load_ema_as_net</span> <span class="o">=</span> <span class="kc">False</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_best_name</span> <span class="o">=</span> <span class="s1">&#39;ckpt_best.pth&#39;</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_best_name</span> <span class="o">=</span> <span class="s2">&quot;ckpt_best.pth&quot;</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">enable_qat</span> <span class="o">=</span> <span class="kc">False</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">qat_params</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">_infinite_train_loader</span> <span class="o">=</span> <span class="kc">False</span>
                
 
                            -
                
 
                            -        <span class="c1"># DETERMINE THE LOCATION OF THE LOSS AND ACCURACY IN THE RESULTS TUPLE OUTPUTED BY THE TEST</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">loss_idx_in_results_tuple</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">acc_idx_in_results_tuple</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_first_backward</span> <span class="o">=</span> <span class="kc">True</span>
                
 
                                     <span class="c1"># METRICS</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">greater_metric_to_watch_is_better</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">metric_to_watch</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_train_metrics_is_better</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="p">{}</span>  <span class="c1"># For each metric, indicates if greater is better</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">greater_valid_metrics_is_better</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                                     <span class="c1"># SETTING THE PROPERTIES FROM THE CONSTRUCTOR</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">experiment_name</span> <span class="o">=</span> <span class="n">experiment_name</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_name</span> <span class="o">=</span> <span class="n">ckpt_name</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">overwrite_local_checkpoint</span> <span class="o">=</span> <span class="n">overwrite_local_checkpoint</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">model_checkpoints_location</span> <span class="o">=</span> <span class="n">model_checkpoints_location</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">_set_dataset_properties</span><span class="p">(</span><span class="n">classes</span><span class="p">,</span> <span class="n">test_loader</span><span class="p">,</span> <span class="n">train_loader</span><span class="p">,</span> <span class="n">valid_loader</span><span class="p">)</span>
                
 
                            -
                
 
                            -        <span class="c1"># CREATING THE LOGGING DIR BASED ON THE INPUT PARAMS TO PREVENT OVERWRITE OF LOCAL VERSION</span>
                
 
                            -        <span class="k">if</span> <span class="n">ckpt_root_dir</span><span class="p">:</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">checkpoints_dir_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">ckpt_root_dir</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">experiment_name</span><span class="p">)</span>
                
 
                            -        <span class="k">elif</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_exists</span><span class="p">(</span><span class="s2">&quot;checkpoints&quot;</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">):</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">checkpoints_dir_path</span> <span class="o">=</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s1">&#39;checkpoints&#39;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">experiment_name</span><span class="p">)</span>
                
 
                            -        <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Illegal checkpoints directory: pass ckpt_root_dir that exists, or add &#39;checkpoints&#39; to&quot;</span>
                
 
                            -                             <span class="s2">&quot;resources.&quot;</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_name</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">checkpoints_dir_path</span> <span class="o">=</span> <span class="n">get_checkpoints_dir_path</span><span class="p">(</span><span class="n">experiment_name</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class="p">)</span>
                
 
                                     <span class="c1"># INITIALIZE THE DEVICE FOR THE MODEL</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">_initialize_device</span><span class="p">(</span><span class="n">requested_device</span><span class="o">=</span><span class="n">device</span><span class="p">,</span> <span class="n">requested_multi_gpu</span><span class="o">=</span><span class="n">multi_gpu</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">post_prediction_callback</span> <span class="o">=</span> <span class="n">post_prediction_callback</span>
                
 
                                     <span class="c1"># SET THE DEFAULTS</span>
                
 
                                     <span class="c1"># TODO: SET DEFAULT TRAINING PARAMS FOR EACH TASK</span>
                
 
                            -        <span class="n">default_results_titles</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;Train Loss&#39;</span><span class="p">,</span> <span class="s1">&#39;Train Acc&#39;</span><span class="p">,</span> <span class="s1">&#39;Train Top5&#39;</span><span class="p">,</span> <span class="s1">&#39;Valid Loss&#39;</span><span class="p">,</span> <span class="s1">&#39;Valid Acc&#39;</span><span class="p">,</span> <span class="s1">&#39;Valid Top5&#39;</span><span cla
                
 
                            +        <span class="n">default_results_titles</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;Train Loss&quot;</span><span class="p">,</span> <span class="s2">&quot;Train Acc&quot;</span><span class="p">,</span> <span class="s2">&quot;Train Top5&quot;</span><span class="p">,</span> <span class="s2">&quot;Valid Loss&quot;</span><span class="p">,</span> <span class="s2">&quot;Valid Acc&quot;</span><span class="p">,</span> <span class="s2">&quot;Valid Top5&quot;</sp
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">results_titles</span> <span class="o">=</span> <span class="n">default_results_titles</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">loss_idx_in_results_tuple</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">acc_idx_in_results_tuple</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span>
                
 
                            -        <span class="n">default_train_metrics</span><span class="p">,</span> <span class="n">default_valid_metrics</span> <span class="o">=</span> <span class="n">MetricCollection</span><span class="p">([</span><span class="n">Accuracy</span><span class="p">(),</span> <span class="n">Top5</span><span class="p">()]),</span> <span class="n">MetricCollection</span><span class="p">(</span>
                
 
                            -            <span class="p">[</span><span class="n">Accuracy</span><span class="p">(),</span> <span class="n">Top5</span><span class="p">()])</span>
                
 
                            -
                
 
                            -        <span class="n">default_loss_logging_items_names</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;Loss&quot;</span><span class="p">]</span>
                
 
                            +        <span class="n">default_train_metrics</span><span class="p">,</span> <span class="n">default_valid_metrics</span> <span class="o">=</span> <span class="n">MetricCollection</span><span class="p">([</span><span class="n">Accuracy</span><span class="p">(),</span> <span class="n">Top5</span><span class="p">()]),</span> <span class="n">MetricCollection</span><span class="p">([</span><span class="n">Accuracy</span><span class="p">(),</span> <span class="n">Top5</span><span class="p">()])</spa
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span> <span class="o">=</span> <span class="n">default_train_metrics</span><span class="p">,</span> <span class="n">default_valid_metrics</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">=</span> <span class="n">default_loss_logging_items_names</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">train_monitored_values</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">valid_monitored_values</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                            -    <span class="k">def</span> <span class="nf">_set_dataset_properties</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">classes</span><span class="p">,</span> <span class="n">test_loader</span><span class="p">,</span> <span class="n">train_loader</span><span class="p">,</span> <span class="n">valid_loader</span><span class="p">):</span>
                
 
                            -        <span class="k">if</span> <span class="nb">any</span><span class="p">([</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">valid_loader</span><span class="p">,</span> <span class="n">classes</span><span class="p">])</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">all</span><span class="p">([</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">valid_loader</span><span class="p">,</span> <span cla
                
 
                            -            <span class="k">raise</span> <span class="n">IllegalDataloaderInitialization</span><span class="p">()</span>
                
 
                            +<div class="viewcode-block" id="Trainer.train_from_config"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.Trainer.train_from_config">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">train_from_config</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">cfg</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">DictConfig</span><span class="p">,</span> <span class="nb">dict</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">nn</span><span class="o">.</span><span 
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Trains according to cfg recipe configuration.</span>
                
 
                            -        <span class="n">dataset_params</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;batch_size&quot;</span><span class="p">:</span> <span class="n">train_loader</span><span class="o">.</span><span class="n">batch_size</span> <span class="k">if</span> <span class="n">train_loader</span> <span class="k">else</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                          <span class="s2">&quot;val_batch_size&quot;</span><span class="p">:</span> <span class="n">valid_loader</span><span class="o">.</span><span class="n">batch_size</span> <span class="k">if</span> <span class="n">valid_loader</span> <span class="k">else</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                          <span class="s2">&quot;test_batch_size&quot;</span><span class="p">:</span> <span class="n">test_loader</span><span class="o">.</span><span class="n">batch_size</span> <span class="k">if</span> <span class="n">test_loader</span> <span class="k">else</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                          <span class="s2">&quot;dataset_dir&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            -                          <span class="s2">&quot;s3_link&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">}</span>
                
 
                            +<span class="sd">        @param cfg: The parsed DictConfig from yaml recipe files or a dictionary</span>
                
 
                            +<span class="sd">        @return: the model and the output of trainer.train(...) (i.e results tuple)</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">if</span> <span class="n">train_loader</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">multi_gpu</span> <span class="o">==</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">DISTRIBUTED_DATA_PARALLEL</span><span class="p">:</span>
                
 
                            -            <span class="k">if</span> <span class="ow">not</span> <span class="nb">all</span><span class="p">([</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">train_loader</span><span class="o">.</span><span class="n">sampler</span><span class="p">,</span> <span class="n">DistributedSampler</span><span class="p">),</span>
                
 
                            -                        <span class="nb">isinstance</span><span class="p">(</span><span class="n">valid_loader</span><span class="o">.</span><span class="n">sampler</span><span class="p">,</span> <span class="n">DistributedSampler</span><span class="p">),</span>
                
 
                            -                        <span class="n">test_loader</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">test_loader</span><span class="o">.</span><span class="n">sampler</span><span class="p">,</span> <span class="n">DistributedSampler</span><span class="p">)]):</span>
                
 
                            -                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;DDP training was selected but the dataloader samplers are not of type DistributedSamplers&quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">setup_device</span><span class="p">(</span><span class="n">multi_gpu</span><span class="o">=</span><span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="s2">&quot;multi_gpu&quot;</span><span class="p">,</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">OFF</span><span class="p">),</span> <span class="n">num_gpus</span>
                
 
                            +
                
 
                            +        <span class="c1"># INSTANTIATE ALL OBJECTS IN CFG</span>
                
 
                            +        <span class="n">cfg</span> <span class="o">=</span> <span class="n">hydra</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">instantiate</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">kwargs</span> <span class="o">=</span> <span class="n">parse_args</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="bp">cls</span><span class="o">.</span><span class="fm">__init__</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">trainer</span> <span class="o">=</span> <span class="n">Trainer</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="c1"># INSTANTIATE DATA LOADERS</span>
                
 
                            +        <span class="n">train_dataloader</span> <span class="o">=</span> <span class="n">dataloaders</span><span class="o">.</span><span class="n">get</span><span class="p">(</span>
                
 
                            +            <span class="n">name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">train_dataloader</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">train_dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="n">cf
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="n">val_dataloader</span> <span class="o">=</span> <span class="n">dataloaders</span><span class="o">.</span><span class="n">get</span><span class="p">(</span>
                
 
                            +            <span class="n">name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">val_dataloader</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="n">cfg</s
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="c1"># BUILD NETWORK</span>
                
 
                            +        <span class="n">model</span> <span class="o">=</span> <span class="n">models</span><span class="o">.</span><span class="n">get</span><span class="p">(</span>
                
 
                            +            <span class="n">model_name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">architecture</span><span class="p">,</span>
                
 
                            +            <span class="n">num_classes</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">arch_params</span><span class="o">.</span><span class="n">num_classes</span><span class="p">,</span>
                
 
                            +            <span class="n">arch_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">arch_params</span><span class="p">,</span>
                
 
                            +            <span class="n">strict_load</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="o">.</span><span class="n">strict_load</span><span class="p">,</span>
                
 
                            +            <span class="n">pretrained_weights</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="o">.</span><span class="n">pretrained_weights</span><span class="p">,</span>
                
 
                            +            <span class="n">checkpoint_path</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="o">.</span><span class="n">checkpoint_path</span><span class="p">,</span>
                
 
                            +            <span class="n">load_backbone</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="o">.</span><span class="n">load_backbone</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +        <span class="n">recipe_logged_cfg</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;recipe_config&quot;</span><span class="p">:</span> <span class="n">OmegaConf</span><span class="o">.</span><span class="n">to_container</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">resolve</span><span class="o">=</span><span class="kc">True</span><span class="p">)}</span>
                
 
                            +        <span class="c1"># TRAIN</span>
                
 
                            +        <span class="n">res</span> <span class="o">=</span> <span class="n">trainer</span><span class="o">.</span><span class="n">train</span><span class="p">(</span>
                
 
                            +            <span class="n">model</span><span class="o">=</span><span class="n">model</span><span class="p">,</span>
                
 
                            +            <span class="n">train_loader</span><span class="o">=</span><span class="n">train_dataloader</span><span class="p">,</span>
                
 
                            +            <span class="n">valid_loader</span><span class="o">=</span><span class="n">val_dataloader</span><span class="p">,</span>
                
 
                            +            <span class="n">training_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">training_hyperparams</span><span class="p">,</span>
                
 
                            +            <span class="n">additional_configs_to_log</span><span class="o">=</span><span class="n">recipe_logged_cfg</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="k">return</span> <span class="n">model</span><span class="p">,</span> <span class="n">res</span></div>
                
 
                            +
                
 
                            +<div class="viewcode-block" id="Trainer.resume_experiment"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.Trainer.resume_experiment">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">resume_experiment</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class=
                
 
                            +        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">        Resume a training that was run using our recipes.</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">test_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o
                
 
                            -            <span class="n">HpmStruct</span><span class="p">(</span><span class="o">**</span><span class="n">dataset_params</span><span class="p">),</span> <span class="n">train_loader</span><span class="p">,</span> <span class="n">valid_loader</span><span class="p">,</span> <span class="n">test_loader</span><span class="p">,</span> <span class="n">classes</span>
                
 
                            +<span class="sd">        :param experiment_name:     Name of the experiment to resume</span>
                
 
                            +<span class="sd">        :param ckpt_root_dir:       Directory including the checkpoints</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Resume training using the checkpoint recipe, ignoring the current recipe&quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">cfg</span> <span class="o">=</span> <span class="n">load_experiment_cfg</span><span class="p">(</span><span class="n">experiment_name</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class="p">)</span>
                
 
                            +        <span class="n">add_params_to_cfg</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">params</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;training_hyperparams.resume=True&quot;</span><span class="p">])</span>
                
 
                            +        <span class="bp">cls</span><span class="o">.</span><span class="n">train_from_config</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="SgModel.connect_dataset_interface"><a class="viewcode-back" href="../../../../super_gradients.training.sg_model.html#super_gradients.training.SgModel.connect_dataset_interface">[docs]</a>    <span class="nd">@resolve_param</span><span class="p">(</span><span class="s1">&#39;dataset_interface&#39;</span><span class="p">,</span> <span class="n">DatasetsFactory</span><span class="p">())</span>
                
 
                            -    <span class="k">def</span> <span class="nf">connect_dataset_interface</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">dataset_interface</span><span class="p">:</span> <span class="n">DatasetInterface</span><span class="p">,</span> <span class="n">data_loader_num_workers</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">8</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="Trainer.evaluate_from_recipe"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.Trainer.evaluate_from_recipe">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">evaluate_from_recipe</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">cfg</span><span class="p">:</span> <span class="n">DictConfig</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        :param dataset_interface: DatasetInterface object</span>
                
 
                            -<span class="sd">        :param data_loader_num_workers: The number of threads to initialize the Data Loaders with</span>
                
 
                            -<span class="sd">            The dataset to be connected</span>
                
 
                            +<span class="sd">        Evaluate according to a cfg recipe configuration.</span>
                
 
                            +
                
 
                            +<span class="sd">        Note:   This script does NOT run training, only validation.</span>
                
 
                            +<span class="sd">                Please make sure that the config refers to a PRETRAINED MODEL either from one of your checkpoint or from pretrained weights from model zoo.</span>
                
 
                            +<span class="sd">        :param cfg: The parsed DictConfig from yaml recipe files or a dictionary</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">:</span>
                
 
                            -            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s2">&quot;Overriding the dataloaders that SgModel was initialized with&quot;</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_interface</span> <span class="o">=</span> <span class="n">dataset_interface</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">test_loader</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> \
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">dataset_interface</span><span class="o">.</span><span class="n">get_data_loaders</span><span class="p">(</span><span class="n">batch_size_factor</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">num_devices</span><span class="p">,</span>
                
 
                            -                                                    <span class="n">num_workers</span><span class="o">=</span><span class="n">data_loader_num_workers</span><span class="p">,</span>
                
 
                            -                                                    <span class="n">distributed_sampler</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">multi_gpu</span> <span class="o">==</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">DISTRIBUTED_DATA_PARALLEL</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_interface</span><span class="o">.</span><span class="n">get_dataset_params</span><span class="p">()</span></div>
                
 
                            +        <span class="n">setup_device</span><span class="p">(</span><span class="n">multi_gpu</span><span class="o">=</span><span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="s2">&quot;multi_gpu&quot;</span><span class="p">,</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">OFF</span><span class="p">),</span> <span class="n">num_gpus</span>
                
 
                            -    <span class="c1"># FIXME - we need to resolve flake8&#39;s &#39;function is too complex&#39; for this function</span>
                
 
                            -<div class="viewcode-block" id="SgModel.build_model"><a class="viewcode-back" href="../../../../super_gradients.training.sg_model.html#super_gradients.training.SgModel.build_model">[docs]</a>    <span class="k">def</span> <span class="nf">build_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>  <span class="c1"># noqa: C901 - too complex</span>
                
 
                            -                    <span class="n">architecture</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">],</span>
                
 
                            -                    <span class="n">arch_params</span><span class="o">=</span><span class="p">{},</span> <span class="n">checkpoint_params</span><span class="o">=</span><span class="p">{},</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            -        <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        :param architecture:               Defines the network&#39;s architecture from models/ALL_ARCHITECTURES</span>
                
 
                            -<span class="sd">        :param arch_params:                Architecture H.P. e.g.: block, num_blocks, num_classes, etc.</span>
                
 
                            -<span class="sd">        :param checkpoint_params:          Dictionary like object with the following key:values:</span>
                
 
                            -
                
 
                            -<span class="sd">            load_checkpoint:            Load a pre-trained checkpoint</span>
                
 
                            -<span class="sd">            strict_load:                See StrictLoad class documentation for details.</span>
                
 
                            -<span class="sd">            source_ckpt_folder_name:    folder name to load the checkpoint from (self.experiment_name if none is given)</span>
                
 
                            -<span class="sd">            load_weights_only:          loads only the weight from the checkpoint and zeroize the training params</span>
                
 
                            -<span class="sd">            load_backbone:              loads the provided checkpoint to self.net.backbone instead of self.net</span>
                
 
                            -<span class="sd">            external_checkpoint_path:   The path to the external checkpoint to be loaded. Can be absolute or relative</span>
                
 
                            -<span class="sd">                                               (ie: path/to/checkpoint.pth). If provided, will automatically attempt to</span>
                
 
                            -<span class="sd">                                               load the checkpoint even if the load_checkpoint flag is not provided.</span>
                
 
                            +        <span class="c1"># INSTANTIATE ALL OBJECTS IN CFG</span>
                
 
                            +        <span class="n">cfg</span> <span class="o">=</span> <span class="n">hydra</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">instantiate</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
                
 
                            -<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="k">if</span> <span class="s1">&#39;num_classes&#39;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">arch_params</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                            -            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_interface</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -                <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s1">&#39;Error&#39;</span><span class="p">,</span> <span class="s1">&#39;Number of classes not defined in arch params and dataset is not defined&#39;</span><span class="p">)</span>
                
 
                            -            <span class="k">else</span><span class="p">:</span>
                
 
                            -                <span class="n">arch_params</span><span class="p">[</span><span class="s1">&#39;num_classes&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">classes</span><span class="p">)</span>
                
 
                            +        <span class="n">kwargs</span> <span class="o">=</span> <span class="n">parse_args</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="bp">cls</span><span class="o">.</span><span class="fm">__init__</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">HpmStruct</span><span class="p">(</span><span class="o">**</span><span class="n">arch_params</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">HpmStruct</span><span class="p">(</span><span class="o">**</span><span class="n">checkpoint_params</span><span class="p">)</span>
                
 
                            +        <span class="n">trainer</span> <span class="o">=</span> <span class="n">Trainer</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">net</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_instantiate_net</span><span class="p">(</span><span class="n">architecture</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span><span class="p">,</span> <span class="n">checkpoint_params</span><span class="p">,</span> <span class="o">*</span><span class=
                
 
                            +        <span class="c1"># INSTANTIATE DATA LOADERS</span>
                
 
                            +        <span class="n">val_dataloader</span> <span class="o">=</span> <span class="n">dataloaders</span><span class="o">.</span><span class="n">get</span><span class="p">(</span>
                
 
                            +            <span class="n">name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">val_dataloader</span><span class="p">,</span> <span class="n">dataset_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">dataset_params</span><span class="o">.</span><span class="n">val_dataset_params</span><span class="p">,</span> <span class="n">dataloader_params</span><span class="o">=</span><span class="n">cfg</s
                
 
                            +        <span class="p">)</span>
                
 
                            -        <span class="c1"># SAVE THE ARCHITECTURE FOR NEURAL ARCHITECTURE SEARCH</span>
                
 
                            +        <span class="n">checkpoints_dir</span> <span class="o">=</span> <span class="n">Path</span><span class="p">(</span><span class="n">get_checkpoints_dir_path</span><span class="p">(</span><span class="n">experiment_name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">experiment_name</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">ckpt
                
 
                            +        <span class="n">checkpoint_path</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">checkpoints_dir</span> <span class="o">/</span> <span class="n">cfg</span><span class="o">.</span><span class="n">training_hyperparams</span><span class="o">.</span><span class="n">ckpt_name</span><span class="p">)</span>
                
 
                            +        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Evaluating checkpoint: </span><span class="si">{</span><span class="n">checkpoint_path</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">architecture</span> <span class="o">=</span> <span class="n">architecture</span>
                
 
                            +        <span class="c1"># BUILD NETWORK</span>
                
 
                            +        <span class="n">model</span> <span class="o">=</span> <span class="n">models</span><span class="o">.</span><span class="n">get</span><span class="p">(</span>
                
 
                            +            <span class="n">model_name</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">architecture</span><span class="p">,</span>
                
 
                            +            <span class="n">num_classes</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">arch_params</span><span class="o">.</span><span class="n">num_classes</span><span class="p">,</span>
                
 
                            +            <span class="n">arch_params</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">arch_params</span><span class="p">,</span>
                
 
                            +            <span class="n">pretrained_weights</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="o">.</span><span class="n">pretrained_weights</span><span class="p">,</span>
                
 
                            +            <span class="n">checkpoint_path</span><span class="o">=</span><span class="n">checkpoint_path</span><span class="p">,</span>
                
 
                            +            <span class="n">load_backbone</span><span class="o">=</span><span class="n">cfg</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="o">.</span><span class="n">load_backbone</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">_net_to_device</span><span class="p">()</span>
                
 
                            +        <span class="c1"># TEST</span>
                
 
                            +        <span class="n">val_results_tuple</span> <span class="o">=</span> <span class="n">trainer</span><span class="o">.</span><span class="n">test</span><span class="p">(</span><span class="n">model</span><span class="o">=</span><span class="n">model</span><span class="p">,</span> <span class="n">test_loader</span><span class="o">=</span><span class="n">val_dataloader</span><span class="p">,</span> <span class="n">test_metrics_list</span><span class="o">=</span><span class="n">cfg</span><span
                
 
                            -        <span class="c1"># SET THE FLAG FOR DIFFERENT PARAMETER GROUP OPTIMIZER UPDATE</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">update_param_groups</span> <span class="o">=</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="s1">&#39;update_param_groups&#39;</span><span class="p">)</span>
                
 
                            +        <span class="n">valid_metrics_dict</span> <span class="o">=</span> <span class="n">get_metrics_dict</span><span class="p">(</span><span class="n">val_results_tuple</span><span class="p">,</span> <span class="n">trainer</span><span class="o">.</span><span class="n">test_metrics</span><span class="p">,</span> <span class="n">trainer</span><span class="o">.</span><span class="n">loss_logging_items_names</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">_load_checkpoint_to_model</span><span class="p">()</span></div>
                
 
                            +        <span class="n">results</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;Validate Results&quot;</span><span class="p">]</span>
                
 
                            +        <span class="n">results</span> <span class="o">+=</span> <span class="p">[</span><span class="sa">f</span><span class="s2">&quot;   - </span><span class="si">{</span><span class="n">metric</span><span class="si">:</span><span class="s2">10</span><span class="si">}</span><span class="s2">: </span><span class="si">{</span><span class="n">value</span><span class="si">}</span><span class="s2">&quot;</span> <span class="k">for</span> <span class="n">metric</span><span class="p">,</span> <spa
                
 
                            +        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">results</span><span class="p">))</span></div>
                
 
                            -    <span class="k">def</span> <span class="nf">_set_ckpt_loading_attributes</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="Trainer.evaluate_checkpoint"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.Trainer.evaluate_checkpoint">[docs]</a>    <span class="nd">@classmethod</span>
                
 
                            +    <span class="k">def</span> <span class="nf">evaluate_checkpoint</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">ckpt_name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;ckpt_latest.pth&quot;</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class=
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                            -<span class="sd">        Sets checkpoint loading related attributes according to self.checkpoint_params</span>
                
 
                            +<span class="sd">        Evaluate a checkpoint resulting from one of your previous experiment, using the same parameters (dataset, valid_metrics,...)</span>
                
 
                            +<span class="sd">        as used during the training of the experiment</span>
                
 
                            +
                
 
                            +<span class="sd">        Note:</span>
                
 
                            +<span class="sd">            The parameters will be unchanged even if the recipe used for that experiment was changed since then.</span>
                
 
                            +<span class="sd">            This is to ensure that validation of the experiment will remain exactly the same as during training.</span>
                
 
                            +
                
 
                            +<span class="sd">        Example, evaluate the checkpoint &quot;average_model.pth&quot; from experiment &quot;my_experiment_name&quot;:</span>
                
 
                            +<span class="sd">            &gt;&gt; evaluate_checkpoint(experiment_name=&quot;my_experiment_name&quot;, ckpt_name=&quot;average_model.pth&quot;)</span>
                
 
                            +
                
 
                            +<span class="sd">        :param experiment_name:     Name of the experiment to validate</span>
                
 
                            +<span class="sd">        :param ckpt_name:           Name of the checkpoint to test (&quot;ckpt_latest.pth&quot;, &quot;average_model.pth&quot; or &quot;ckpt_best.pth&quot; for instance)</span>
                
 
                            +<span class="sd">        :param ckpt_root_dir:       Directory including the checkpoints</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">checkpoint</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">strict_load</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;strict_load&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="n">Stric
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">load_ema_as_net</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;load_ema_as_net&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">source_ckpt_folder_name</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;source_ckpt_folder_name&#39;</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">load_checkpoint</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;load_checkpoint&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">load_backbone</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;load_backbone&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">external_checkpoint_path</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;external_checkpoint_path&#39;</span><span class="p">)</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">load_checkpoint</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">external_checkpoint_path</span><span class="p">:</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">load_weights_only</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;load_weights_only&#39;</span><span class="p">,</span>
                
 
                            -                                                          <span class="n">default_val</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_name</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span><span class="p">,</span> <span class="s1">&#39;ckpt_name&#39;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="bp">self</sp
                
 
                            +        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Evaluate checkpoint&quot;</span><span class="p">)</span>
                
 
                            +        <span class="n">cfg</span> <span class="o">=</span> <span class="n">load_experiment_cfg</span><span class="p">(</span><span class="n">experiment_name</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class="p">)</span>
                
 
                            +        <span class="n">add_params_to_cfg</span><span class="p">(</span><span class="n">cfg</span><span class="p">,</span> <span class="n">params</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;training_hyperparams.resume=True&quot;</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;ckpt_name=</span><span class="si">{</span><span class="n">ckpt_name</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">])</span>
                
 
                            +        <span class="bp">cls</span><span class="o">.</span><span class="n">evaluate_from_recipe</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span></div>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_set_dataset_params</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span> <span class="o">=</span> <span class="p">{</span>
                
 
                            +            <span class="s2">&quot;train_dataset_params&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="o">.</span><span class="n">dataset</span><span class="o">.</span><span class="n">dataset_params</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="o">.</span><span class
                
 
                            +            <span class="s2">&quot;train_dataloader_params&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="o">.</span><span class="n">dataloader_params</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span> <span class="s2">&quot;dataloader_params&quot;</span><span
                
 
                            +            <span class="s2">&quot;valid_dataset_params&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span><span class="o">.</span><span class="n">dataset</span><span class="o">.</span><span class="n">dataset_params</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span><span class="o">.</span><span class
                
 
                            +            <span class="s2">&quot;valid_dataloader_params&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span><span class="o">.</span><span class="n">dataloader_params</span> <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span><span class="p">,</span> <span class="s2">&quot;dataloader_params&quot;</span><span
                
 
                            +        <span class="p">}</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span> <span class="o">=</span> <span class="n">HpmStruct</span><span class="p">(</span><span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">dataset_params</span><span class="p">)</span>
                
 
                                 <span class="k">def</span> <span class="nf">_net_to_device</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
@@ -376,20 +421,17 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
                
 
                                     <span class="c1"># FOR MULTI-GPU TRAINING (not distributed)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span><span class="o">.</span><span class="n">sync_bn</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span><span class="p">,</span> <span class="s1">&#39;sync_bn&#39;</span><span class="p">,</span> <span class="n">default_val</span><
                
 
                            +        <span class="n">sync_bn</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="p">,</span> <span class="s2">&quot;sync_bn&quot;</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">multi_gpu</span> <span class="o">==</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">DATA_PARALLEL</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">net</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">DataParallel</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="p">,</span> <span class="n">device_ids</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class
                
 
                                     <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">multi_gpu</span> <span class="o">==</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">DISTRIBUTED_DATA_PARALLEL</span><span class="p">:</span>
                
 
                            -            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span><span class="o">.</span><span class="n">sync_bn</span><span class="p">:</span>
                
 
                            +            <span class="k">if</span> <span class="n">sync_bn</span><span class="p">:</span>
                
 
                                             <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">ddp_silent_mode</span><span class="p">:</span>
                
 
                            -                    <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s1">&#39;DDP - Using Sync Batch Norm... Training time will be affected accordingly&#39;</span><span class="p">)</span>
                
 
                            +                    <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;DDP - Using Sync Batch Norm... Training time will be affected accordingly&quot;</span><span class="p">)</span>
                
 
                                             <span class="bp">self</span><span class="o">.</span><span class="n">net</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">SyncBatchNorm</span><span class="o">.</span><span class="n">convert_sync_batchnorm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="p">)</span><span class="o">.</span><span class="n">to</spa
                
 
                            -            <span class="n">local_rank</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;:&#39;</span><span class="p">)[</span><span class="mi">1</span><span class="p">])</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">net</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">parallel</span><span class="o">.</span><span class="n">DistributedDataParallel</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="p">,</span>
                
 
                            -                                                                 <span class="n">device_ids</span><span class="o">=</span><span class="p">[</span><span class="n">local_rank</span><span class="p">],</span>
                
 
                            -                                                                 <span class="n">output_device</span><span class="o">=</span><span class="n">local_rank</span><span class="p">,</span>
                
 
                            -                                                                 <span class="n">find_unused_parameters</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +            <span class="n">local_rank</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;:&quot;</span><span class="p">)[</span><span class="mi">1</span><span class="p">])</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">net</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">parallel</span><span class="o">.</span><span class="n">DistributedDataParallel</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="p">,</span> <span class="n">device_ids</span><span class="o">=</sp
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">net</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">WrappedModel</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="p">)</span>
                
@@ -404,8 +446,7 @@
 
                                     <span class="c1"># SET THE MODEL IN training STATE</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">train</span><span class="p">()</span>
                
 
                                     <span class="c1"># THE DISABLE FLAG CONTROLS WHETHER THE PROGRESS BAR IS SILENT OR PRINTS THE LOGS</span>
                
 
                            -        <span class="n">progress_bar_train_loader</span> <span class="o">=</span> <span class="n">tqdm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">bar_format</span><span class="o">=</span><span class="s2">&quot;</span><span class="si">{l_bar}{bar:10}{r_bar}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">dynamic_ncols</span><span class="o">=</span><span 
                
 
                            -                                         <span class="n">disable</span><span class="o">=</span><span class="n">silent_mode</span><span class="p">)</span>
                
 
                            +        <span class="n">progress_bar_train_loader</span> <span class="o">=</span> <span class="n">tqdm</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span> <span class="n">bar_format</span><span class="o">=</span><span class="s2">&quot;</span><span class="si">{l_bar}{bar:10}{r_bar}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">dynamic_ncols</span><span class="o">=</span><span 
                
 
                                     <span class="n">progress_bar_train_loader</span><span class="o">.</span><span class="n">set_description</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Train epoch </span><span class="si">{</span><span class="n">epoch</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
                
 
                                     <span class="c1"># RESET/INIT THE METRIC LOGGERS</span>
                
@@ -414,21 +455,23 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
                
 
                                     <span class="n">loss_avg_meter</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">AverageMeter</span><span class="p">()</span>
                
 
                            -        <span class="n">context</span> <span class="o">=</span> <span class="n">PhaseContext</span><span class="p">(</span><span class="n">epoch</span><span class="o">=</span><span class="n">epoch</span><span class="p">,</span>
                
 
                            -                               <span class="n">optimizer</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span>
                
 
                            -                               <span class="n">metrics_compute_fn</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">,</span>
                
 
                            -                               <span class="n">loss_avg_meter</span><span class="o">=</span><span class="n">loss_avg_meter</span><span class="p">,</span>
                
 
                            -                               <span class="n">criterion</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">criterion</span><span class="p">,</span>
                
 
                            -                               <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span>
                
 
                            -                               <span class="n">lr_warmup_epochs</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span><span class="p">,</span>
                
 
                            -                               <span class="n">sg_logger</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="p">,</span>
                
 
                            -                               <span class="n">train_loader</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span>
                
 
                            -                               <span class="n">context_methods</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_context_methods</span><span class="p">(</span><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_END</span><span class="p">),</span>
                
 
                            -                               <span class="n">ddp_silent_mode</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ddp_silent_mode</span><span class="p">)</span>
                
 
                            +        <span class="n">context</span> <span class="o">=</span> <span class="n">PhaseContext</span><span class="p">(</span>
                
 
                            +            <span class="n">epoch</span><span class="o">=</span><span class="n">epoch</span><span class="p">,</span>
                
 
                            +            <span class="n">optimizer</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span>
                
 
                            +            <span class="n">metrics_compute_fn</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">,</span>
                
 
                            +            <span class="n">loss_avg_meter</span><span class="o">=</span><span class="n">loss_avg_meter</span><span class="p">,</span>
                
 
                            +            <span class="n">criterion</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">criterion</span><span class="p">,</span>
                
 
                            +            <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span>
                
 
                            +            <span class="n">lr_warmup_epochs</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span><span class="p">,</span>
                
 
                            +            <span class="n">sg_logger</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="p">,</span>
                
 
                            +            <span class="n">train_loader</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="p">,</span>
                
 
                            +            <span class="n">context_methods</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_get_context_methods</span><span class="p">(</span><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_END</span><span class="p">),</span>
                
 
                            +            <span class="n">ddp_silent_mode</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ddp_silent_mode</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                                     <span class="k">for</span> <span class="n">batch_idx</span><span class="p">,</span> <span class="n">batch_items</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">progress_bar_train_loader</span><span class="p">):</span>
                
 
                                         <span class="n">batch_items</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">tensor_container_to_device</span><span class="p">(</span><span class="n">batch_items</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">non_blocking</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            -            <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">additional_batch_items</span> <span class="o">=</span> <span class="n">sg_model_utils</span><span class="o">.</span><span class="n">unpack_batch_items</span><span class="p">(</span><span class="n">batch_items</span><span class="p">)</span>
                
 
                            +            <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">additional_batch_items</span> <span class="o">=</span> <span class="n">sg_trainer_utils</span><span class="o">.</span><span class="n">unpack_batch_items</span><span class="p">(</span><span class="n">batch_items</span><span class="p">)</span>
                
 
                                         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">pre_prediction_callback</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                             <span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pre_prediction_callback</span><span class="p">(</span><span class="n">inputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">)</span>
                
@@ -440,12 +483,7 @@
 
                                             <span class="c1"># COMPUTE THE LOSS FOR BACK PROP + EXTRA METRICS COMPUTED DURING THE LOSS FORWARD PASS</span>
                
 
                                             <span class="n">loss</span><span class="p">,</span> <span class="n">loss_log_items</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_losses</span><span class="p">(</span><span class="n">outputs</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span>
                
 
                            -            <span class="n">context</span><span class="o">.</span><span class="n">update_context</span><span class="p">(</span><span class="n">batch_idx</span><span class="o">=</span><span class="n">batch_idx</span><span class="p">,</span>
                
 
                            -                                   <span class="n">inputs</span><span class="o">=</span><span class="n">inputs</span><span class="p">,</span>
                
 
                            -                                   <span class="n">preds</span><span class="o">=</span><span class="n">outputs</span><span class="p">,</span>
                
 
                            -                                   <span class="n">target</span><span class="o">=</span><span class="n">targets</span><span class="p">,</span>
                
 
                            -                                   <span class="n">loss_log_items</span><span class="o">=</span><span class="n">loss_log_items</span><span class="p">,</span>
                
 
                            -                                   <span class="o">**</span><span class="n">additional_batch_items</span><span class="p">)</span>
                
 
                            +            <span class="n">context</span><span class="o">.</span><span class="n">update_context</span><span class="p">(</span><span class="n">batch_idx</span><span class="o">=</span><span class="n">batch_idx</span><span class="p">,</span> <span class="n">inputs</span><span class="o">=</span><span class="n">inputs</span><span class="p">,</span> <span class="n">preds</span><span class="o">=</span><span class="n">outputs</span><span class="p">,</span> <span class="n">target</span><span class="o">
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">phase_callback_handler</span><span class="p">(</span><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_END</span><span class="p">,</span> <span class="n">context</span><span class="p">)</span>
                
@@ -457,13 +495,12 @@
 
                                         <span class="c1"># COMPUTE THE RUNNING USER METRICS AND LOSS RUNNING ITEMS. RESULT TUPLE IS THEIR CONCATENATION.</span>
                
 
                                         <span class="n">logging_values</span> <span class="o">=</span> <span class="n">loss_avg_meter</span><span class="o">.</span><span class="n">average</span> <span class="o">+</span> <span class="n">get_metrics_results_tuple</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">)</span>
                
 
                            -            <span class="n">gpu_memory_utilization</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">memory_cached</span><span class="p">()</span> <span class="o">/</span> <span class="mf">1E9</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 cl
                
 
                            +            <span class="n">gpu_memory_utilization</span> <span class="o">=</span> <span class="n">get_gpu_mem_utilization</span><span class="p">()</span> <span class="o">/</span> <span class="mf">1e9</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="mi">0</span>
                
 
                                         <span class="c1"># RENDER METRICS PROGRESS</span>
                
 
                            -            <span class="n">pbar_message_dict</span> <span class="o">=</span> <span class="n">get_train_loop_description_dict</span><span class="p">(</span><span class="n">logging_values</span><span class="p">,</span>
                
 
                            -                                                                <span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">,</span>
                
 
                            -                                                                <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span><span class="p">,</span>
                
 
                            -                                                                <span class="n">gpu_mem</span><span class="o">=</span><span class="n">gpu_memory_utilization</span><span class="p">)</span>
                
 
                            +            <span class="n">pbar_message_dict</span> <span class="o">=</span> <span class="n">get_train_loop_description_dict</span><span class="p">(</span>
                
 
                            +                <span class="n">logging_values</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span><span class="p">,</span> <span class="n">gpu_mem</span><span class="o">=</span><span class="n">gpu_memory_utilization</span>
                
 
                            +            <span class="p">)</span>
                
 
                                         <span class="n">progress_bar_train_loader</span><span class="o">.</span><span class="n">set_postfix</span><span class="p">(</span><span class="o">**</span><span class="n">pbar_message_dict</span><span class="p">)</span>
                
@@ -475,8 +512,9 @@
 
                                     <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">ddp_silent_mode</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">upload</span><span class="p">()</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">train_monitored_values</span> <span class="o">=</span> <span class="n">sg_model_utils</span><span class="o">.</span><span class="n">update_monitored_values_dict</span><span class="p">(</span>
                
 
                            -            <span class="n">monitored_values_dict</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">train_monitored_values</span><span class="p">,</span> <span class="n">new_values_dict</span><span class="o">=</span><span class="n">pbar_message_dict</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">train_monitored_values</span> <span class="o">=</span> <span class="n">sg_trainer_utils</span><span class="o">.</span><span class="n">update_monitored_values_dict</span><span class="p">(</span>
                
 
                            +            <span class="n">monitored_values_dict</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">train_monitored_values</span><span class="p">,</span> <span class="n">new_values_dict</span><span class="o">=</span><span class="n">pbar_message_dict</span>
                
 
                            +        <span class="p">)</span>
                
 
                                     <span class="k">return</span> <span class="n">logging_values</span>
                
@@ -489,12 +527,50 @@
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="n">loss_logging_items</span> <span class="o">=</span> <span class="n">loss</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span>
                
 
                            +        <span class="c1"># ON FIRST BACKWARD, DERRIVE THE LOGGING TITLES.</span>
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">_first_backward</span><span class="p">:</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">_init_loss_logging_names</span><span class="p">(</span><span class="n">loss_logging_items</span><span class="p">)</span>
                
 
                            +            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">metric_to_watch</span><span class="p">:</span>
                
 
                            +                <span class="bp">self</span><span class="o">.</span><span class="n">_init_monitored_items</span><span class="p">()</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">_first_backward</span> <span class="o">=</span> <span class="kc">False</span>
                
 
                            +
                
 
                                     <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">loss_logging_items</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span><span class="p">):</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Loss output length must match loss_logging_items_names. Got &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span>
                
 
                            -                <span class="nb">len</span><span class="p">(</span><span class="n">loss_logging_items</span><span class="p">))</span> <span class="o">+</span> <span class="s1">&#39;, and &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span><span class="p">)))</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                
 
                            +                <span class="s2">&quot;Loss output length must match loss_logging_items_names. Got &quot;</span>
                
 
                            +                <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">loss_logging_items</span><span class="p">))</span>
                
 
                            +                <span class="o">+</span> <span class="s2">&quot;, and &quot;</span>
                
 
                            +                <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span><span class="p">))</span>
                
 
                            +            <span class="p">)</span>
                
 
                                     <span class="c1"># RETURN AND THE LOSS LOGGING ITEMS COMPUTED DURING LOSS FORWARD PASS</span>
                
 
                                     <span class="k">return</span> <span class="n">loss</span><span class="p">,</span> <span class="n">loss_logging_items</span>
                
 
                            +    <span class="k">def</span> <span class="nf">_init_monitored_items</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">+</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span><span class="p">))</span><span class="o">.</span><span class=
                
 
                            +        <span class="c1"># Instantiate the values to monitor (loss/metric)</span>
                
 
                            +        <span class="k">for</span> <span class="n">loss_name</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span><span class="p">:</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">train_monitored_values</span><span class="p">[</span><span class="n">loss_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">MonitoredValue</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">loss_name</span><span class="p">,</span> <span class="n">greater_is_better</span><span class="o">=</span><span class="kc">False</span><span class="p">)</spa
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">valid_monitored_values</span><span class="p">[</span><span class="n">loss_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">MonitoredValue</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">loss_name</span><span class="p">,</span> <span class="n">greater_is_better</span><span class="o">=</span><span class="kc">False</span><span class="p">)</spa
                
 
                            +
                
 
                            +        <span class="k">for</span> <span class="n">metric_name</span> <span class="ow">in</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">):</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">train_monitored_values</span><span class="p">[</span><span class="n">metric_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">MonitoredValue</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">metric_name</span><span class="p">,</span> <span class="n">greater_is_better</span><span class="o">=</span><span class="bp">self</span><span class="o">.</
                
 
                            +
                
 
                            +        <span class="k">for</span> <span class="n">metric_name</span> <span class="ow">in</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span><span class="p">):</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">valid_monitored_values</span><span class="p">[</span><span class="n">metric_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">MonitoredValue</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">metric_name</span><span class="p">,</span> <span class="n">greater_is_better</span><span class="o">=</span><span class="bp">self</span><span class="o">.</
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">results_titles</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;Train_&quot;</span> <span class="o">+</span> <span class="n">t</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">+</span> <span class="n">get_metrics_titles</span><span class="p
                
 
                            +            <span class="s2">&quot;Valid_&quot;</span> <span class="o">+</span> <span class="n">t</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">+</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span><span class="p">)</span>
                
 
                            +        <span class="p">]</span>
                
 
                            +
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">average_best_models</span><span class="p">:</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">model_weight_averaging</span> <span class="o">=</span> <span class="n">ModelWeightAveraging</span><span class="p">(</span>
                
 
                            +                <span class="bp">self</span><span class="o">.</span><span class="n">checkpoints_dir_path</span><span class="p">,</span>
                
 
                            +                <span class="n">greater_is_better</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">greater_metric_to_watch_is_better</span><span class="p">,</span>
                
 
                            +                <span class="n">source_ckpt_folder_name</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">source_ckpt_folder_name</span><span class="p">,</span>
                
 
                            +                <span class="n">metric_to_watch</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_to_watch</span><span class="p">,</span>
                
 
                            +                <span class="n">metric_idx</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span><span class="p">,</span>
                
 
                            +                <span class="n">load_checkpoint</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">load_checkpoint</span><span class="p">,</span>
                
 
                            +            <span class="p">)</span>
                
 
                            +
                
 
                                 <span class="k">def</span> <span class="nf">_backward_step</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">loss</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">epoch</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">:</span> <span class="nb">int</span><span 
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Run backprop on the loss and perform a step</span>
                
@@ -527,44 +603,43 @@
 
                                         <span class="c1"># RUN PHASE CALLBACKS</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">phase_callback_handler</span><span class="p">(</span><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_STEP</span><span class="p">,</span> <span class="n">context</span><span class="p">)</span>
                
 
                            -    <span class="k">def</span> <span class="nf">_save_checkpoint</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">optimizer</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">epoch</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="p">:</span> <span class
                
 
                            -                         <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +    <span class="k">def</span> <span class="nf">_save_checkpoint</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">optimizer</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">epoch</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="p">:</span> <span class
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Save the current state dict as latest (always), best (if metric was improved), epoch# (if determined in training</span>
                
 
                             <span class="sd">        params)</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="c1"># WHEN THE validation_results_tuple IS NONE WE SIMPLY SAVE THE state_dict AS LATEST AND Return</span>
                
 
                                     <span class="k">if</span> <span class="n">validation_results_tuple</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="s1">&#39;ckpt_latest_weights_only.pth&#39;</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="p">{</span><span class="s1">&#39;net&#39;</span><span class="p">:</span> <span class="bp">self</span><span clas
                
 
                            -                                          <span class="n">global_step</span><span class="o">=</span><span class="n">epoch</span><span class="p">)</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="s2">&quot;ckpt_latest_weights_only.pth&quot;</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="p">{</span><span class="s2">&quot;net&quot;</span><span class="p">:</span> <span class="bp">self</span><span 
                
 
                                         <span class="k">return</span>
                
 
                                     <span class="c1"># COMPUTE THE CURRENT metric</span>
                
 
                                     <span class="c1"># IF idx IS A LIST - SUM ALL THE VALUES STORED IN THE LIST&#39;S INDICES</span>
                
 
                            -        <span class="n">metric</span> <span class="o">=</span> <span class="n">validation_results_tuple</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span><span class="p">]</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="k">else</span> \
                
 
                            -            <span class="nb">sum</span><span class="p">([</span><span class="n">validation_results_tuple</span><span class="p">[</span><span class="n">idx</span><span class="p">]</span> <span class="k">for</span> <span class="n">idx</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span><span class="p">])</span>
                
 
                            +        <span class="n">metric</span> <span class="o">=</span> <span class="p">(</span>
                
 
                            +            <span class="n">validation_results_tuple</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span><span class="p">]</span>
                
 
                            +            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span>
                
 
                            +            <span class="k">else</span> <span class="nb">sum</span><span class="p">([</span><span class="n">validation_results_tuple</span><span class="p">[</span><span class="n">idx</span><span class="p">]</span> <span class="k">for</span> <span class="n">idx</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span><span class="p">])</span>
                
 
                            +        <span class="p">)</span>
                
 
                                     <span class="c1"># BUILD THE state_dict</span>
                
 
                            -        <span class="n">state</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;net&#39;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">state_dict</span><span class="p">(),</span> <span class="s1">&#39;acc&#39;</span><span class="p">:</span> <span class="n">metric</span><span class="p">,</span> <span class="s1">&#39;epoch&#39;</span><span class="p">:</span> <span cla
                
 
                            +        <span class="n">state</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;net&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">state_dict</span><span class="p">(),</span> <span class="s2">&quot;acc&quot;</span><span class="p">:</span> <span class="n">metric</span><span class="p">,</span> <span class="s2">&quot;epoch&quot;</span><span class="p">:</span> <sp
                
 
                                     <span class="k">if</span> <span class="n">optimizer</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="n">state</span><span class="p">[</span><span class="s1">&#39;optimizer_state_dict&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
                
 
                            +            <span class="n">state</span><span class="p">[</span><span class="s2">&quot;optimizer_state_dict&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scaler</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="n">state</span><span class="p">[</span><span class="s1">&#39;scaler_state_dict&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">scaler</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
                
 
                            +            <span class="n">state</span><span class="p">[</span><span class="s2">&quot;scaler_state_dict&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">scaler</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">ema</span><span class="p">:</span>
                
 
                            -            <span class="n">state</span><span class="p">[</span><span class="s1">&#39;ema_net&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ema_model</span><span class="o">.</span><span class="n">ema</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
                
 
                            +            <span class="n">state</span><span class="p">[</span><span class="s2">&quot;ema_net&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ema_model</span><span class="o">.</span><span class="n">ema</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
                
 
                                     <span class="c1"># SAVES CURRENT MODEL AS ckpt_latest</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="s1">&#39;ckpt_latest.pth&#39;</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="n">state</span><span class="p">,</span> <span class="n">global_step</span><span class="o">=</span><span class="n">epoch</span><s
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="s2">&quot;ckpt_latest.pth&quot;</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="n">state</span><span class="p">,</span> <span class="n">global_step</span><span class="o">=</span><span class="n">epoch</span>
                
 
                                     <span class="c1"># SAVE MODEL AT SPECIFIC EPOCHS DETERMINED BY save_ckpt_epoch_list</span>
                
 
                                     <span class="k">if</span> <span class="n">epoch</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">save_ckpt_epoch_list</span><span class="p">:</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s1">&#39;ckpt_epoch_</span><span class="si">{</span><span class="n">epoch</span><span class="si">}</span><span class="s1">.pth&#39;</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span clas
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;ckpt_epoch_</span><span class="si">{</span><span class="n">epoch</span><span class="si">}</span><span class="s2">.pth&quot;</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span cl
                
 
                                     <span class="c1"># OVERRIDE THE BEST CHECKPOINT AND best_metric IF metric GOT BETTER THAN THE PREVIOUS BEST</span>
                
 
                            -        <span class="k">if</span> <span class="p">(</span><span class="n">metric</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">best_metric</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_metric_to_watch_is_better</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span>
                
 
                            -                <span class="n">metric</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">best_metric</span> <span class="ow">and</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_metric_to_watch_is_better</span><span class="p">):</span>
                
 
                            +        <span class="k">if</span> <span class="p">(</span><span class="n">metric</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">best_metric</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_metric_to_watch_is_better</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">metric</span> <span class="o">&lt;</span> <span class="bp">self</spa
                
 
                                         <span class="c1"># STORE THE CURRENT metric AS BEST</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">best_metric</span> <span class="o">=</span> <span class="n">metric</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">_save_best_checkpoint</span><span class="p">(</span><span class="n">epoch</span><span class="p">,</span> <span class="n">state</span><span class="p">)</span>
                
@@ -578,16 +653,50 @@
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">average_best_models</span><span class="p">:</span>
                
 
                                         <span class="n">net_for_averaging</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ema_model</span><span class="o">.</span><span class="n">ema</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">ema</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">net</span>
                
 
                            -            <span class="n">averaged_model_sd</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">model_weight_averaging</span><span class="o">.</span><span class="n">get_average_model</span><span class="p">(</span><span class="n">net_for_averaging</span><span class="p">,</span>
                
 
                            -                                                                              <span class="n">validation_results_tuple</span><span class="o">=</span><span class="n">validation_results_tuple</span><span class="p">)</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">average_model_checkpoint_filename</span><span class="p">,</span>
                
 
                            -                                          <span class="n">state_dict</span><span class="o">=</span><span class="p">{</span><span class="s1">&#39;net&#39;</span><span class="p">:</span> <span class="n">averaged_model_sd</span><span class="p">},</span> <span class="n">global_step</span><span class="o">=</span><span class="n">epoch</span><span class="p">)</span>
                
 
                            +            <span class="n">averaged_model_sd</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">model_weight_averaging</span><span class="o">.</span><span class="n">get_average_model</span><span class="p">(</span><span class="n">net_for_averaging</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="o">=</span><span class="n">validation_results_tuple</span><span class="p">)</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">average_model_checkpoint_filename</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="p">{</span><span class="s2">&quot;net&quot;</span><span class="p"
                
 
                                 <span class="k">def</span> <span class="nf">_save_best_checkpoint</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">epoch</span><span class="p">,</span> <span class="n">state</span><span class="p">):</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_checkpoint</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ckpt_best_name</span><span class="p">,</span> <span class="n">state_dict</span><span class="o">=</span><span class="n">state</span><span class="p">,</span> <span class="n">global_step</span><span class="o
                
 
                            +    <span class="k">def</span> <span class="nf">_prep_net_for_train</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">_init_arch_params</span><span class="p">()</span>
                
 
                            +
                
 
                            +        <span class="c1"># TODO: REMOVE THE BELOW LINE (FOR BACKWARD COMPATIBILITY)</span>
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span> <span class="o">=</span> <span class="n">HpmStruct</span><span class="p">(</span><span class="n">load_checkpoint</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">resume</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_net_to_device</span><span class="p">()</span>
                
 
                            +
                
 
                            +        <span class="c1"># SET THE FLAG FOR DIFFERENT PARAMETER GROUP OPTIMIZER UPDATE</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">update_param_groups</span> <span class="o">=</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="s2">&quot;update_param_groups&quot;</span><span class="p">)</span>
                
 
                            +
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">checkpoint</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">strict_load</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="p">,</span> <span class="s2">&quot;resume_strict_load&quot;</span><span class="p">,</span> <span class="n">StrictLoad</span><span class="o">.</span><span class="n"
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">load_ema_as_net</span> <span class="o">=</span> <span class="kc">False</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">load_checkpoint</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="p">,</span> <span class="s2">&quot;resume&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">external_checkpoint_path</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="p">,</span> <span class="s2">&quot;resume_path&quot;</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">load_checkpoint</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">load_checkpoint</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">external_checkpoint_path</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_name</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">get_param</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="p">,</span> <span class="s2">&quot;ckpt_name&quot;</span><span class="p">,</span> <span class="s2">&quot;ckpt_latest.pth&quot;</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_load_checkpoint_to_model</span><span class="p">()</span>
                
 
                            +
                
 
                            +    <span class="k">def</span> <span class="nf">_init_arch_params</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +        <span class="n">default_arch_params</span> <span class="o">=</span> <span class="n">HpmStruct</span><span class="p">()</span>
                
 
                            +        <span class="n">arch_params</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="p">,</span> <span class="s2">&quot;arch_params&quot;</span><span class="p">,</span> <span class="n">default_arch_params</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span> <span class="o">=</span> <span class="n">default_arch_params</span>
                
 
                            +        <span class="k">if</span> <span class="n">arch_params</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span><span class="o">.</span><span class="n">override</span><span class="p">(</span><span class="o">**</span><span class="n">arch_params</span><span class="o">.</span><span class="n">to_dict</span><span class="p">())</span>
                
 
                            +
                
 
                                 <span class="c1"># FIXME - we need to resolve flake8&#39;s &#39;function is too complex&#39; for this function</span>
                
 
                            -<div class="viewcode-block" id="SgModel.train"><a class="viewcode-back" href="../../../../super_gradients.training.sg_model.html#super_gradients.training.SgModel.train">[docs]</a>    <span class="k">def</span> <span class="nf">train</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">training_params</span><span class="p">:</span> <span class="nb">dict</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()):</span>  <span cla
                
 
                            +<div class="viewcode-block" id="Trainer.train"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.Trainer.train">[docs]</a>    <span class="k">def</span> <span class="nf">train</span><span class="p">(</span>
                
 
                            +        <span class="bp">self</span><span class="p">,</span>
                
 
                            +        <span class="n">model</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span>
                
 
                            +        <span class="n">training_params</span><span class="p">:</span> <span class="nb">dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">train_loader</span><span class="p">:</span> <span class="n">DataLoader</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">valid_loader</span><span class="p">:</span> <span class="n">DataLoader</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +        <span class="n">additional_configs_to_log</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +    <span class="p">):</span>  <span class="c1"># noqa: C901</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        train - Trains the Model</span>
                
@@ -596,8 +705,28 @@
 
                             <span class="sd">          the data loaders, as dictionary. The phase context will hold the additional items, under an attribute with</span>
                
 
                             <span class="sd">          the same name as the key in this dictionary. Then such items can be accessed through phase callbacks.</span>
                
 
                            +<span class="sd">            :param additional_configs_to_log: Dict, dictionary containing configs that will be added to the training&#39;s</span>
                
 
                            +<span class="sd">                sg_logger. Format should be {&quot;Config_title_1&quot;: {...}, &quot;Config_title_2&quot;:{..}}.</span>
                
 
                            +<span class="sd">            :param model: torch.nn.Module, model to train.</span>
                
 
                            +<span class="sd">            :param train_loader: Dataloader for train set.</span>
                
 
                            +<span class="sd">            :param valid_loader: Dataloader for validation.</span>
                
 
                             <span class="sd">            :param training_params:</span>
                
 
                            +
                
 
                            +<span class="sd">                - `resume` : bool (default=False)</span>
                
 
                            +
                
 
                            +<span class="sd">                    Whether to continue training from ckpt with the same experiment name</span>
                
 
                            +<span class="sd">                     (i.e resume from CKPT_ROOT_DIR/EXPERIMENT_NAME/CKPT_NAME)</span>
                
 
                            +
                
 
                            +<span class="sd">                - `ckpt_name` : str (default=ckpt_latest.pth)</span>
                
 
                            +
                
 
                            +<span class="sd">                    The checkpoint (.pth file) filename in CKPT_ROOT_DIR/EXPERIMENT_NAME/ to use when resume=True and</span>
                
 
                            +<span class="sd">                     resume_path=None</span>
                
 
                            +
                
 
                            +<span class="sd">                - `resume_path`: str (default=None)</span>
                
 
                            +
                
 
                            +<span class="sd">                    Explicit checkpoint path (.pth file) to use to resume training.</span>
                
 
                            +
                
 
                             <span class="sd">                - `max_epochs` : int</span>
                
 
                             <span class="sd">                    Number of epochs to run training.</span>
                
@@ -659,10 +788,54 @@
 
                             <span class="sd">                    where the computed loss is the sum of a few components we would like to log- these entries in</span>
                
 
                             <span class="sd">                    loss_items).</span>
                
 
                            -<span class="sd">                    When training, set the loss_logging_items_names parameter in train_params to be a list of</span>
                
 
                            -<span class="sd">                    strings, of length n_items who&#39;s ith element is the name of the ith entry in loss_items. Then</span>
                
 
                            -<span class="sd">                    each item will be logged, rendered on tensorboard and &quot;watched&quot; (i.e saving model checkpoints</span>
                
 
                            -<span class="sd">                    according to it).</span>
                
 
                            +<span class="sd">                    IMPORTANT:When dealing with external loss classes, to logg/monitor the loss_items as described</span>
                
 
                            +<span class="sd">                    above by specific string name:</span>
                
 
                            +
                
 
                            +<span class="sd">                    Set a &quot;component_names&quot; property in the loss class, whos instance is passed through train_params,</span>
                
 
                            +<span class="sd">                     to be a list of strings, of length n_items who&#39;s ith element is the name of the ith entry in loss_items.</span>
                
 
                            +<span class="sd">                     Then each item will be logged, rendered on tensorboard and &quot;watched&quot; (i.e saving model checkpoints</span>
                
 
                            +<span class="sd">                     according to it) under &lt;LOSS_CLASS.__name__&gt;&quot;/&quot;&lt;COMPONENT_NAME&gt;. If a single item is returned rather then a</span>
                
 
                            +<span class="sd">                     tuple, it would be logged under &lt;LOSS_CLASS.__name__&gt;. When there is no such attributed, the items</span>
                
 
                            +<span class="sd">                     will be named &lt;LOSS_CLASS.__name__&gt;&quot;/&quot;Loss_&quot;&lt;IDX&gt; according to the length of loss_items</span>
                
 
                            +
                
 
                            +<span class="sd">                    For example:</span>
                
 
                            +<span class="sd">                        class MyLoss(_Loss):</span>
                
 
                            +<span class="sd">                            ...</span>
                
 
                            +<span class="sd">                            def forward(self, inputs, targets):</span>
                
 
                            +<span class="sd">                                ...</span>
                
 
                            +<span class="sd">                                total_loss = comp1 + comp2</span>
                
 
                            +<span class="sd">                                loss_items = torch.cat((total_loss.unsqueeze(0),comp1.unsqueeze(0), comp2.unsqueeze(0)).detach()</span>
                
 
                            +<span class="sd">                                return total_loss, loss_items</span>
                
 
                            +<span class="sd">                            ...</span>
                
 
                            +<span class="sd">                            @property</span>
                
 
                            +<span class="sd">                            def component_names(self):</span>
                
 
                            +<span class="sd">                                return [&quot;total_loss&quot;, &quot;my_1st_component&quot;, &quot;my_2nd_component&quot;]</span>
                
 
                            +
                
 
                            +<span class="sd">                    Trainer.train(...</span>
                
 
                            +<span class="sd">                                    train_params={&quot;loss&quot;:MyLoss(),</span>
                
 
                            +<span class="sd">                                                    ...</span>
                
 
                            +<span class="sd">                                                    &quot;metric_to_watch&quot;: &quot;MyLoss/my_1st_component&quot;}</span>
                
 
                            +
                
 
                            +<span class="sd">                        This will write to log and monitor MyLoss/total_loss, MyLoss/my_1st_component,</span>
                
 
                            +<span class="sd">                         MyLoss/my_2nd_component.</span>
                
 
                            +
                
 
                            +<span class="sd">                   For example:</span>
                
 
                            +<span class="sd">                        class MyLoss2(_Loss):</span>
                
 
                            +<span class="sd">                            ...</span>
                
 
                            +<span class="sd">                            def forward(self, inputs, targets):</span>
                
 
                            +<span class="sd">                                ...</span>
                
 
                            +<span class="sd">                                total_loss = comp1 + comp2</span>
                
 
                            +<span class="sd">                                loss_items = torch.cat((total_loss.unsqueeze(0),comp1.unsqueeze(0), comp2.unsqueeze(0)).detach()</span>
                
 
                            +<span class="sd">                                return total_loss, loss_items</span>
                
 
                            +<span class="sd">                            ...</span>
                
 
                            +
                
 
                            +<span class="sd">                    Trainer.train(...</span>
                
 
                            +<span class="sd">                                    train_params={&quot;loss&quot;:MyLoss(),</span>
                
 
                            +<span class="sd">                                                    ...</span>
                
 
                            +<span class="sd">                                                    &quot;metric_to_watch&quot;: &quot;MyLoss2/loss_0&quot;}</span>
                
 
                            +
                
 
                            +<span class="sd">                        This will write to log and monitor MyLoss2/loss_0, MyLoss2/loss_1, MyLoss2/loss_2</span>
                
 
                            +<span class="sd">                        as they have been named by their positional index in loss_items.</span>
                
 
                             <span class="sd">                    Since running logs will save the loss_items in some internal state, it is recommended that</span>
                
 
                             <span class="sd">                    loss_items are detached from their computational graph for memory efficiency.</span>
                
@@ -711,7 +884,7 @@
 
                             <span class="sd">                        is a list referring to the names of each entry in the output metric (torch tensor of size n)</span>
                
 
                             <span class="sd">                        one of &quot;loss_logging_items_names&quot; i.e which will correspond to an item returned during the</span>
                
 
                            -<span class="sd">                        loss function&#39;s forward pass.</span>
                
 
                            +<span class="sd">                        loss function&#39;s forward pass (see loss docs abov).</span>
                
 
                             <span class="sd">                    At the end of each epoch, if a new best metric_to_watch value is achieved, the models checkpoint</span>
                
 
                             <span class="sd">                    is saved in YOUR_PYTHON_PATH/checkpoints/ckpt_best.pth</span>
                
@@ -799,11 +972,6 @@
 
                             <span class="sd">                    will be added to the tensorboard along with some sample images from the dataset. Currently only</span>
                
 
                             <span class="sd">                    detection datasets are supported for analysis.</span>
                
 
                            -<span class="sd">                -  `save_full_train_log` : bool (default=False)</span>
                
 
                            -
                
 
                            -<span class="sd">                    When set, a full log (of all super_gradients modules, including uncaught exceptions from any other</span>
                
 
                            -<span class="sd">                     module) of the training will be saved in the checkpoint directory under full_train_log.log</span>
                
 
                            -
                
 
                             <span class="sd">                -  `sg_logger` : Union[AbstractSGLogger, str] (defauls=base_sg_logger)</span>
                
 
                             <span class="sd">                    Define the SGLogger object for this training process. The SGLogger handles all disk writes, logs, TensorBoard, remote logging</span>
                
@@ -855,49 +1023,49 @@
 
                             <span class="sd">                        num_calib_batches: int, number of batches to collect the statistics from.</span>
                
 
                            -<span class="sd">                        percentile: float, percentile value to use when SgModel,quant_modules_calib_method=&#39;percentile&#39;.</span>
                
 
                            +<span class="sd">                        percentile: float, percentile value to use when Trainer,quant_modules_calib_method=&#39;percentile&#39;.</span>
                
 
                             <span class="sd">                         Discarded when other methods are used (Default=99.99).</span>
                
 
                             <span class="sd">        :return:</span>
                
 
                             <span class="sd">        &quot;&quot;&quot;</span>
                
 
                                     <span class="k">global</span> <span class="n">logger</span>
                
 
                            +        <span class="k">if</span> <span class="n">training_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="n">training_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span> <span class="o">=</span> <span class="n">train_loader</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span> <span class="o">=</span> <span class="n">valid_loader</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_set_dataset_params</span><span class="p">()</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">net</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s1">&#39;Model&#39;</span><span class="p">,</span> <span class="s1">&#39;No model found&#39;</span><span class="p">)</span>
                
 
                            -        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">dataset_interface</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s1">&#39;Data&#39;</span><span class="p">,</span> <span class="s1">&#39;No dataset found&#39;</span><span class="p">)</span>
                
 
                            -
                
 
                            +        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">multi_gpu</span> <span class="o">==</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">DISTRIBUTED_DATA_PARALLEL</span><span class="p">:</span>
                
 
                            +            <span class="c1"># Note: the dataloader uses sampler of the batch_sampler when it is not None.</span>
                
 
                            +            <span class="n">train_sampler</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="o">.</span><span class="n">batch_sampler</span><span class="o">.</span><span class="n">sampler</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span><span class="o">.</span><span class="n">batch_sampler</span> <span class="ow">is</span> <span class="ow">not</
                
 
                            +            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">train_sampler</span><span class="p">,</span> <span class="n">SequentialSampler</span><span class="p">):</span>
                
 
                            +                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                
 
                            +                    <span class="s2">&quot;You are using a SequentialSampler on you training dataloader, while working on DDP. &quot;</span>
                
 
                            +                    <span class="s2">&quot;This cancels the DDP benefits since it makes each process iterate through the entire dataset&quot;</span>
                
 
                            +                <span class="p">)</span>
                
 
                            +            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">train_sampler</span><span class="p">,</span> <span class="p">(</span><span class="n">DistributedSampler</span><span class="p">,</span> <span class="n">InfiniteSampler</span><span class="p">,</span> <span class="n">RepeatAugSampler</span><span class="p">)):</span>
                
 
                            +                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
                
 
                            +                    <span class="s2">&quot;The training sampler you are using might not support DDP. &quot;</span>
                
 
                            +                    <span class="s2">&quot;If it doesnt, please use one of the following sampler: DistributedSampler, InfiniteSampler, RepeatAugSampler&quot;</span>
                
 
                            +                <span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span> <span class="o">=</span> <span class="n">TrainingParams</span><span class="p">()</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">override</span><span class="p">(</span><span class="o">**</span><span class="n">training_params</span><span class="p">)</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">net</span> <span class="o">=</span> <span class="n">model</span>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_prep_net_for_train</span><span class="p">()</span>
                
 
                            +
                
 
                                     <span class="c1"># SET RANDOM SEED</span>
                
 
                            -        <span class="n">random_seed</span><span class="p">(</span><span class="n">is_ddp</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">multi_gpu</span> <span class="o">==</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">DISTRIBUTED_DATA_PARALLEL</span><span class="p">,</span>
                
 
                            -                    <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">seed</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">seed</span><span class="p">)</span>
                
 
                            +        <span class="n">random_seed</span><span class="p">(</span><span class="n">is_ddp</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">multi_gpu</span> <span class="o">==</span> <span class="n">MultiGPUMode</span><span class="o">.</span><span class="n">DISTRIBUTED_DATA_PARALLEL</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><sp
                
 
                                     <span class="n">silent_mode</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">silent_mode</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">ddp_silent_mode</span>
                
 
                                     <span class="c1"># METRICS</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">_set_train_metrics</span><span class="p">(</span><span class="n">train_metrics_list</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">train_metrics_list</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">_set_valid_metrics</span><span class="p">(</span><span class="n">valid_metrics_list</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">valid_metrics_list</span><span class="p">)</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">loss_logging_items_names</span>
                
 
                            -
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">results_titles</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;Train_&quot;</span> <span class="o">+</span> <span class="n">t</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span>
                
 
                            -                               <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">+</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">train_metrics</span><span class="p">)]</span> <span class="o">+</span> \
                
 
                            -                              <span class="p">[</span><span class="s2">&quot;Valid_&quot;</span> <span class="o">+</span> <span class="n">t</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span>
                
 
                            -                               <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">+</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span><span class="p">)]</span>
                
 
                                     <span class="c1"># Store the metric to follow (loss\accuracy) and initialize as the worst value</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">metric_to_watch</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">metric_to_watch</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">greater_metric_to_watch_is_better</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">greater_metric_to_watch_is_better</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span> <span class="o">=</span> <span class="p">(</span>
                
 
                            -            <span class="bp">self</span><span class="o">.</span><span class="n">loss_logging_items_names</span> <span class="o">+</span> <span class="n">get_metrics_titles</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span><span class="p">))</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_to_watch</span><span class="p">)<
                
 
                            -
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.training_hyperparams.training_hyperparams &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.training_hyperparams.training_hyperparams</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.training_hyperparams.training_hyperparams</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">hydra</span>
<span class="kn">import</span> <span class="nn">pkg_resources</span>
<span class="kn">from</span> <span class="nn">hydra</span> <span class="kn">import</span> <span class="n">compose</span><span class="p">,</span> <span class="n">initialize_config_dir</span>
<span class="kn">from</span> <span class="nn">hydra.core.global_hydra</span> <span class="kn">import</span> <span class="n">GlobalHydra</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">override_default_params_without_nones</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.hydra_utils</span> <span class="kn">import</span> <span class="n">normalize_path</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>


<div class="viewcode-block" id="get"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.get">[docs]</a><span class="k">def</span> <span class="nf">get</span><span class="p">(</span><span class="n">config_name</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Dict</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Class for creating training hyper parameters dictionary, taking defaults from yaml</span>
<span class="sd">     files in src/super_gradients/recipes.</span>

<span class="sd">    :param overriding_params: Dict, dictionary like object containing entries to override in the recipe&#39;s training</span>
<span class="sd">     hyper parameters dictionary.</span>
<span class="sd">    :param config_name: yaml config filename in recipes (for example coco2017_yolox).</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">overriding_params</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">overriding_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
    <span class="n">GlobalHydra</span><span class="o">.</span><span class="n">instance</span><span class="p">()</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
    <span class="n">sg_recipes_dir</span> <span class="o">=</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s2">&quot;super_gradients.recipes&quot;</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
    <span class="k">with</span> <span class="n">initialize_config_dir</span><span class="p">(</span><span class="n">config_dir</span><span class="o">=</span><span class="n">normalize_path</span><span class="p">(</span><span class="n">sg_recipes_dir</span><span class="p">),</span> <span class="n">version_base</span><span class="o">=</span><span class="s2">&quot;1.2&quot;</span><span class="p">):</span>
        <span class="n">cfg</span> <span class="o">=</span> <span class="n">compose</span><span class="p">(</span><span class="n">config_name</span><span class="o">=</span><span class="n">normalize_path</span><span class="p">(</span><span class="n">config_name</span><span class="p">))</span>
        <span class="n">cfg</span> <span class="o">=</span> <span class="n">hydra</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">instantiate</span><span class="p">(</span><span class="n">cfg</span><span class="p">)</span>
        <span class="n">training_params</span> <span class="o">=</span> <span class="n">cfg</span><span class="o">.</span><span class="n">training_hyperparams</span>
        <span class="n">training_params</span> <span class="o">=</span> <span class="n">override_default_params_without_nones</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">,</span> <span class="n">training_params</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">training_params</span></div>


<div class="viewcode-block" id="cifar10_resnet_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.cifar10_resnet_train_params">[docs]</a><span class="k">def</span> <span class="nf">cifar10_resnet_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;cifar10_resnet&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_ddrnet_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.cityscapes_ddrnet_train_params">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_ddrnet_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;cityscapes_ddrnet&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_regseg48_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.cityscapes_regseg48_train_params">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_regseg48_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;cityscapes_regseg48&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_stdc_base_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.cityscapes_stdc_base_train_params">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_stdc_base_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;cityscapes_stdc_base&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_stdc_seg50_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.cityscapes_stdc_seg50_train_params">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_stdc_seg50_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;cityscapes_stdc_seg50&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="cityscapes_stdc_seg75_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.cityscapes_stdc_seg75_train_params">[docs]</a><span class="k">def</span> <span class="nf">cityscapes_stdc_seg75_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;cityscapes_stdc_seg75&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="coco2017_ssd_lite_mobilenet_v2_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.coco2017_ssd_lite_mobilenet_v2_train_params">[docs]</a><span class="k">def</span> <span class="nf">coco2017_ssd_lite_mobilenet_v2_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;coco2017_ssd_lite_mobilenet_v2&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="coco2017_yolox_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.coco2017_yolox_train_params">[docs]</a><span class="k">def</span> <span class="nf">coco2017_yolox_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;coco2017_yolox&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="coco_segmentation_shelfnet_lw_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.coco_segmentation_shelfnet_lw_train_params">[docs]</a><span class="k">def</span> <span class="nf">coco_segmentation_shelfnet_lw_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;coco_segmentation_shelfnet_lw&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_efficientnet_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_efficientnet_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_efficientnet_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_efficientnet&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv2_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_mobilenetv2_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv2_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_mobilenetv2&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv3_base_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_mobilenetv3_base_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv3_base_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_mobilenetv3_base&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv3_large_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_mobilenetv3_large_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv3_large_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_mobilenetv3_large&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_mobilenetv3_small_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_mobilenetv3_small_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_mobilenetv3_small_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_mobilenetv3_small&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_regnetY_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_regnetY_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_regnetY_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_regnetY&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_repvgg_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_repvgg_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_repvgg_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_repvgg&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_resnet50_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_resnet50_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_resnet50_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_resnet50&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_resnet50_kd_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_resnet50_kd_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_resnet50_kd_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_resnet50_kd&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_vit_base_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_vit_base_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_vit_base_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_vit_base&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>


<div class="viewcode-block" id="imagenet_vit_large_train_params"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.training_hyperparams.imagenet_vit_large_train_params">[docs]</a><span class="k">def</span> <span class="nf">imagenet_vit_large_train_params</span><span class="p">(</span><span class="n">overriding_params</span><span class="p">:</span> <span class="n">Dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">get</span><span class="p">(</span><span class="s2">&quot;imagenet_vit_large&quot;</span><span class="p">,</span> <span class="n">overriding_params</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.transforms.transforms &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.transforms.transforms</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.transforms.transforms</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">collections</span>
<span class="kn">import</span> <span class="nn">math</span>
<span class="kn">import</span> <span class="nn">random</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Sequence</span><span class="p">,</span> <span class="n">Dict</span>

<span class="kn">from</span> <span class="nn">PIL</span> <span class="kn">import</span> <span class="n">Image</span><span class="p">,</span> <span class="n">ImageFilter</span><span class="p">,</span> <span class="n">ImageOps</span>
<span class="kn">from</span> <span class="nn">torchvision</span> <span class="kn">import</span> <span class="n">transforms</span> <span class="k">as</span> <span class="n">transforms</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">cv2</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">get_mosaic_coordinate</span><span class="p">,</span> <span class="n">adjust_box_anns</span><span class="p">,</span> <span class="n">xyxy2cxcywh</span><span class="p">,</span> <span class="n">cxcywh2xyxy</span><span class="p">,</span> <span class="n">DetectionTargetsFormat</span>

<span class="n">image_resample</span> <span class="o">=</span> <span class="n">Image</span><span class="o">.</span><span class="n">BILINEAR</span>
<span class="n">mask_resample</span> <span class="o">=</span> <span class="n">Image</span><span class="o">.</span><span class="n">NEAREST</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SegmentationTransform</span><span class="p">:</span>
    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot;{&quot;</span><span class="p">,</span> <span class="s2">&quot;(&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot;}&quot;</span><span class="p">,</span> <span class="s2">&quot;)&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SegResize</span><span class="p">(</span><span class="n">SegmentationTransform</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">h</span> <span class="o">=</span> <span class="n">h</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">w</span> <span class="o">=</span> <span class="n">w</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">resize</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">w</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">h</span><span class="p">),</span> <span class="n">image_resample</span><span class="p">)</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">resize</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">w</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">h</span><span class="p">),</span> <span class="n">mask_resample</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">sample</span>


<span class="k">class</span> <span class="nc">SegRandomFlip</span><span class="p">(</span><span class="n">SegmentationTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Randomly flips the image and mask (synchronously) with probability &#39;prob&#39;.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">prob</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">):</span>
        <span class="k">assert</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">prob</span> <span class="o">&lt;=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;Probability value must be between 0 and 1, found </span><span class="si">{</span><span class="n">prob</span><span class="si">}</span><span class="s2">&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">=</span> <span class="n">prob</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span><span class="p">:</span>
            <span class="n">image</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="n">Image</span><span class="o">.</span><span class="n">FLIP_LEFT_RIGHT</span><span class="p">)</span>
            <span class="n">mask</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="n">Image</span><span class="o">.</span><span class="n">FLIP_LEFT_RIGHT</span><span class="p">)</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span>

        <span class="k">return</span> <span class="n">sample</span>


<span class="k">class</span> <span class="nc">SegRescale</span><span class="p">(</span><span class="n">SegmentationTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Rescales the image and mask (synchronously) while preserving aspect ratio.</span>
<span class="sd">    The rescaling can be done according to scale_factor, short_size or long_size.</span>
<span class="sd">    If more than one argument is given, the rescaling mode is determined by this order: scale_factor, then short_size,</span>
<span class="sd">    then long_size.</span>

<span class="sd">    Args:</span>
<span class="sd">        scale_factor: rescaling is done by multiplying input size by scale_factor:</span>
<span class="sd">            out_size = (scale_factor * w, scale_factor * h)</span>
<span class="sd">        short_size: rescaling is done by determining the scale factor by the ratio short_size / min(h, w).</span>
<span class="sd">        long_size: rescaling is done by determining the scale factor by the ratio long_size / max(h, w).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">short_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">long_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">scale_factor</span> <span class="o">=</span> <span class="n">scale_factor</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">short_size</span> <span class="o">=</span> <span class="n">short_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">long_size</span> <span class="o">=</span> <span class="n">long_size</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">check_valid_arguments</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>
        <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">size</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale_factor</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">scale</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale_factor</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">short_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">short_size</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">)</span>
            <span class="n">scale</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">short_size</span> <span class="o">/</span> <span class="n">short_size</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">long_size</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">)</span>
            <span class="n">scale</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">long_size</span> <span class="o">/</span> <span class="n">long_size</span>

        <span class="n">out_size</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">scale</span> <span class="o">*</span> <span class="n">w</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">scale</span> <span class="o">*</span> <span class="n">h</span><span class="p">)</span>

        <span class="n">image</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">out_size</span><span class="p">,</span> <span class="n">image_resample</span><span class="p">)</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">out_size</span><span class="p">,</span> <span class="n">mask_resample</span><span class="p">)</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span>

        <span class="k">return</span> <span class="n">sample</span>

    <span class="k">def</span> <span class="nf">check_valid_arguments</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale_factor</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">short_size</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">long_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Must assign one rescale argument: scale_factor, short_size or long_size&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale_factor</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale_factor</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Scale factor must be a positive number, found: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">scale_factor</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">short_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">short_size</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Short size must be a positive number, found: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">short_size</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">long_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">long_size</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Long size must be a positive number, found: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">long_size</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SegRandomRescale</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Random rescale the image and mask (synchronously) while preserving aspect ratio.</span>
<span class="sd">    Scale factor is randomly picked between scales [min, max]</span>
<span class="sd">    Args:</span>
<span class="sd">        scales: scale range tuple (min, max), if scales is a float range will be defined as (1, scales) if scales &gt; 1,</span>
<span class="sd">            otherwise (scales, 1). must be a positive number.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">scales</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">List</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">)):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">scales</span> <span class="o">=</span> <span class="n">scales</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">check_valid_arguments</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>
        <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">size</span>

        <span class="n">scale</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">scales</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">scales</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
        <span class="n">out_size</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">scale</span> <span class="o">*</span> <span class="n">w</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">scale</span> <span class="o">*</span> <span class="n">h</span><span class="p">)</span>

        <span class="n">image</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">out_size</span><span class="p">,</span> <span class="n">image_resample</span><span class="p">)</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">out_size</span><span class="p">,</span> <span class="n">mask_resample</span><span class="p">)</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span>

        <span class="k">return</span> <span class="n">sample</span>

    <span class="k">def</span> <span class="nf">check_valid_arguments</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Check the scale values are valid. if order is wrong, flip the order and return the right scale values.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">,</span> <span class="n">collections</span><span class="o">.</span><span class="n">abc</span><span class="o">.</span><span class="n">Iterable</span><span class="p">):</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span> <span class="o">&lt;=</span> <span class="mi">1</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">scales</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">scales</span> <span class="o">=</span> <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;SegRandomRescale scale values must be positive numbers, found: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">[</span><span class="mi">1</span><span class="p">]:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">scales</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">scales</span>


<span class="k">class</span> <span class="nc">SegRandomRotate</span><span class="p">(</span><span class="n">SegmentationTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Randomly rotates image and mask (synchronously) between &#39;min_deg&#39; and &#39;max_deg&#39;.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">min_deg</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="o">-</span><span class="mi">10</span><span class="p">,</span> <span class="n">max_deg</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">10</span><span class="p">,</span> <span class="n">fill_mask</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="n">fill_image</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">List</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_deg</span> <span class="o">=</span> <span class="n">min_deg</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_deg</span> <span class="o">=</span> <span class="n">max_deg</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span> <span class="o">=</span> <span class="n">fill_mask</span>
        <span class="c1"># grey color in RGB mode</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span> <span class="o">=</span> <span class="p">(</span><span class="n">fill_image</span><span class="p">,</span> <span class="n">fill_image</span><span class="p">,</span> <span class="n">fill_image</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">check_valid_arguments</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>

        <span class="n">deg</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">min_deg</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_deg</span><span class="p">)</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">rotate</span><span class="p">(</span><span class="n">deg</span><span class="p">,</span> <span class="n">resample</span><span class="o">=</span><span class="n">image_resample</span><span class="p">,</span> <span class="n">fillcolor</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span><span class="p">)</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">rotate</span><span class="p">(</span><span class="n">deg</span><span class="p">,</span> <span class="n">resample</span><span class="o">=</span><span class="n">mask_resample</span><span class="p">,</span> <span class="n">fillcolor</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span><span class="p">)</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span>

        <span class="k">return</span> <span class="n">sample</span>

    <span class="k">def</span> <span class="nf">check_valid_arguments</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span> <span class="o">=</span> <span class="n">_validate_fill_values_arguments</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SegCropImageAndMask</span><span class="p">(</span><span class="n">SegmentationTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Crops image and mask (synchronously).</span>
<span class="sd">    In &quot;center&quot; mode a center crop is performed while, in &quot;random&quot; mode the drop will be positioned around</span>
<span class="sd">     random coordinates.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">crop_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">List</span><span class="p">],</span> <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">        :param crop_size: tuple of (width, height) for the final crop size, if is scalar size is a</span>
<span class="sd">            square (crop_size, crop_size)</span>
<span class="sd">        :param mode: how to choose the center of the crop, &#39;center&#39; for the center of the input image,</span>
<span class="sd">            &#39;random&#39; center the point is chosen randomally</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span> <span class="o">=</span> <span class="n">crop_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">=</span> <span class="n">mode</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">check_valid_arguments</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>

        <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">size</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;random&quot;</span><span class="p">:</span>
            <span class="n">x1</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">w</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
            <span class="n">y1</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">h</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">x1</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">((</span><span class="n">w</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">/</span> <span class="mf">2.0</span><span class="p">))</span>
            <span class="n">y1</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">((</span><span class="n">h</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="o">/</span> <span class="mf">2.0</span><span class="p">))</span>

        <span class="n">image</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">crop</span><span class="p">((</span><span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">y1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">crop</span><span class="p">((</span><span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">y1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span>

        <span class="k">return</span> <span class="n">sample</span>

    <span class="k">def</span> <span class="nf">check_valid_arguments</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;center&quot;</span><span class="p">,</span> <span class="s2">&quot;random&quot;</span><span class="p">]:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Unsupported mode: found: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">mode</span><span class="si">}</span><span class="s2">, expected: &#39;center&#39; or &#39;random&#39;&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">,</span> <span class="n">collections</span><span class="o">.</span><span class="n">abc</span><span class="o">.</span><span class="n">Iterable</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&lt;=</span> <span class="mi">0</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Crop size must be positive numbers, found: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SegRandomGaussianBlur</span><span class="p">(</span><span class="n">SegmentationTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Adds random Gaussian Blur to image with probability &#39;prob&#39;.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">prob</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">):</span>
        <span class="k">assert</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">prob</span> <span class="o">&lt;=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="s2">&quot;Probability value must be between 0 and 1&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">=</span> <span class="n">prob</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>

        <span class="k">if</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span><span class="p">:</span>
            <span class="n">image</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">filter</span><span class="p">(</span><span class="n">ImageFilter</span><span class="o">.</span><span class="n">GaussianBlur</span><span class="p">(</span><span class="n">radius</span><span class="o">=</span><span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()))</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span>

        <span class="k">return</span> <span class="n">sample</span>


<span class="k">class</span> <span class="nc">SegPadShortToCropSize</span><span class="p">(</span><span class="n">SegmentationTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Pads image to &#39;crop_size&#39;.</span>
<span class="sd">    Should be called only after &quot;SegRescale&quot; or &quot;SegRandomRescale&quot; in augmentations pipeline.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">crop_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">List</span><span class="p">],</span> <span class="n">fill_mask</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="n">fill_image</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">List</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param crop_size: tuple of (width, height) for the final crop size, if is scalar size is a</span>
<span class="sd">            square (crop_size, crop_size)</span>
<span class="sd">        :param fill_mask: value to fill mask labels background.</span>
<span class="sd">        :param fill_image: grey value to fill image padded background.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># CHECK IF CROP SIZE IS A ITERABLE OR SCALAR</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span> <span class="o">=</span> <span class="n">crop_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span> <span class="o">=</span> <span class="n">fill_mask</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">fill_image</span><span class="p">)</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">fill_image</span><span class="p">,</span> <span class="n">Sequence</span><span class="p">)</span> <span class="k">else</span> <span class="n">fill_image</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">check_valid_arguments</span><span class="p">()</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span>
        <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">size</span>

        <span class="c1"># pad images from center symmetrically</span>
        <span class="k">if</span> <span class="n">w</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="ow">or</span> <span class="n">h</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]:</span>
            <span class="n">padh</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">h</span><span class="p">)</span> <span class="o">/</span> <span class="mi">2</span> <span class="k">if</span> <span class="n">h</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="k">else</span> <span class="mi">0</span>
            <span class="n">pad_top</span><span class="p">,</span> <span class="n">pad_bottom</span> <span class="o">=</span> <span class="n">math</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">padh</span><span class="p">),</span> <span class="n">math</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">padh</span><span class="p">)</span>
            <span class="n">padw</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">w</span><span class="p">)</span> <span class="o">/</span> <span class="mi">2</span> <span class="k">if</span> <span class="n">w</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">else</span> <span class="mi">0</span>
            <span class="n">pad_left</span><span class="p">,</span> <span class="n">pad_right</span> <span class="o">=</span> <span class="n">math</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">padw</span><span class="p">),</span> <span class="n">math</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">padw</span><span class="p">)</span>

            <span class="n">image</span> <span class="o">=</span> <span class="n">ImageOps</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="n">border</span><span class="o">=</span><span class="p">(</span><span class="n">pad_left</span><span class="p">,</span> <span class="n">pad_top</span><span class="p">,</span> <span class="n">pad_right</span><span class="p">,</span> <span class="n">pad_bottom</span><span class="p">),</span> <span class="n">fill</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span><span class="p">)</span>
            <span class="n">mask</span> <span class="o">=</span> <span class="n">ImageOps</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">mask</span><span class="p">,</span> <span class="n">border</span><span class="o">=</span><span class="p">(</span><span class="n">pad_left</span><span class="p">,</span> <span class="n">pad_top</span><span class="p">,</span> <span class="n">pad_right</span><span class="p">,</span> <span class="n">pad_bottom</span><span class="p">),</span> <span class="n">fill</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span><span class="p">)</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;mask&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mask</span>

        <span class="k">return</span> <span class="n">sample</span>

    <span class="k">def</span> <span class="nf">check_valid_arguments</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">,</span> <span class="n">collections</span><span class="o">.</span><span class="n">abc</span><span class="o">.</span><span class="n">Iterable</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&lt;=</span> <span class="mi">0</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Crop size must be positive numbers, found: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">crop_size</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span> <span class="o">=</span> <span class="n">_validate_fill_values_arguments</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">fill_mask</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fill_image</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">SegColorJitter</span><span class="p">(</span><span class="n">transforms</span><span class="o">.</span><span class="n">ColorJitter</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">):</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="nb">super</span><span class="p">(</span><span class="n">SegColorJitter</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__call__</span><span class="p">(</span><span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">sample</span>


<span class="k">def</span> <span class="nf">_validate_fill_values_arguments</span><span class="p">(</span><span class="n">fill_mask</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">fill_image</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">List</span><span class="p">]):</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">fill_image</span><span class="p">,</span> <span class="n">collections</span><span class="o">.</span><span class="n">abc</span><span class="o">.</span><span class="n">Iterable</span><span class="p">):</span>
        <span class="c1"># If fill_image is single value, turn to grey color in RGB mode.</span>
        <span class="n">fill_image</span> <span class="o">=</span> <span class="p">(</span><span class="n">fill_image</span><span class="p">,</span> <span class="n">fill_image</span><span class="p">,</span> <span class="n">fill_image</span><span class="p">)</span>
    <span class="k">elif</span> <span class="nb">len</span><span class="p">(</span><span class="n">fill_image</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">3</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;fill_image must be an RGB tuple of size equal to 3, found: </span><span class="si">{</span><span class="n">fill_image</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="c1"># assert values are integers</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">fill_mask</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">or</span> <span class="ow">not</span> <span class="nb">all</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">fill_image</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Fill value must be integers,&quot;</span> <span class="sa">f</span><span class="s2">&quot; found: fill_image = </span><span class="si">{</span><span class="n">fill_image</span><span class="si">}</span><span class="s2">, fill_mask = </span><span class="si">{</span><span class="n">fill_mask</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="c1"># assert values in range 0-255</span>
    <span class="k">if</span> <span class="nb">min</span><span class="p">(</span><span class="n">fill_image</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="nb">max</span><span class="p">(</span><span class="n">fill_image</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">255</span> <span class="ow">or</span> <span class="n">fill_mask</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">fill_mask</span> <span class="o">&gt;</span> <span class="mi">255</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Fill value must be a value from 0 to 255,&quot;</span> <span class="sa">f</span><span class="s2">&quot; found: fill_image = </span><span class="si">{</span><span class="n">fill_image</span><span class="si">}</span><span class="s2">, fill_mask = </span><span class="si">{</span><span class="n">fill_mask</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">fill_mask</span><span class="p">,</span> <span class="n">fill_image</span>


<span class="k">class</span> <span class="nc">DetectionTransform</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Detection transform base class.</span>

<span class="sd">    Complex transforms that require extra data loading can use the the additional_samples_count attribute in a</span>
<span class="sd">     similar fashion to what&#39;s been done in COCODetectionDataset:</span>

<span class="sd">    self._load_additional_inputs_for_transform(sample, transform)</span>

<span class="sd">    # after the above call, sample[&quot;additional_samples&quot;] holds a list of additional inputs and targets.</span>

<span class="sd">    sample = transform(sample)</span>



<span class="sd">    Attributes:</span>
<span class="sd">        additional_samples_count: (int) additional samples to be loaded.</span>
<span class="sd">        non_empty_targets: (bool) whether the additianl targets can have empty targets or not.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">additional_samples_count</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="n">non_empty_targets</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">additional_samples_count</span> <span class="o">=</span> <span class="n">additional_samples_count</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">non_empty_targets</span> <span class="o">=</span> <span class="n">non_empty_targets</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">dict</span><span class="p">,</span> <span class="nb">list</span><span class="p">]):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot;{&quot;</span><span class="p">,</span> <span class="s2">&quot;(&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot;}&quot;</span><span class="p">,</span> <span class="s2">&quot;)&quot;</span><span class="p">)</span>


<div class="viewcode-block" id="DetectionMosaic"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.transforms.DetectionMosaic">[docs]</a><span class="k">class</span> <span class="nc">DetectionMosaic</span><span class="p">(</span><span class="n">DetectionTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    DetectionMosaic detection transform</span>

<span class="sd">    Attributes:</span>
<span class="sd">        input_dim: (tuple) input dimension.</span>
<span class="sd">        prob: (float) probability of applying mosaic.</span>
<span class="sd">        enable_mosaic: (bool) whether to apply mosaic at all (regardless of prob) (default=True).</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dim</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">,</span> <span class="n">prob</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">enable_mosaic</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DetectionMosaic</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><span class="n">additional_samples_count</span><span class="o">=</span><span class="mi">3</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">=</span> <span class="n">prob</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span> <span class="o">=</span> <span class="n">input_dim</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">enable_mosaic</span> <span class="o">=</span> <span class="n">enable_mosaic</span>

<div class="viewcode-block" id="DetectionMosaic.close"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.transforms.DetectionMosaic.close">[docs]</a>    <span class="k">def</span> <span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">additional_samples_count</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">enable_mosaic</span> <span class="o">=</span> <span class="kc">False</span></div>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">dict</span><span class="p">,</span> <span class="nb">list</span><span class="p">]):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">enable_mosaic</span> <span class="ow">and</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span><span class="p">:</span>
            <span class="n">mosaic_labels</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">mosaic_labels_seg</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="n">input_h</span><span class="p">,</span> <span class="n">input_w</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</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">input_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>

            <span class="c1"># yc, xc = s, s  # mosaic center x, y</span>
            <span class="n">yc</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">input_h</span><span class="p">,</span> <span class="mf">1.5</span> <span class="o">*</span> <span class="n">input_h</span><span class="p">))</span>
            <span class="n">xc</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">input_w</span><span class="p">,</span> <span class="mf">1.5</span> <span class="o">*</span> <span class="n">input_w</span><span class="p">))</span>

            <span class="c1"># 3 additional samples, total of 4</span>
            <span class="n">all_samples</span> <span class="o">=</span> <span class="p">[</span><span class="n">sample</span><span class="p">]</span> <span class="o">+</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;additional_samples&quot;</span><span class="p">]</span>

            <span class="k">for</span> <span class="n">i_mosaic</span><span class="p">,</span> <span class="n">mosaic_sample</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">all_samples</span><span class="p">):</span>
                <span class="n">img</span><span class="p">,</span> <span class="n">_labels</span> <span class="o">=</span> <span class="n">mosaic_sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">mosaic_sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span>
                <span class="n">_labels_seg</span> <span class="o">=</span> <span class="n">mosaic_sample</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;target_seg&quot;</span><span class="p">)</span>

                <span class="n">h0</span><span class="p">,</span> <span class="n">w0</span> <span class="o">=</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">]</span>  <span class="c1"># orig hw</span>
                <span class="n">scale</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">*</span> <span class="n">input_h</span> <span class="o">/</span> <span class="n">h0</span><span class="p">,</span> <span class="mf">1.0</span> <span class="o">*</span> <span class="n">input_w</span> <span class="o">/</span> <span class="n">w0</span><span class="p">)</span>
                <span class="n">img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">w0</span> <span class="o">*</span> <span class="n">scale</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">h0</span> <span class="o">*</span> <span class="n">scale</span><span class="p">)),</span> <span class="n">interpolation</span><span class="o">=</span><span class="n">cv2</span><span class="o">.</span><span class="n">INTER_LINEAR</span><span class="p">)</span>
                <span class="c1"># generate output mosaic image</span>
                <span class="p">(</span><span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">c</span><span class="p">)</span> <span class="o">=</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">3</span><span class="p">]</span>
                <span class="k">if</span> <span class="n">i_mosaic</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">mosaic_img</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">full</span><span class="p">((</span><span class="n">input_h</span> <span class="o">*</span> <span class="mi">2</span><span class="p">,</span> <span class="n">input_w</span> <span class="o">*</span> <span class="mi">2</span><span class="p">,</span> <span class="n">c</span><span class="p">),</span> <span class="mi">114</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span>

                <span class="c1"># suffix l means large image, while s means small image in mosaic aug.</span>
                <span class="p">(</span><span class="n">l_x1</span><span class="p">,</span> <span class="n">l_y1</span><span class="p">,</span> <span class="n">l_x2</span><span class="p">,</span> <span class="n">l_y2</span><span class="p">),</span> <span class="p">(</span><span class="n">s_x1</span><span class="p">,</span> <span class="n">s_y1</span><span class="p">,</span> <span class="n">s_x2</span><span class="p">,</span> <span class="n">s_y2</span><span class="p">)</span> <span class="o">=</span> <span class="n">get_mosaic_coordinate</span><span class="p">(</span><span class="n">i_mosaic</span><span class="p">,</span> <span class="n">xc</span><span class="p">,</span> <span class="n">yc</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">input_h</span><span class="p">,</span> <span class="n">input_w</span><span class="p">)</span>

                <span class="n">mosaic_img</span><span class="p">[</span><span class="n">l_y1</span><span class="p">:</span><span class="n">l_y2</span><span class="p">,</span> <span class="n">l_x1</span><span class="p">:</span><span class="n">l_x2</span><span class="p">]</span> <span class="o">=</span> <span class="n">img</span><span class="p">[</span><span class="n">s_y1</span><span class="p">:</span><span class="n">s_y2</span><span class="p">,</span> <span class="n">s_x1</span><span class="p">:</span><span class="n">s_x2</span><span class="p">]</span>
                <span class="n">padw</span><span class="p">,</span> <span class="n">padh</span> <span class="o">=</span> <span class="n">l_x1</span> <span class="o">-</span> <span class="n">s_x1</span><span class="p">,</span> <span class="n">l_y1</span> <span class="o">-</span> <span class="n">s_y1</span>

                <span class="n">labels</span> <span class="o">=</span> <span class="n">_labels</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>

                <span class="c1"># Normalized xywh to pixel xyxy format</span>
                <span class="k">if</span> <span class="n">_labels</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">labels</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">scale</span> <span class="o">*</span> <span class="n">_labels</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">padw</span>
                    <span class="n">labels</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">scale</span> <span class="o">*</span> <span class="n">_labels</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">padh</span>
                    <span class="n">labels</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">scale</span> <span class="o">*</span> <span class="n">_labels</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">padw</span>
                    <span class="n">labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="n">scale</span> <span class="o">*</span> <span class="n">_labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">+</span> <span class="n">padh</span>
                <span class="n">mosaic_labels</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">labels</span><span class="p">)</span>

                <span class="k">if</span> <span class="n">_labels_seg</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">labels_seg</span> <span class="o">=</span> <span class="n">_labels_seg</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                    <span class="k">if</span> <span class="n">_labels</span><span class="o">.</span><span class="n">size</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                        <span class="n">labels_seg</span><span class="p">[:,</span> <span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">scale</span> <span class="o">*</span> <span class="n">labels_seg</span><span class="p">[:,</span> <span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">padw</span>
                        <span class="n">labels_seg</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="o">=</span> <span class="n">scale</span> <span class="o">*</span> <span class="n">labels_seg</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="o">+</span> <span class="n">padh</span>
                    <span class="n">mosaic_labels_seg</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">labels_seg</span><span class="p">)</span>

            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">mosaic_labels</span><span class="p">):</span>
                <span class="n">mosaic_labels</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">(</span><span class="n">mosaic_labels</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
                <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">input_w</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">])</span>
                <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">],</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">input_h</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">])</span>
                <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">],</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">input_w</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">])</span>
                <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">],</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">input_h</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mosaic_labels</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">])</span>

            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">mosaic_labels_seg</span><span class="p">):</span>
                <span class="n">mosaic_labels_seg</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">(</span><span class="n">mosaic_labels_seg</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
                <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">mosaic_labels_seg</span><span class="p">[:,</span> <span class="p">::</span><span class="mi">2</span><span class="p">],</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">input_w</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mosaic_labels_seg</span><span class="p">[:,</span> <span class="p">::</span><span class="mi">2</span><span class="p">])</span>
                <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">mosaic_labels_seg</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">0</span><span class="p">,</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">input_h</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">mosaic_labels_seg</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="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mosaic_img</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mosaic_labels</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;info&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">mosaic_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">mosaic_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">mosaic_labels_seg</span><span class="p">):</span>
                <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target_seg&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">mosaic_labels_seg</span>

        <span class="k">return</span> <span class="n">sample</span></div>


<div class="viewcode-block" id="DetectionRandomAffine"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.transforms.DetectionRandomAffine">[docs]</a><span class="k">class</span> <span class="nc">DetectionRandomAffine</span><span class="p">(</span><span class="n">DetectionTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    DetectionRandomAffine detection transform</span>

<span class="sd">    Attributes:</span>
<span class="sd">     target_size: (tuple) desired output shape.</span>

<span class="sd">     degrees:  (Union[tuple, float]) degrees for random rotation, when float the random values are drawn uniformly</span>
<span class="sd">        from (-degrees, degrees)</span>

<span class="sd">     translate:  (Union[tuple, float]) translate size (in pixels) for random translation, when float the random values</span>
<span class="sd">        are drawn uniformly from (-translate, translate)</span>

<span class="sd">     scales: (Union[tuple, float]) values for random rescale, when float the random values are drawn uniformly</span>
<span class="sd">        from (0.1-scales, 0.1+scales)</span>

<span class="sd">     shear: (Union[tuple, float]) degrees for random shear, when float the random values are drawn uniformly</span>
<span class="sd">        from (shear, shear)</span>

<span class="sd">     enable: (bool) whether to apply the below transform at all.</span>

<span class="sd">     filter_box_candidates: (bool) whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio (default=False).</span>

<span class="sd">     wh_thr: (float) edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller</span>
<span class="sd">      then this values will be filtered out. (default=2)</span>

<span class="sd">     ar_thr: (float) aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger</span>
<span class="sd">      then this values will be filtered out. (default=20)</span>

<span class="sd">     area_thr:(float) threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True.</span>
<span class="sd">      Bounding boxes with such ratio smaller then this value will be filtered out. (default=0.1)</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span> <span class="n">degrees</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">translate</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">scales</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">shear</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">target_size</span><span class="o">=</span><span class="p">(</span><span class="mi">640</span><span class="p">,</span> <span class="mi">640</span><span class="p">),</span> <span class="n">filter_box_candidates</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">wh_thr</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">ar_thr</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span> <span class="n">area_thr</span><span class="o">=</span><span class="mf">0.1</span>
    <span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DetectionRandomAffine</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">degrees</span> <span class="o">=</span> <span class="n">degrees</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">translate</span> <span class="o">=</span> <span class="n">translate</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">scale</span> <span class="o">=</span> <span class="n">scales</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shear</span> <span class="o">=</span> <span class="n">shear</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">target_size</span> <span class="o">=</span> <span class="n">target_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">enable</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">filter_box_candidates</span> <span class="o">=</span> <span class="n">filter_box_candidates</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">wh_thr</span> <span class="o">=</span> <span class="n">wh_thr</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ar_thr</span> <span class="o">=</span> <span class="n">ar_thr</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">area_thr</span> <span class="o">=</span> <span class="n">area_thr</span>

<div class="viewcode-block" id="DetectionRandomAffine.close"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.transforms.DetectionRandomAffine.close">[docs]</a>    <span class="k">def</span> <span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">enable</span> <span class="o">=</span> <span class="kc">False</span></div>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">enable</span><span class="p">:</span>
            <span class="n">img</span><span class="p">,</span> <span class="n">target</span> <span class="o">=</span> <span class="n">random_affine</span><span class="p">(</span>
                <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span>
                <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">],</span>
                <span class="n">sample</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;target_seg&quot;</span><span class="p">),</span>
                <span class="n">target_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">target_size</span><span class="p">,</span>
                <span class="n">degrees</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">degrees</span><span class="p">,</span>
                <span class="n">translate</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">translate</span><span class="p">,</span>
                <span class="n">scales</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">scale</span><span class="p">,</span>
                <span class="n">shear</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">shear</span><span class="p">,</span>
                <span class="n">filter_box_candidates</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">filter_box_candidates</span><span class="p">,</span>
                <span class="n">wh_thr</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">wh_thr</span><span class="p">,</span>
                <span class="n">area_thr</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">area_thr</span><span class="p">,</span>
                <span class="n">ar_thr</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ar_thr</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">img</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">target</span>
        <span class="k">return</span> <span class="n">sample</span></div>


<span class="k">class</span> <span class="nc">DetectionMixup</span><span class="p">(</span><span class="n">DetectionTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Mixup detection transform</span>

<span class="sd">    Attributes:</span>
<span class="sd">        input_dim: (tuple) input dimension.</span>
<span class="sd">        mixup_scale: (tuple) scale range for the additional loaded image for mixup.</span>
<span class="sd">        prob: (float) probability of applying mixup.</span>
<span class="sd">        enable_mixup: (bool) whether to apply mixup at all (regardless of prob) (default=True).</span>
<span class="sd">        flip_prob: (float) prbability to apply horizontal flip to the additional sample.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dim</span><span class="p">,</span> <span class="n">mixup_scale</span><span class="p">,</span> <span class="n">prob</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span> <span class="n">enable_mixup</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">flip_prob</span><span class="o">=</span><span class="mf">0.5</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DetectionMixup</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><span class="n">additional_samples_count</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">non_empty_targets</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span> <span class="o">=</span> <span class="n">input_dim</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mixup_scale</span> <span class="o">=</span> <span class="n">mixup_scale</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">=</span> <span class="n">prob</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">enable_mixup</span> <span class="o">=</span> <span class="n">enable_mixup</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">flip_prob</span> <span class="o">=</span> <span class="n">flip_prob</span>

    <span class="k">def</span> <span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">additional_samples_count</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">enable_mixup</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">enable_mixup</span> <span class="ow">and</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span><span class="p">:</span>
            <span class="n">origin_img</span><span class="p">,</span> <span class="n">origin_labels</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span>
            <span class="n">cp_sample</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;additional_samples&quot;</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">img</span><span class="p">,</span> <span class="n">cp_labels</span> <span class="o">=</span> <span class="n">cp_sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">cp_sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span>
            <span class="n">cp_boxes</span> <span class="o">=</span> <span class="n">cp_labels</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span>

            <span class="n">img</span><span class="p">,</span> <span class="n">cp_boxes</span> <span class="o">=</span> <span class="n">_mirror</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">cp_boxes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">flip_prob</span><span class="p">)</span>
            <span class="c1"># PLUG IN TARGET THE FLIPPED BOXES</span>
            <span class="n">cp_labels</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">cp_boxes</span>

            <span class="n">jit_factor</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="o">*</span><span class="bp">self</span><span class="o">.</span><span class="n">mixup_scale</span><span class="p">)</span>

            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
                <span class="n">cp_img</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">input_dim</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">input_dim</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">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span> <span class="o">*</span> <span class="mi">114</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">cp_img</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span> <span class="o">*</span> <span class="mi">114</span>

            <span class="n">cp_scale_ratio</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</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">input_dim</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
            <span class="n">resized_img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span>
                <span class="n">img</span><span class="p">,</span>
                <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">cp_scale_ratio</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">cp_scale_ratio</span><span class="p">)),</span>
                <span class="n">interpolation</span><span class="o">=</span><span class="n">cv2</span><span class="o">.</span><span class="n">INTER_LINEAR</span><span class="p">,</span>
            <span class="p">)</span>

            <span class="n">cp_img</span><span class="p">[:</span> <span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">cp_scale_ratio</span><span class="p">),</span> <span class="p">:</span> <span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">cp_scale_ratio</span><span class="p">)]</span> <span class="o">=</span> <span class="n">resized_img</span>

            <span class="n">cp_img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span>
                <span class="n">cp_img</span><span class="p">,</span>
                <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">cp_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">jit_factor</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">cp_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">jit_factor</span><span class="p">)),</span>
            <span class="p">)</span>
            <span class="n">cp_scale_ratio</span> <span class="o">*=</span> <span class="n">jit_factor</span>

            <span class="n">origin_h</span><span class="p">,</span> <span class="n">origin_w</span> <span class="o">=</span> <span class="n">cp_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">]</span>
            <span class="n">target_h</span><span class="p">,</span> <span class="n">target_w</span> <span class="o">=</span> <span class="n">origin_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">]</span>
            <span class="n">padded_img</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="nb">max</span><span class="p">(</span><span class="n">origin_h</span><span class="p">,</span> <span class="n">target_h</span><span class="p">),</span> <span class="nb">max</span><span class="p">(</span><span class="n">origin_w</span><span class="p">,</span> <span class="n">target_w</span><span class="p">),</span> <span class="mi">3</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span>
            <span class="n">padded_img</span><span class="p">[:</span><span class="n">origin_h</span><span class="p">,</span> <span class="p">:</span><span class="n">origin_w</span><span class="p">]</span> <span class="o">=</span> <span class="n">cp_img</span>

            <span class="n">x_offset</span><span class="p">,</span> <span class="n">y_offset</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span>
            <span class="k">if</span> <span class="n">padded_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="n">target_h</span><span class="p">:</span>
                <span class="n">y_offset</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">padded_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">target_h</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">padded_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">&gt;</span> <span class="n">target_w</span><span class="p">:</span>
                <span class="n">x_offset</span> <span class="o">=</span> <span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">padded_img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">target_w</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">padded_cropped_img</span> <span class="o">=</span> <span class="n">padded_img</span><span class="p">[</span><span class="n">y_offset</span> <span class="p">:</span> <span class="n">y_offset</span> <span class="o">+</span> <span class="n">target_h</span><span class="p">,</span> <span class="n">x_offset</span> <span class="p">:</span> <span class="n">x_offset</span> <span class="o">+</span> <span class="n">target_w</span><span class="p">]</span>

            <span class="n">cp_bboxes_origin_np</span> <span class="o">=</span> <span class="n">adjust_box_anns</span><span class="p">(</span><span class="n">cp_labels</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span><span class="o">.</span><span class="n">copy</span><span class="p">(),</span> <span class="n">cp_scale_ratio</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">origin_w</span><span class="p">,</span> <span class="n">origin_h</span><span class="p">)</span>
            <span class="n">cp_bboxes_transformed_np</span> <span class="o">=</span> <span class="n">cp_bboxes_origin_np</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
            <span class="n">cp_bboxes_transformed_np</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">cp_bboxes_transformed_np</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">x_offset</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">target_w</span><span class="p">)</span>
            <span class="n">cp_bboxes_transformed_np</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="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">cp_bboxes_transformed_np</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="o">-</span> <span class="n">y_offset</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">target_h</span><span class="p">)</span>

            <span class="n">cls_labels</span> <span class="o">=</span> <span class="n">cp_labels</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">:</span><span class="mi">5</span><span class="p">]</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
            <span class="n">box_labels</span> <span class="o">=</span> <span class="n">cp_bboxes_transformed_np</span>
            <span class="n">labels</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">hstack</span><span class="p">((</span><span class="n">box_labels</span><span class="p">,</span> <span class="n">cls_labels</span><span class="p">))</span>
            <span class="n">origin_labels</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">((</span><span class="n">origin_labels</span><span class="p">,</span> <span class="n">labels</span><span class="p">))</span>
            <span class="n">origin_img</span> <span class="o">=</span> <span class="n">origin_img</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
            <span class="n">origin_img</span> <span class="o">=</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="n">origin_img</span> <span class="o">+</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="n">padded_cropped_img</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>

            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">origin_img</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">),</span> <span class="n">origin_labels</span>
        <span class="k">return</span> <span class="n">sample</span>


<div class="viewcode-block" id="DetectionPaddedRescale"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.transforms.DetectionPaddedRescale">[docs]</a><span class="k">class</span> <span class="nc">DetectionPaddedRescale</span><span class="p">(</span><span class="n">DetectionTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Preprocessing transform to be applied last of all transforms for validation.</span>

<span class="sd">    Image- Rescales and pads to self.input_dim.</span>
<span class="sd">    Targets- pads targets to max_targets, moves the class label to first index, converts boxes format- xyxy -&gt; cxcywh.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        input_dim: (tuple) final input dimension (default=(640,640))</span>
<span class="sd">        swap: image axis&#39;s to be rearranged.</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_dim</span><span class="p">,</span> <span class="n">swap</span><span class="o">=</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">max_targets</span><span class="o">=</span><span class="mi">50</span><span class="p">,</span> <span class="n">pad_value</span><span class="o">=</span><span class="mi">114</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">swap</span> <span class="o">=</span> <span class="n">swap</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span> <span class="o">=</span> <span class="n">input_dim</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span> <span class="o">=</span> <span class="n">max_targets</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pad_value</span> <span class="o">=</span> <span class="n">pad_value</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">]):</span>
        <span class="n">img</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">crowd_targets</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;crowd_target&quot;</span><span class="p">)</span>
        <span class="n">img</span><span class="p">,</span> <span class="n">r</span> <span class="o">=</span> <span class="n">rescale_and_pad_to_size</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">swap</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">pad_value</span><span class="p">)</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">img</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_rescale_target</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="n">r</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">crowd_targets</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;crowd_target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_rescale_target</span><span class="p">(</span><span class="n">crowd_targets</span><span class="p">,</span> <span class="n">r</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">sample</span>

    <span class="k">def</span> <span class="nf">_rescale_target</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">targets</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">,</span> <span class="n">r</span><span class="p">:</span> <span class="nb">float</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;SegRescale the target according to a coefficient used to rescale the image.</span>
<span class="sd">        This is done to have images and targets at the same scale.</span>

<span class="sd">        :param targets:  Targets to rescale, shape (batch_size, 6)</span>
<span class="sd">        :param r:        SegRescale coefficient that was applied to the image</span>

<span class="sd">        :return:         Rescaled targets, shape (batch_size, 6)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">targets</span> <span class="o">=</span> <span class="n">targets</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span><span class="p">,</span> <span class="mi">5</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
        <span class="n">boxes</span><span class="p">,</span> <span class="n">labels</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">],</span> <span class="n">targets</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">]</span>
        <span class="n">boxes</span> <span class="o">=</span> <span class="n">xyxy2cxcywh</span><span class="p">(</span><span class="n">boxes</span><span class="p">)</span>
        <span class="n">boxes</span> <span class="o">*=</span> <span class="n">r</span>
        <span class="n">boxes</span> <span class="o">=</span> <span class="n">cxcywh2xyxy</span><span class="p">(</span><span class="n">boxes</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">((</span><span class="n">boxes</span><span class="p">,</span> <span class="n">labels</span><span class="p">[:,</span> <span class="n">np</span><span class="o">.</span><span class="n">newaxis</span><span class="p">]),</span> <span class="mi">1</span><span class="p">)</span></div>


<span class="k">class</span> <span class="nc">DetectionHorizontalFlip</span><span class="p">(</span><span class="n">DetectionTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Horizontal Flip for Detection</span>

<span class="sd">    Attributes:</span>
<span class="sd">        prob: float: probability of applying horizontal flip</span>
<span class="sd">        max_targets: int: max objects in single image, padding target to this size in case of empty image.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">prob</span><span class="p">,</span> <span class="n">max_targets</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">120</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DetectionHorizontalFlip</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">=</span> <span class="n">prob</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span> <span class="o">=</span> <span class="n">max_targets</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">):</span>
        <span class="n">image</span><span class="p">,</span> <span class="n">targets</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span>
        <span class="n">boxes</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">boxes</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">targets</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span><span class="p">,</span> <span class="mi">5</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
            <span class="n">boxes</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span>
        <span class="n">image</span><span class="p">,</span> <span class="n">boxes</span> <span class="o">=</span> <span class="n">_mirror</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="n">boxes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span><span class="p">)</span>
        <span class="n">targets</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">targets</span>
        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">image</span>
        <span class="k">return</span> <span class="n">sample</span>


<div class="viewcode-block" id="DetectionHSV"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.transforms.DetectionHSV">[docs]</a><span class="k">class</span> <span class="nc">DetectionHSV</span><span class="p">(</span><span class="n">DetectionTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Detection HSV transform.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">prob</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">hgain</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">sgain</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">vgain</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">bgr_channels</span><span class="o">=</span><span class="p">(</span><span class="mi">0</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="nb">super</span><span class="p">(</span><span class="n">DetectionHSV</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">prob</span> <span class="o">=</span> <span class="n">prob</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">hgain</span> <span class="o">=</span> <span class="n">hgain</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sgain</span> <span class="o">=</span> <span class="n">sgain</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">vgain</span> <span class="o">=</span> <span class="n">vgain</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bgr_channels</span> <span class="o">=</span> <span class="n">bgr_channels</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">:</span> <span class="nb">dict</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">prob</span><span class="p">:</span>
            <span class="n">augment_hsv</span><span class="p">(</span><span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">hgain</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">sgain</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">vgain</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">bgr_channels</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">sample</span></div>


<div class="viewcode-block" id="DetectionTargetsFormatTransform"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.transforms.DetectionTargetsFormatTransform">[docs]</a><span class="k">class</span> <span class="nc">DetectionTargetsFormatTransform</span><span class="p">(</span><span class="n">DetectionTransform</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Detection targets format transform</span>

<span class="sd">    Converts targets in input_format to output_format.</span>
<span class="sd">    Attributes:</span>
<span class="sd">        input_format: DetectionTargetsFormat: input target format</span>
<span class="sd">        output_format: DetectionTargetsFormat: output target format</span>
<span class="sd">        min_bbox_edge_size: int: bboxes with edge size lower then this values will be removed.</span>
<span class="sd">        max_targets: int: max objects in single image, padding target to this size.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">input_format</span><span class="p">:</span> <span class="n">DetectionTargetsFormat</span> <span class="o">=</span> <span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">XYXY_LABEL</span><span class="p">,</span>
        <span class="n">output_format</span><span class="p">:</span> <span class="n">DetectionTargetsFormat</span> <span class="o">=</span> <span class="n">DetectionTargetsFormat</span><span class="o">.</span><span class="n">LABEL_CXCYWH</span><span class="p">,</span>
        <span class="n">min_bbox_edge_size</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
        <span class="n">max_targets</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">120</span><span class="p">,</span>
    <span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DetectionTargetsFormatTransform</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_format</span> <span class="o">=</span> <span class="n">input_format</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">output_format</span> <span class="o">=</span> <span class="n">output_format</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_bbox_edge_size</span> <span class="o">=</span> <span class="n">min_bbox_edge_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span> <span class="o">=</span> <span class="n">max_targets</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">sample</span><span class="p">):</span>
        <span class="n">normalized_input</span> <span class="o">=</span> <span class="s2">&quot;NORMALIZED&quot;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_format</span><span class="o">.</span><span class="n">value</span>
        <span class="n">normalized_output</span> <span class="o">=</span> <span class="s2">&quot;NORMALIZED&quot;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_format</span><span class="o">.</span><span class="n">value</span>
        <span class="n">normalize</span> <span class="o">=</span> <span class="ow">not</span> <span class="n">normalized_input</span> <span class="ow">and</span> <span class="n">normalized_output</span>
        <span class="n">denormalize</span> <span class="o">=</span> <span class="n">normalized_input</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">normalized_output</span>

        <span class="n">label_first_in_input</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_format</span><span class="o">.</span><span class="n">value</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;_&quot;</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;LABEL&quot;</span>
        <span class="n">label_first_in_output</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_format</span><span class="o">.</span><span class="n">value</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;_&quot;</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;LABEL&quot;</span>

        <span class="n">input_xyxy_format</span> <span class="o">=</span> <span class="s2">&quot;XYXY&quot;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_format</span><span class="o">.</span><span class="n">value</span>
        <span class="n">output_xyxy_format</span> <span class="o">=</span> <span class="s2">&quot;XYXY&quot;</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_format</span><span class="o">.</span><span class="n">value</span>
        <span class="n">convert2xyxy</span> <span class="o">=</span> <span class="ow">not</span> <span class="n">input_xyxy_format</span> <span class="ow">and</span> <span class="n">output_xyxy_format</span>
        <span class="n">convert2cxcy</span> <span class="o">=</span> <span class="n">input_xyxy_format</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">output_xyxy_format</span>

        <span class="n">image</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">crowd_targets</span> <span class="o">=</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;image&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">],</span> <span class="n">sample</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;crowd_target&quot;</span><span class="p">)</span>

        <span class="n">_</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">shape</span>

        <span class="k">def</span> <span class="nf">_format_target</span><span class="p">(</span><span class="n">targets_in</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">label_first_in_input</span><span class="p">:</span>
                <span class="n">labels</span><span class="p">,</span> <span class="n">boxes</span> <span class="o">=</span> <span class="n">targets_in</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">targets_in</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">:]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">boxes</span><span class="p">,</span> <span class="n">labels</span> <span class="o">=</span> <span class="n">targets_in</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">],</span> <span class="n">targets_in</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">]</span>

            <span class="k">if</span> <span class="n">convert2cxcy</span><span class="p">:</span>
                <span class="n">boxes</span> <span class="o">=</span> <span class="n">xyxy2cxcywh</span><span class="p">(</span><span class="n">boxes</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">convert2xyxy</span><span class="p">:</span>
                <span class="n">boxes</span> <span class="o">=</span> <span class="n">cxcywh2xyxy</span><span class="p">(</span><span class="n">boxes</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">normalize</span><span class="p">:</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">w</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">h</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">/</span> <span class="n">w</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="n">h</span>

            <span class="k">elif</span> <span class="n">denormalize</span><span class="p">:</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">w</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">h</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="n">w</span>
                <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">*</span> <span class="n">h</span>

            <span class="n">min_bbox_edge_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">min_bbox_edge_size</span> <span class="o">/</span> <span class="nb">max</span><span class="p">(</span><span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">)</span> <span class="k">if</span> <span class="n">normalized_output</span> <span class="k">else</span> <span class="bp">self</span><span class="o">.</span><span class="n">min_bbox_edge_size</span>

            <span class="n">cxcywh_boxes</span> <span class="o">=</span> <span class="n">boxes</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">output_xyxy_format</span> <span class="k">else</span> <span class="n">xyxy2cxcywh</span><span class="p">(</span><span class="n">boxes</span><span class="o">.</span><span class="n">copy</span><span class="p">())</span>

            <span class="n">mask_b</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">minimum</span><span class="p">(</span><span class="n">cxcywh_boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">],</span> <span class="n">cxcywh_boxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">])</span> <span class="o">&gt;</span> <span class="n">min_bbox_edge_size</span>
            <span class="n">boxes_t</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[</span><span class="n">mask_b</span><span class="p">]</span>
            <span class="n">labels_t</span> <span class="o">=</span> <span class="n">labels</span><span class="p">[</span><span class="n">mask_b</span><span class="p">]</span>

            <span class="n">labels_t</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">expand_dims</span><span class="p">(</span><span class="n">labels_t</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">targets_t</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">hstack</span><span class="p">((</span><span class="n">labels_t</span><span class="p">,</span> <span class="n">boxes_t</span><span class="p">))</span> <span class="k">if</span> <span class="n">label_first_in_output</span> <span class="k">else</span> <span class="n">np</span><span class="o">.</span><span class="n">hstack</span><span class="p">((</span><span class="n">boxes_t</span><span class="p">,</span> <span class="n">labels_t</span><span class="p">))</span>
            <span class="n">padded_targets</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span><span class="p">,</span> <span class="mi">5</span><span class="p">))</span>
            <span class="n">padded_targets</span><span class="p">[</span><span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">targets_t</span><span class="p">))[:</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span><span class="p">]]</span> <span class="o">=</span> <span class="n">targets_t</span><span class="p">[:</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_targets</span><span class="p">]</span>
            <span class="n">padded_targets</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ascontiguousarray</span><span class="p">(</span><span class="n">padded_targets</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>

            <span class="k">return</span> <span class="n">padded_targets</span>

        <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">_format_target</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">crowd_targets</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">sample</span><span class="p">[</span><span class="s2">&quot;crowd_target&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">_format_target</span><span class="p">(</span><span class="n">crowd_targets</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">sample</span></div>


<span class="k">def</span> <span class="nf">get_aug_params</span><span class="p">(</span><span class="n">value</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">float</span><span class="p">],</span> <span class="n">center</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Generates a random value for augmentations as described below</span>

<span class="sd">    :param value: Union[tuple, float] defines the range of values for generation. Wen tuple-</span>
<span class="sd">     drawn uniformly between (value[0], value[1]), and (center - value, center + value) when float</span>
<span class="sd">    :param center: float, defines center to subtract when value is float.</span>
<span class="sd">    :return: generated value</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="nb">float</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">center</span> <span class="o">-</span> <span class="n">value</span><span class="p">,</span> <span class="n">center</span> <span class="o">+</span> <span class="n">value</span><span class="p">)</span>
    <span class="k">elif</span> <span class="nb">len</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">value</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">value</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;Affine params should be either a sequence containing two values</span><span class="se">\</span>
<span class="s2">                          or single float values. Got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                <span class="n">value</span>
            <span class="p">)</span>
        <span class="p">)</span>


<span class="k">def</span> <span class="nf">get_affine_matrix</span><span class="p">(</span>
    <span class="n">target_size</span><span class="p">,</span>
    <span class="n">degrees</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="n">translate</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
    <span class="n">scales</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
    <span class="n">shear</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
<span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a random affine transform matrix.</span>

<span class="sd">    :param target_size: (tuple) desired output shape.</span>

<span class="sd">    :param degrees:  (Union[tuple, float]) degrees for random rotation, when float the random values are drawn uniformly</span>
<span class="sd">     from (-degrees, degrees)</span>

<span class="sd">    :param translate:  (Union[tuple, float]) translate size (in pixels) for random translation, when float the random values</span>
<span class="sd">     are drawn uniformly from (-translate, translate)</span>

<span class="sd">    :param scales: (Union[tuple, float]) values for random rescale, when float the random values are drawn uniformly</span>
<span class="sd">     from (0.1-scales, 0.1+scales)</span>

<span class="sd">    :param shear: (Union[tuple, float]) degrees for random shear, when float the random values are drawn uniformly</span>
<span class="sd">     from (shear, shear)</span>

<span class="sd">    :return: affine_transform_matrix, drawn_scale</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">twidth</span><span class="p">,</span> <span class="n">theight</span> <span class="o">=</span> <span class="n">target_size</span>

    <span class="c1"># Rotation and Scale</span>
    <span class="n">angle</span> <span class="o">=</span> <span class="n">get_aug_params</span><span class="p">(</span><span class="n">degrees</span><span class="p">)</span>
    <span class="n">scale</span> <span class="o">=</span> <span class="n">get_aug_params</span><span class="p">(</span><span class="n">scales</span><span class="p">,</span> <span class="n">center</span><span class="o">=</span><span class="mf">1.0</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">scale</span> <span class="o">&lt;=</span> <span class="mf">0.0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Argument scale should be positive&quot;</span><span class="p">)</span>

    <span class="n">R</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">getRotationMatrix2D</span><span class="p">(</span><span class="n">angle</span><span class="o">=</span><span class="n">angle</span><span class="p">,</span> <span class="n">center</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">scale</span><span class="o">=</span><span class="n">scale</span><span class="p">)</span>

    <span class="n">M</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">([</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">])</span>
    <span class="c1"># Shear</span>
    <span class="n">shear_x</span> <span class="o">=</span> <span class="n">math</span><span class="o">.</span><span class="n">tan</span><span class="p">(</span><span class="n">get_aug_params</span><span class="p">(</span><span class="n">shear</span><span class="p">)</span> <span class="o">*</span> <span class="n">math</span><span class="o">.</span><span class="n">pi</span> <span class="o">/</span> <span class="mi">180</span><span class="p">)</span>
    <span class="n">shear_y</span> <span class="o">=</span> <span class="n">math</span><span class="o">.</span><span class="n">tan</span><span class="p">(</span><span class="n">get_aug_params</span><span class="p">(</span><span class="n">shear</span><span class="p">)</span> <span class="o">*</span> <span class="n">math</span><span class="o">.</span><span class="n">pi</span> <span class="o">/</span> <span class="mi">180</span><span class="p">)</span>

    <span class="n">M</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">R</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">shear_y</span> <span class="o">*</span> <span class="n">R</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
    <span class="n">M</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">R</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">shear_x</span> <span class="o">*</span> <span class="n">R</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>

    <span class="c1"># Translation</span>
    <span class="n">translation_x</span> <span class="o">=</span> <span class="n">get_aug_params</span><span class="p">(</span><span class="n">translate</span><span class="p">)</span> <span class="o">*</span> <span class="n">twidth</span>  <span class="c1"># x translation (pixels)</span>
    <span class="n">translation_y</span> <span class="o">=</span> <span class="n">get_aug_params</span><span class="p">(</span><span class="n">translate</span><span class="p">)</span> <span class="o">*</span> <span class="n">theight</span>  <span class="c1"># y translation (pixels)</span>

    <span class="n">M</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">translation_x</span>
    <span class="n">M</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="o">=</span> <span class="n">translation_y</span>

    <span class="k">return</span> <span class="n">M</span><span class="p">,</span> <span class="n">scale</span>


<span class="k">def</span> <span class="nf">apply_affine_to_bboxes</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="n">targets_seg</span><span class="p">,</span> <span class="n">target_size</span><span class="p">,</span> <span class="n">M</span><span class="p">):</span>
    <span class="n">num_gts</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span>
    <span class="n">twidth</span><span class="p">,</span> <span class="n">theight</span> <span class="o">=</span> <span class="n">target_size</span>
    <span class="c1"># targets_seg = [B x w x h]</span>
    <span class="c1"># if any is_not_nan in axis = 1</span>
    <span class="n">seg_is_present_mask</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">logical_or</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="o">~</span><span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">targets_seg</span><span class="p">),</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">num_gts_masks</span> <span class="o">=</span> <span class="n">seg_is_present_mask</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
    <span class="n">num_gts_boxes</span> <span class="o">=</span> <span class="n">num_gts</span> <span class="o">-</span> <span class="n">num_gts_masks</span>

    <span class="k">if</span> <span class="n">num_gts_boxes</span><span class="p">:</span>
        <span class="c1"># warp corner points</span>
        <span class="n">corner_points</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="n">num_gts_boxes</span> <span class="o">*</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">3</span><span class="p">))</span>
        <span class="c1"># x1y1, x2y2, x1y2, x2y1</span>
        <span class="n">corner_points</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="o">~</span><span class="n">seg_is_present_mask</span><span class="p">][:,</span> <span class="p">[</span><span class="mi">0</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">3</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">]]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">num_gts_boxes</span> <span class="o">*</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
        <span class="n">corner_points</span> <span class="o">=</span> <span class="n">corner_points</span> <span class="o">@</span> <span class="n">M</span><span class="o">.</span><span class="n">T</span>  <span class="c1"># apply affine transform</span>
        <span class="n">corner_points</span> <span class="o">=</span> <span class="n">corner_points</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">num_gts_boxes</span><span class="p">,</span> <span class="mi">8</span><span class="p">)</span>

        <span class="c1"># create new boxes</span>
        <span class="n">corner_xs</span> <span class="o">=</span> <span class="n">corner_points</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span>
        <span class="n">corner_ys</span> <span class="o">=</span> <span class="n">corner_points</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="n">new_bboxes</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">((</span><span class="n">np</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">corner_xs</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">corner_ys</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">corner_xs</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">corner_ys</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">T</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">new_bboxes</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">4</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">float</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">num_gts_masks</span><span class="p">:</span>
        <span class="c1"># warp segmentation points</span>
        <span class="n">num_seg_points</span> <span class="o">=</span> <span class="n">targets_seg</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">//</span> <span class="mi">2</span>
        <span class="n">corner_points_seg</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="n">num_gts_masks</span> <span class="o">*</span> <span class="n">num_seg_points</span><span class="p">,</span> <span class="mi">3</span><span class="p">))</span>
        <span class="n">corner_points_seg</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">targets_seg</span><span class="p">[</span><span class="n">seg_is_present_mask</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">num_gts_masks</span> <span class="o">*</span> <span class="n">num_seg_points</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
        <span class="n">corner_points_seg</span> <span class="o">=</span> <span class="n">corner_points_seg</span> <span class="o">@</span> <span class="n">M</span><span class="o">.</span><span class="n">T</span>
        <span class="n">corner_points_seg</span> <span class="o">=</span> <span class="n">corner_points_seg</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">num_gts_masks</span><span class="p">,</span> <span class="n">num_seg_points</span> <span class="o">*</span> <span class="mi">2</span><span class="p">)</span>

        <span class="c1"># create new boxes</span>
        <span class="n">seg_points_xs</span> <span class="o">=</span> <span class="n">corner_points_seg</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span>
        <span class="n">seg_points_ys</span> <span class="o">=</span> <span class="n">corner_points_seg</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="n">new_tight_bboxes</span> <span class="o">=</span> <span class="p">(</span>
            <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">((</span><span class="n">np</span><span class="o">.</span><span class="n">nanmin</span><span class="p">(</span><span class="n">seg_points_xs</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">nanmin</span><span class="p">(</span><span class="n">seg_points_ys</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">nanmax</span><span class="p">(</span><span class="n">seg_points_xs</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">nanmax</span><span class="p">(</span><span class="n">seg_points_ys</span><span class="p">,</span> <span class="mi">1</span><span class="p">)))</span>
            <span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
            <span class="o">.</span><span class="n">T</span>
        <span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">new_tight_bboxes</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">4</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">float</span><span class="p">)</span>

    <span class="n">targets</span><span class="p">[</span><span class="o">~</span><span class="n">seg_is_present_mask</span><span class="p">,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">new_bboxes</span>
    <span class="n">targets</span><span class="p">[</span><span class="n">seg_is_present_mask</span><span class="p">,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">new_tight_bboxes</span>

    <span class="c1"># clip boxes</span>
    <span class="n">targets</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">twidth</span><span class="p">)</span>
    <span class="n">targets</span><span class="p">[:,</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="o">=</span> <span class="n">targets</span><span class="p">[:,</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="o">.</span><span class="n">clip</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">theight</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">targets</span>


<span class="k">def</span> <span class="nf">random_affine</span><span class="p">(</span>
    <span class="n">img</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span>
    <span class="n">targets</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span> <span class="o">=</span> <span class="p">(),</span>
    <span class="n">targets_seg</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">target_size</span><span class="p">:</span> <span class="nb">tuple</span> <span class="o">=</span> <span class="p">(</span><span class="mi">640</span><span class="p">,</span> <span class="mi">640</span><span class="p">),</span>
    <span class="n">degrees</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">]</span> <span class="o">=</span> <span class="mi">10</span><span class="p">,</span>
    <span class="n">translate</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
    <span class="n">scales</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
    <span class="n">shear</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">]</span> <span class="o">=</span> <span class="mi">10</span><span class="p">,</span>
    <span class="n">filter_box_candidates</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="n">wh_thr</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
    <span class="n">ar_thr</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span>
    <span class="n">area_thr</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
<span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Performs random affine transform to img, targets</span>
<span class="sd">    :param img:         Input image</span>
<span class="sd">    :param targets:     Input target</span>
<span class="sd">    :param targets_seg: Targets derived from segmentation masks</span>
<span class="sd">    :param target_size: Desired output shape</span>
<span class="sd">    :param degrees:     Degrees for random rotation, when float the random values are drawn uniformly</span>
<span class="sd">                            from (-degrees, degrees).</span>
<span class="sd">    :param translate:   Translate size (in pixels) for random translation, when float the random values</span>
<span class="sd">                            are drawn uniformly from (-translate, translate)</span>
<span class="sd">    :param scales:      Values for random rescale, when float the random values are drawn uniformly</span>
<span class="sd">                            from (0.1-scales, 0.1+scales)</span>
<span class="sd">    :param shear:       Degrees for random shear, when float the random values are drawn uniformly</span>
<span class="sd">                                from (shear, shear)</span>

<span class="sd">    :param filter_box_candidates:    whether to filter out transformed bboxes by edge size, area ratio, and aspect ratio.</span>
<span class="sd">    :param wh_thr: (float) edge size threshold when filter_box_candidates = True. Bounding oxes with edges smaller</span>
<span class="sd">      then this values will be filtered out. (default=2)</span>

<span class="sd">    :param ar_thr: (float) aspect ratio threshold filter_box_candidates = True. Bounding boxes with aspect ratio larger</span>
<span class="sd">      then this values will be filtered out. (default=20)</span>

<span class="sd">    :param area_thr:(float) threshold for area ratio between original image and the transformed one, when when filter_box_candidates = True.</span>
<span class="sd">      Bounding boxes with such ratio smaller then this value will be filtered out. (default=0.1)</span>
<span class="sd">    :return:            Image and Target with applied random affine</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">targets_seg</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">targets</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="mi">0</span><span class="p">))</span> <span class="k">if</span> <span class="n">targets_seg</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">targets_seg</span>
    <span class="n">M</span><span class="p">,</span> <span class="n">scale</span> <span class="o">=</span> <span class="n">get_affine_matrix</span><span class="p">(</span><span class="n">target_size</span><span class="p">,</span> <span class="n">degrees</span><span class="p">,</span> <span class="n">translate</span><span class="p">,</span> <span class="n">scales</span><span class="p">,</span> <span class="n">shear</span><span class="p">)</span>

    <span class="n">img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">warpAffine</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">M</span><span class="p">,</span> <span class="n">dsize</span><span class="o">=</span><span class="n">target_size</span><span class="p">,</span> <span class="n">borderValue</span><span class="o">=</span><span class="p">(</span><span class="mi">114</span><span class="p">,</span> <span class="mi">114</span><span class="p">,</span> <span class="mi">114</span><span class="p">))</span>

    <span class="c1"># Transform label coordinates</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">targets_orig</span> <span class="o">=</span> <span class="n">targets</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="n">targets</span> <span class="o">=</span> <span class="n">apply_affine_to_bboxes</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="n">targets_seg</span><span class="p">,</span> <span class="n">target_size</span><span class="p">,</span> <span class="n">M</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">filter_box_candidates</span><span class="p">:</span>
            <span class="n">box_candidates_ids</span> <span class="o">=</span> <span class="n">_filter_box_candidates</span><span class="p">(</span><span class="n">targets_orig</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">],</span> <span class="n">targets</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">],</span> <span class="n">wh_thr</span><span class="o">=</span><span class="n">wh_thr</span><span class="p">,</span> <span class="n">ar_thr</span><span class="o">=</span><span class="n">ar_thr</span><span class="p">,</span> <span class="n">area_thr</span><span class="o">=</span><span class="n">area_thr</span><span class="p">)</span>
            <span class="n">targets</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="n">box_candidates_ids</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">img</span><span class="p">,</span> <span class="n">targets</span>


<span class="k">def</span> <span class="nf">_filter_box_candidates</span><span class="p">(</span><span class="n">box1</span><span class="p">,</span> <span class="n">box2</span><span class="p">,</span> <span class="n">wh_thr</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">ar_thr</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span> <span class="n">area_thr</span><span class="o">=</span><span class="mf">0.1</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    compute candidate boxes</span>
<span class="sd">        :param box1:        before augment</span>
<span class="sd">        :param box2:        after augment</span>
<span class="sd">        :param wh_thr:      wh_thr (pixels)</span>
<span class="sd">        :param ar_thr:      aspect_ratio_thr</span>
<span class="sd">        :param area_thr:    area_ratio</span>
<span class="sd">    :return:</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">box1</span> <span class="o">=</span> <span class="n">box1</span><span class="o">.</span><span class="n">T</span>
    <span class="n">box2</span> <span class="o">=</span> <span class="n">box2</span><span class="o">.</span><span class="n">T</span>
    <span class="n">w1</span><span class="p">,</span> <span class="n">h1</span> <span class="o">=</span> <span class="n">box1</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">box1</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">box1</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">-</span> <span class="n">box1</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
    <span class="n">w2</span><span class="p">,</span> <span class="n">h2</span> <span class="o">=</span> <span class="n">box2</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">box2</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">box2</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">-</span> <span class="n">box2</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
    <span class="n">ar</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="n">w2</span> <span class="o">/</span> <span class="p">(</span><span class="n">h2</span> <span class="o">+</span> <span class="mf">1e-16</span><span class="p">),</span> <span class="n">h2</span> <span class="o">/</span> <span class="p">(</span><span class="n">w2</span> <span class="o">+</span> <span class="mf">1e-16</span><span class="p">))</span>  <span class="c1"># aspect ratio</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">w2</span> <span class="o">&gt;</span> <span class="n">wh_thr</span><span class="p">)</span> <span class="o">&amp;</span> <span class="p">(</span><span class="n">h2</span> <span class="o">&gt;</span> <span class="n">wh_thr</span><span class="p">)</span> <span class="o">&amp;</span> <span class="p">(</span><span class="n">w2</span> <span class="o">*</span> <span class="n">h2</span> <span class="o">/</span> <span class="p">(</span><span class="n">w1</span> <span class="o">*</span> <span class="n">h1</span> <span class="o">+</span> <span class="mf">1e-16</span><span class="p">)</span> <span class="o">&gt;</span> <span class="n">area_thr</span><span class="p">)</span> <span class="o">&amp;</span> <span class="p">(</span><span class="n">ar</span> <span class="o">&lt;</span> <span class="n">ar_thr</span><span class="p">)</span>  <span class="c1"># candidates</span>


<span class="k">def</span> <span class="nf">_mirror</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="n">boxes</span><span class="p">,</span> <span class="n">prob</span><span class="o">=</span><span class="mf">0.5</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Horizontal flips image and bboxes with probability prob.</span>

<span class="sd">    :param image: (np.array) image to be flipped.</span>
<span class="sd">    :param boxes: (np.array) bboxes to be modified.</span>
<span class="sd">    :param prob: probability to perform flipping.</span>
<span class="sd">    :return: flipped_image, flipped_bboxes</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">flipped_boxes</span> <span class="o">=</span> <span class="n">boxes</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
    <span class="n">_</span><span class="p">,</span> <span class="n">width</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">image</span><span class="o">.</span><span class="n">shape</span>
    <span class="k">if</span> <span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">()</span> <span class="o">&lt;</span> <span class="n">prob</span><span class="p">:</span>
        <span class="n">image</span> <span class="o">=</span> <span class="n">image</span><span class="p">[:,</span> <span class="p">::</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">flipped_boxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">width</span> <span class="o">-</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">::</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">image</span><span class="p">,</span> <span class="n">flipped_boxes</span>


<span class="k">def</span> <span class="nf">augment_hsv</span><span class="p">(</span><span class="n">img</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">,</span> <span class="n">hgain</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">sgain</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">vgain</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">bgr_channels</span><span class="o">=</span><span class="p">(</span><span class="mi">0</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="n">hsv_augs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="o">-</span><span class="mi">1</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="o">*</span> <span class="p">[</span><span class="n">hgain</span><span class="p">,</span> <span class="n">sgain</span><span class="p">,</span> <span class="n">vgain</span><span class="p">]</span>  <span class="c1"># random gains</span>
    <span class="n">hsv_augs</span> <span class="o">*=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">randint</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>  <span class="c1"># random selection of h, s, v</span>
    <span class="n">hsv_augs</span> <span class="o">=</span> <span class="n">hsv_augs</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">int16</span><span class="p">)</span>
    <span class="n">img_hsv</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">cvtColor</span><span class="p">(</span><span class="n">img</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="n">bgr_channels</span><span class="p">],</span> <span class="n">cv2</span><span class="o">.</span><span class="n">COLOR_BGR2HSV</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">int16</span><span class="p">)</span>

    <span class="n">img_hsv</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">img_hsv</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">hsv_augs</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">%</span> <span class="mi">180</span>
    <span class="n">img_hsv</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">img_hsv</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">hsv_augs</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">255</span><span class="p">)</span>
    <span class="n">img_hsv</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">img_hsv</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">hsv_augs</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">255</span><span class="p">)</span>

    <span class="n">img</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="n">bgr_channels</span><span class="p">]</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">cvtColor</span><span class="p">(</span><span class="n">img_hsv</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> <span class="n">cv2</span><span class="o">.</span><span class="n">COLOR_HSV2BGR</span><span class="p">)</span>  <span class="c1"># no return needed</span>


<span class="k">def</span> <span class="nf">rescale_and_pad_to_size</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="n">input_size</span><span class="p">,</span> <span class="n">swap</span><span class="o">=</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">pad_val</span><span class="o">=</span><span class="mi">114</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Rescales image according to minimum ratio between the target height /image height, target width / image width,</span>
<span class="sd">    and pads the image to the target size.</span>

<span class="sd">    :param img: Image to be rescaled</span>
<span class="sd">    :param input_size: Target size</span>
<span class="sd">    :param swap: Axis&#39;s to be rearranged.</span>
<span class="sd">    :return: rescaled image, ratio</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
        <span class="n">padded_img</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="n">input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span> <span class="o">*</span> <span class="n">pad_val</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">padded_img</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">input_size</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span> <span class="o">*</span> <span class="n">pad_val</span>

    <span class="n">r</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">input_size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">input_size</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">/</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
    <span class="n">resized_img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span>
        <span class="n">img</span><span class="p">,</span>
        <span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">r</span><span class="p">),</span> <span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">r</span><span class="p">)),</span>
        <span class="n">interpolation</span><span class="o">=</span><span class="n">cv2</span><span class="o">.</span><span class="n">INTER_LINEAR</span><span class="p">,</span>
    <span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span>
    <span class="n">padded_img</span><span class="p">[:</span> <span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">r</span><span class="p">),</span> <span class="p">:</span> <span class="nb">int</span><span class="p">(</span><span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">r</span><span class="p">)]</span> <span class="o">=</span> <span class="n">resized_img</span>

    <span class="n">padded_img</span> <span class="o">=</span> <span class="n">padded_img</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="n">swap</span><span class="p">)</span>
    <span class="n">padded_img</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ascontiguousarray</span><span class="p">(</span><span class="n">padded_img</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">padded_img</span><span class="p">,</span> <span class="n">r</span>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.callbacks &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.callbacks</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.callbacks</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">copy</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">time</span>
<span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
<span class="kn">import</span> <span class="nn">math</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">onnx</span>
<span class="kn">import</span> <span class="nn">onnxruntime</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">signal</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>

<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">DetectionVisualization</span><span class="p">,</span> <span class="n">DetectionPostPredictionCallback</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.segmentation_utils</span> <span class="kn">import</span> <span class="n">BinarySegmentationVisualization</span>
<span class="kn">import</span> <span class="nn">cv2</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>

<span class="k">try</span><span class="p">:</span>
    <span class="kn">from</span> <span class="nn">deci_lab_client.client</span> <span class="kn">import</span> <span class="n">DeciPlatformClient</span>
    <span class="kn">from</span> <span class="nn">deci_lab_client.models</span> <span class="kn">import</span> <span class="n">ModelBenchmarkState</span>
    <span class="kn">from</span> <span class="nn">deci_lab_client.models.model_metadata</span> <span class="kn">import</span> <span class="n">ModelMetadata</span>

    <span class="n">_imported_deci_lab_failure</span> <span class="o">=</span> <span class="kc">None</span>
<span class="k">except</span> <span class="p">(</span><span class="ne">ImportError</span><span class="p">,</span> <span class="ne">NameError</span><span class="p">,</span> <span class="ne">ModuleNotFoundError</span><span class="p">)</span> <span class="k">as</span> <span class="n">import_err</span><span class="p">:</span>
    <span class="n">logger</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;Failed to import deci_lab_client&quot;</span><span class="p">)</span>
    <span class="n">_imported_deci_lab_failure</span> <span class="o">=</span> <span class="n">import_err</span>


<div class="viewcode-block" id="Phase"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.Phase">[docs]</a><span class="k">class</span> <span class="nc">Phase</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
    <span class="n">PRE_TRAINING</span> <span class="o">=</span> <span class="s2">&quot;PRE_TRAINING&quot;</span>
    <span class="n">TRAIN_BATCH_END</span> <span class="o">=</span> <span class="s2">&quot;TRAIN_BATCH_END&quot;</span>
    <span class="n">TRAIN_BATCH_STEP</span> <span class="o">=</span> <span class="s2">&quot;TRAIN_BATCH_STEP&quot;</span>
    <span class="n">TRAIN_EPOCH_START</span> <span class="o">=</span> <span class="s2">&quot;TRAIN_EPOCH_START&quot;</span>
    <span class="n">TRAIN_EPOCH_END</span> <span class="o">=</span> <span class="s2">&quot;TRAIN_EPOCH_END&quot;</span>
    <span class="n">VALIDATION_BATCH_END</span> <span class="o">=</span> <span class="s2">&quot;VALIDATION_BATCH_END&quot;</span>
    <span class="n">VALIDATION_EPOCH_END</span> <span class="o">=</span> <span class="s2">&quot;VALIDATION_EPOCH_END&quot;</span>
    <span class="n">VALIDATION_END_BEST_EPOCH</span> <span class="o">=</span> <span class="s2">&quot;VALIDATION_END_BEST_EPOCH&quot;</span>
    <span class="n">TEST_BATCH_END</span> <span class="o">=</span> <span class="s2">&quot;TEST_BATCH_END&quot;</span>
    <span class="n">TEST_END</span> <span class="o">=</span> <span class="s2">&quot;TEST_END&quot;</span>
    <span class="n">POST_TRAINING</span> <span class="o">=</span> <span class="s2">&quot;POST_TRAINING&quot;</span></div>


<div class="viewcode-block" id="ContextSgMethods"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.ContextSgMethods">[docs]</a><span class="k">class</span> <span class="nc">ContextSgMethods</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Class for delegating SgModel&#39;s methods, so that only the relevant ones are (&quot;phase wise&quot;) are accessible.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">methods</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">attr</span><span class="p">,</span> <span class="n">attr_val</span> <span class="ow">in</span> <span class="n">methods</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">attr</span><span class="p">,</span> <span class="n">attr_val</span><span class="p">)</span></div>


<div class="viewcode-block" id="PhaseContext"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.PhaseContext">[docs]</a><span class="k">class</span> <span class="nc">PhaseContext</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Represents the input for phase callbacks, and is constantly updated after callback calls.</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">epoch</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">batch_idx</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">optimizer</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">metrics_dict</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">inputs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">preds</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">metrics_compute_fn</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">loss_avg_meter</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">loss_log_items</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">criterion</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">device</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">experiment_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">ckpt_dir</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">net</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">lr_warmup_epochs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">sg_logger</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">train_loader</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">valid_loader</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">training_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">ddp_silent_mode</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">checkpoint_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">architecture</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">arch_params</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">metric_idx_in_results_tuple</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">metric_to_watch</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">valid_metrics</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">context_methods</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">epoch</span> <span class="o">=</span> <span class="n">epoch</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span> <span class="o">=</span> <span class="n">batch_idx</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">optimizer</span> <span class="o">=</span> <span class="n">optimizer</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">inputs</span> <span class="o">=</span> <span class="n">inputs</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">preds</span> <span class="o">=</span> <span class="n">preds</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">target</span> <span class="o">=</span> <span class="n">target</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">metrics_dict</span> <span class="o">=</span> <span class="n">metrics_dict</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">metrics_compute_fn</span> <span class="o">=</span> <span class="n">metrics_compute_fn</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">loss_avg_meter</span> <span class="o">=</span> <span class="n">loss_avg_meter</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">loss_log_items</span> <span class="o">=</span> <span class="n">loss_log_items</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">criterion</span> <span class="o">=</span> <span class="n">criterion</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">=</span> <span class="n">device</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">stop_training</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">experiment_name</span> <span class="o">=</span> <span class="n">experiment_name</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_dir</span> <span class="o">=</span> <span class="n">ckpt_dir</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">net</span> <span class="o">=</span> <span class="n">net</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">=</span> <span class="n">lr_warmup_epochs</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sg_logger</span> <span class="o">=</span> <span class="n">sg_logger</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader</span> <span class="o">=</span> <span class="n">train_loader</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valid_loader</span> <span class="o">=</span> <span class="n">valid_loader</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span> <span class="o">=</span> <span class="n">training_params</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ddp_silent_mode</span> <span class="o">=</span> <span class="n">ddp_silent_mode</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">checkpoint_params</span> <span class="o">=</span> <span class="n">checkpoint_params</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">architecture</span> <span class="o">=</span> <span class="n">architecture</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">arch_params</span> <span class="o">=</span> <span class="n">arch_params</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">metric_idx_in_results_tuple</span> <span class="o">=</span> <span class="n">metric_idx_in_results_tuple</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">metric_to_watch</span> <span class="o">=</span> <span class="n">metric_to_watch</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">valid_metrics</span> <span class="o">=</span> <span class="n">valid_metrics</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">context_methods</span> <span class="o">=</span> <span class="n">context_methods</span>

<div class="viewcode-block" id="PhaseContext.update_context"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.PhaseContext.update_context">[docs]</a>    <span class="k">def</span> <span class="nf">update_context</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">attr</span><span class="p">,</span> <span class="n">attr_val</span> <span class="ow">in</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">attr</span><span class="p">,</span> <span class="n">attr_val</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="PhaseCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.PhaseCallback">[docs]</a><span class="k">class</span> <span class="nc">PhaseCallback</span><span class="p">:</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">phase</span> <span class="o">=</span> <span class="n">phase</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span></div>


<div class="viewcode-block" id="ModelConversionCheckCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.ModelConversionCheckCallback">[docs]</a><span class="k">class</span> <span class="nc">ModelConversionCheckCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Pre-training callback that verifies model conversion to onnx given specified conversion parameters.</span>

<span class="sd">    The model is converted, then inference is applied with onnx runtime.</span>

<span class="sd">    Use this callback wit hthe same args as DeciPlatformCallback to prevent conversion fails at the end of training.</span>

<span class="sd">    Attributes:</span>

<span class="sd">        model_meta_data: (ModelMetadata) model&#39;s meta-data object.</span>

<span class="sd">        The following parameters may be passed as kwargs in order to control the conversion to onnx:</span>
<span class="sd">        :param opset_version (default=11)</span>
<span class="sd">        :param do_constant_folding (default=True)</span>
<span class="sd">        :param dynamic_axes (default=</span>
<span class="sd">                                        {&#39;input&#39;: {0: &#39;batch_size&#39;},</span>
<span class="sd">                                        # Variable length axes</span>
<span class="sd">                                        &#39;output&#39;: {0: &#39;batch_size&#39;}}</span>

<span class="sd">                                        )</span>
<span class="sd">        :param input_names (default=[&quot;input&quot;])</span>
<span class="sd">        :param output_names (default=[&quot;output&quot;])</span>
<span class="sd">        :param rtol (default=1e-03)</span>
<span class="sd">        :param atol (default=1e-05)</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_meta_data</span><span class="p">:</span> <span class="n">ModelMetadata</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ModelConversionCheckCallback</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><span class="n">phase</span><span class="o">=</span><span class="n">Phase</span><span class="o">.</span><span class="n">PRE_TRAINING</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span> <span class="o">=</span> <span class="n">model_meta_data</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">opset_version</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;opset_version&quot;</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">do_constant_folding</span> <span class="o">=</span> <span class="p">(</span>
            <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;do_constant_folding&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span> <span class="k">if</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;do_constant_folding&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span> <span class="k">else</span> <span class="kc">True</span>
        <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">input_names</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;input_names&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="p">[</span><span class="s2">&quot;input&quot;</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">output_names</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;output_names&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="p">[</span><span class="s2">&quot;output&quot;</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dynamic_axes</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;dynamic_axes&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="p">{</span><span class="s2">&quot;input&quot;</span><span class="p">:</span> <span class="p">{</span><span class="mi">0</span><span class="p">:</span> <span class="s2">&quot;batch_size&quot;</span><span class="p">},</span> <span class="s2">&quot;output&quot;</span><span class="p">:</span> <span class="p">{</span><span class="mi">0</span><span class="p">:</span> <span class="s2">&quot;batch_size&quot;</span><span class="p">}}</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;rtol&quot;</span><span class="p">,</span> <span class="mf">1e-03</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">atol</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;atol&quot;</span><span class="p">,</span> <span class="mf">1e-05</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="n">model</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="p">)</span>
        <span class="n">model</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span>
        <span class="n">model</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>  <span class="c1"># Put model into eval mode</span>

        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;prep_model_for_conversion&quot;</span><span class="p">):</span>
            <span class="n">model</span><span class="o">.</span><span class="n">prep_model_for_conversion</span><span class="p">(</span><span class="n">input_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span><span class="o">.</span><span class="n">input_dimensions</span><span class="p">)</span>

        <span class="n">x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span><span class="o">.</span><span class="n">primary_batch_size</span><span class="p">,</span> <span class="o">*</span><span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span><span class="o">.</span><span class="n">input_dimensions</span><span class="p">,</span> <span class="n">requires_grad</span><span class="o">=</span><span class="kc">False</span>
        <span class="p">)</span>

        <span class="n">tmp_model_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">ckpt_dir</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span><span class="o">.</span><span class="n">name</span> <span class="o">+</span> <span class="s2">&quot;_tmp.onnx&quot;</span><span class="p">)</span>

        <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">():</span>
            <span class="n">torch_out</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>

        <span class="n">torch</span><span class="o">.</span><span class="n">onnx</span><span class="o">.</span><span class="n">export</span><span class="p">(</span>
            <span class="n">model</span><span class="p">,</span>  <span class="c1"># Model being run</span>
            <span class="n">x</span><span class="p">,</span>  <span class="c1"># Model input (or a tuple for multiple inputs)</span>
            <span class="n">tmp_model_path</span><span class="p">,</span>  <span class="c1"># Where to save the model (can be a file or file-like object)</span>
            <span class="n">export_params</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>  <span class="c1"># Store the trained parameter weights inside the model file</span>
            <span class="n">opset_version</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">opset_version</span><span class="p">,</span>
            <span class="n">do_constant_folding</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">do_constant_folding</span><span class="p">,</span>
            <span class="n">input_names</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">input_names</span><span class="p">,</span>
            <span class="n">output_names</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">output_names</span><span class="p">,</span>
            <span class="n">dynamic_axes</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dynamic_axes</span><span class="p">,</span>
        <span class="p">)</span>

        <span class="n">onnx_model</span> <span class="o">=</span> <span class="n">onnx</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">tmp_model_path</span><span class="p">)</span>
        <span class="n">onnx</span><span class="o">.</span><span class="n">checker</span><span class="o">.</span><span class="n">check_model</span><span class="p">(</span><span class="n">onnx_model</span><span class="p">)</span>

        <span class="n">ort_session</span> <span class="o">=</span> <span class="n">onnxruntime</span><span class="o">.</span><span class="n">InferenceSession</span><span class="p">(</span>
            <span class="n">tmp_model_path</span><span class="p">,</span> <span class="n">providers</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;CUDAExecutionProvider&quot;</span><span class="p">,</span> <span class="s2">&quot;CPUExecutionProvider&quot;</span><span class="p">]</span>
        <span class="p">)</span>

        <span class="c1"># compute ONNX Runtime output prediction</span>
        <span class="n">ort_inputs</span> <span class="o">=</span> <span class="p">{</span><span class="n">ort_session</span><span class="o">.</span><span class="n">get_inputs</span><span class="p">()[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">name</span><span class="p">:</span> <span class="n">x</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>
        <span class="n">ort_outs</span> <span class="o">=</span> <span class="n">ort_session</span><span class="o">.</span><span class="n">run</span><span class="p">(</span><span class="kc">None</span><span class="p">,</span> <span class="n">ort_inputs</span><span class="p">)</span>

        <span class="c1"># TODO: Ideally we don&#39;t want to check this but have the certainty of just calling torch_out.cpu()</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">torch_out</span><span class="p">,</span> <span class="n">List</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">torch_out</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="n">torch_out</span> <span class="o">=</span> <span class="n">torch_out</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="c1"># compare ONNX Runtime and PyTorch results</span>
        <span class="n">np</span><span class="o">.</span><span class="n">testing</span><span class="o">.</span><span class="n">assert_allclose</span><span class="p">(</span><span class="n">torch_out</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> <span class="n">ort_outs</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">rtol</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">rtol</span><span class="p">,</span> <span class="n">atol</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">atol</span><span class="p">)</span>

        <span class="n">os</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">tmp_model_path</span><span class="p">)</span>

        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Exported model has been tested with ONNXRuntime, and the result looks good!&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="DeciLabUploadCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.DeciLabUploadCallback">[docs]</a><span class="k">class</span> <span class="nc">DeciLabUploadCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Post-training callback for uploading and optimizing a model.</span>

<span class="sd">    Attributes:</span>

<span class="sd">        email: (str) username for Deci platform.</span>
<span class="sd">        model_meta_data: (ModelMetadata) model&#39;s meta-data object.</span>
<span class="sd">        optimization_request_form: (dict) optimization request form object.</span>
<span class="sd">        password: (str) default=None, should only be used for testing.</span>
<span class="sd">        ckpt_name: (str) default=&quot;ckpt_best&quot; refers to the filename of the checkpoint, inside the checkpoint directory.</span>

<span class="sd">        The following parameters may be passed as kwargs in order to control the conversion to onnx:</span>
<span class="sd">        :param opset_version</span>
<span class="sd">        :param do_constant_folding</span>
<span class="sd">        :param dynamic_axes</span>
<span class="sd">        :param input_names</span>
<span class="sd">        :param output_names</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_meta_data</span><span class="p">,</span> <span class="n">optimization_request_form</span><span class="p">,</span> <span class="n">auth_token</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">ckpt_name</span><span class="o">=</span><span class="s2">&quot;ckpt_best.pth&quot;</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">phase</span><span class="o">=</span><span class="n">Phase</span><span class="o">.</span><span class="n">POST_TRAINING</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">_imported_deci_lab_failure</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="n">_imported_deci_lab_failure</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span> <span class="o">=</span> <span class="n">model_meta_data</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">optimization_request_form</span> <span class="o">=</span> <span class="n">optimization_request_form</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">conversion_kwargs</span> <span class="o">=</span> <span class="n">kwargs</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_name</span> <span class="o">=</span> <span class="n">ckpt_name</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">platform_client</span> <span class="o">=</span> <span class="n">DeciPlatformClient</span><span class="p">(</span><span class="s2">&quot;api.deci.ai&quot;</span><span class="p">,</span> <span class="mi">443</span><span class="p">,</span> <span class="n">https</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">platform_client</span><span class="o">.</span><span class="n">login</span><span class="p">(</span><span class="n">token</span><span class="o">=</span><span class="n">auth_token</span><span class="p">)</span>

<div class="viewcode-block" id="DeciLabUploadCallback.log_optimization_failed"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.DeciLabUploadCallback.log_optimization_failed">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">log_optimization_failed</span><span class="p">():</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;We couldn&#39;t finish your model optimization. Visit https://console.deci.ai for details&quot;</span><span class="p">)</span></div>

<div class="viewcode-block" id="DeciLabUploadCallback.upload_model"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.DeciLabUploadCallback.upload_model">[docs]</a>    <span class="k">def</span> <span class="nf">upload_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">            This function will upload the trained model to the Deci Lab</span>

<span class="sd">            Args:</span>
<span class="sd">                model: The resulting model from the training process</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">platform_client</span><span class="o">.</span><span class="n">add_model</span><span class="p">(</span>
            <span class="n">add_model_request</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span><span class="p">,</span>
            <span class="n">optimization_request</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">optimization_request_form</span><span class="p">,</span>
            <span class="n">local_loaded_model</span><span class="o">=</span><span class="n">model</span><span class="p">,</span>
        <span class="p">)</span></div>

<div class="viewcode-block" id="DeciLabUploadCallback.get_optimization_status"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.DeciLabUploadCallback.get_optimization_status">[docs]</a>    <span class="k">def</span> <span class="nf">get_optimization_status</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">optimized_model_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">            This function will do fetch the optimized version of the trained model and check on its benchmark status.</span>
<span class="sd">            The status will be checked against the server every 30 seconds and the process will timeout after 30 minutes</span>
<span class="sd">            or log about the successful optimization - whichever happens first.</span>
<span class="sd">            Args:</span>
<span class="sd">                optimized_model_name (str): Optimized model name</span>
<span class="sd">            Returns:</span>
<span class="sd">                bool: whether or not the optimized model has been benchmarked</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">def</span> <span class="nf">handler</span><span class="p">(</span><span class="n">_signum</span><span class="p">,</span> <span class="n">_frame</span><span class="p">):</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="s2">&quot;Process timed out. Visit https://console.deci.ai for details&quot;</span><span class="p">)</span>
            <span class="k">return</span> <span class="kc">False</span>

        <span class="n">signal</span><span class="o">.</span><span class="n">signal</span><span class="p">(</span><span class="n">signal</span><span class="o">.</span><span class="n">SIGALRM</span><span class="p">,</span> <span class="n">handler</span><span class="p">)</span>
        <span class="n">signal</span><span class="o">.</span><span class="n">alarm</span><span class="p">(</span><span class="mi">1800</span><span class="p">)</span>

        <span class="n">finished</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">while</span> <span class="ow">not</span> <span class="n">finished</span><span class="p">:</span>
            <span class="n">optimized_model</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">platform_client</span><span class="o">.</span><span class="n">get_model_by_name</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">optimized_model_name</span><span class="p">)</span><span class="o">.</span><span class="n">data</span>
            <span class="k">if</span> <span class="n">optimized_model</span><span class="o">.</span><span class="n">benchmark_state</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">[</span><span class="n">ModelBenchmarkState</span><span class="o">.</span><span class="n">IN_PROGRESS</span><span class="p">,</span> <span class="n">ModelBenchmarkState</span><span class="o">.</span><span class="n">PENDING</span><span class="p">]:</span>
                <span class="n">finished</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">30</span><span class="p">)</span>

        <span class="n">signal</span><span class="o">.</span><span class="n">alarm</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
        <span class="k">return</span> <span class="kc">True</span></div>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">            This function will attempt to upload the trained model and schedule an optimization for it.</span>
<span class="sd">            Args:</span>
<span class="sd">                context (PhaseContext): Training phase context</span>
<span class="sd">            Returns:</span>
<span class="sd">                bool: whether or not the optimized model has been benchmarked</span>
<span class="sd">            &quot;&quot;&quot;</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">model</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">net</span><span class="p">)</span>
            <span class="n">model_state_dict_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">ckpt_dir</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">ckpt_name</span><span class="p">)</span>
            <span class="n">model_state_dict</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">model_state_dict_path</span><span class="p">)[</span><span class="s2">&quot;net&quot;</span><span class="p">]</span>
            <span class="n">model</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">state_dict</span><span class="o">=</span><span class="n">model_state_dict</span><span class="p">)</span>

            <span class="n">model</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;prep_model_for_conversion&quot;</span><span class="p">):</span>
                <span class="n">model</span><span class="o">.</span><span class="n">prep_model_for_conversion</span><span class="p">(</span><span class="n">input_size</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span><span class="o">.</span><span class="n">input_dimensions</span><span class="p">)</span>

            <span class="bp">self</span><span class="o">.</span><span class="n">upload_model</span><span class="p">(</span><span class="n">model</span><span class="o">=</span><span class="n">model</span><span class="p">)</span>
            <span class="n">model_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">model_meta_data</span><span class="o">.</span><span class="n">name</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Successfully added </span><span class="si">{</span><span class="n">model_name</span><span class="si">}</span><span class="s2"> to the model repository&quot;</span><span class="p">)</span>

            <span class="n">optimized_model_name</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">model_name</span><span class="si">}</span><span class="s2">_1_1&quot;</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;We&#39;ll wait for the scheduled optimization to finish. Please don&#39;t close this window&quot;</span><span class="p">)</span>
            <span class="n">success</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_optimization_status</span><span class="p">(</span><span class="n">optimized_model_name</span><span class="o">=</span><span class="n">optimized_model_name</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">success</span><span class="p">:</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Successfully finished your model optimization. Visit https://console.deci.ai for details&quot;</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">DeciLabUploadCallback</span><span class="o">.</span><span class="n">log_optimization_failed</span><span class="p">()</span>
        <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">ex</span><span class="p">:</span>
            <span class="n">DeciLabUploadCallback</span><span class="o">.</span><span class="n">log_optimization_failed</span><span class="p">()</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="n">ex</span><span class="p">)</span></div>


<div class="viewcode-block" id="LRCallbackBase"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.LRCallbackBase">[docs]</a><span class="k">class</span> <span class="nc">LRCallbackBase</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Base class for hard coded learning rate scheduling regimes, implemented as callbacks.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">,</span> <span class="n">initial_lr</span><span class="p">,</span> <span class="n">update_param_groups</span><span class="p">,</span> <span class="n">train_loader_len</span><span class="p">,</span> <span class="n">net</span><span class="p">,</span> <span class="n">training_params</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">LRCallbackBase</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><span class="n">phase</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">=</span> <span class="n">initial_lr</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr</span> <span class="o">=</span> <span class="n">initial_lr</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">update_param_groups</span> <span class="o">=</span> <span class="n">update_param_groups</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span> <span class="o">=</span> <span class="n">train_loader_len</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">net</span> <span class="o">=</span> <span class="n">net</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span> <span class="o">=</span> <span class="n">training_params</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_lr_scheduling_enabled</span><span class="p">(</span><span class="n">context</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">perform_scheduling</span><span class="p">(</span><span class="n">context</span><span class="p">)</span>

<div class="viewcode-block" id="LRCallbackBase.is_lr_scheduling_enabled"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.LRCallbackBase.is_lr_scheduling_enabled">[docs]</a>    <span class="k">def</span> <span class="nf">is_lr_scheduling_enabled</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Predicate that controls whether to perform lr scheduling based on values in context.</span>

<span class="sd">        @param context: PhaseContext: current phase&#39;s context.</span>
<span class="sd">        @return: bool, whether to apply lr scheduling or not.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

<div class="viewcode-block" id="LRCallbackBase.perform_scheduling"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.LRCallbackBase.perform_scheduling">[docs]</a>    <span class="k">def</span> <span class="nf">perform_scheduling</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Performs lr scheduling based on values in context.</span>

<span class="sd">        @param context: PhaseContext: current phase&#39;s context.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

<div class="viewcode-block" id="LRCallbackBase.update_lr"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.LRCallbackBase.update_lr">[docs]</a>    <span class="k">def</span> <span class="nf">update_lr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">,</span> <span class="n">epoch</span><span class="p">,</span> <span class="n">batch_idx</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">update_param_groups</span><span class="p">:</span>
            <span class="n">param_groups</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">update_param_groups</span><span class="p">(</span>
                <span class="n">optimizer</span><span class="o">.</span><span class="n">param_groups</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">lr</span><span class="p">,</span> <span class="n">epoch</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span>
            <span class="p">)</span>
            <span class="n">optimizer</span><span class="o">.</span><span class="n">param_groups</span> <span class="o">=</span> <span class="n">param_groups</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># UPDATE THE OPTIMIZERS PARAMETER</span>
            <span class="k">for</span> <span class="n">param_group</span> <span class="ow">in</span> <span class="n">optimizer</span><span class="o">.</span><span class="n">param_groups</span><span class="p">:</span>
                <span class="n">param_group</span><span class="p">[</span><span class="s2">&quot;lr&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">lr</span></div></div>


<div class="viewcode-block" id="WarmupLRCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.WarmupLRCallback">[docs]</a><span class="k">class</span> <span class="nc">WarmupLRCallback</span><span class="p">(</span><span class="n">LRCallbackBase</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    LR scheduling callback for linear step warmup.</span>
<span class="sd">    LR climbs from warmup_initial_lr with even steps to initial lr. When warmup_initial_lr is None- LR climb starts from</span>
<span class="sd">     initial_lr/(1+warmup_epochs).</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">WarmupLRCallback</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><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_EPOCH_START</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">warmup_initial_lr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">warmup_initial_lr</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">/</span> <span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">+</span> <span class="mi">1</span>
        <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">warmup_step_size</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">warmup_initial_lr</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span>

<div class="viewcode-block" id="WarmupLRCallback.perform_scheduling"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.WarmupLRCallback.perform_scheduling">[docs]</a>    <span class="k">def</span> <span class="nf">perform_scheduling</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">warmup_initial_lr</span> <span class="o">+</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">warmup_step_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">update_lr</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span></div>

<div class="viewcode-block" id="WarmupLRCallback.is_lr_scheduling_enabled"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.WarmupLRCallback.is_lr_scheduling_enabled">[docs]</a>    <span class="k">def</span> <span class="nf">is_lr_scheduling_enabled</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">&gt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span></div></div>


<div class="viewcode-block" id="StepLRCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.StepLRCallback">[docs]</a><span class="k">class</span> <span class="nc">StepLRCallback</span><span class="p">(</span><span class="n">LRCallbackBase</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Hard coded step learning rate scheduling (i.e at specific milestones).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">lr_updates</span><span class="p">,</span> <span class="n">lr_decay_factor</span><span class="p">,</span> <span class="n">step_lr_update_freq</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">StepLRCallback</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><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_EPOCH_END</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">step_lr_update_freq</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">lr_updates</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                <span class="s2">&quot;Only one of [lr_updates, step_lr_update_freq] should be passed to StepLRCallback constructor&quot;</span>
            <span class="p">)</span>

        <span class="k">if</span> <span class="n">step_lr_update_freq</span><span class="p">:</span>
            <span class="n">max_epochs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
            <span class="n">warmup_epochs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span>
            <span class="n">lr_updates</span> <span class="o">=</span> <span class="p">[</span>
                <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">step_lr_update_freq</span> <span class="o">*</span> <span class="n">x</span><span class="p">))</span>
                <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">max_epochs</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">warmup_epochs</span> <span class="o">&lt;=</span> <span class="nb">int</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">step_lr_update_freq</span> <span class="o">*</span> <span class="n">x</span><span class="p">))</span> <span class="o">&lt;</span> <span class="n">max_epochs</span>
            <span class="p">]</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
                <span class="s2">&quot;Specific lr_updates were passed along with cooldown_epochs &gt; 0,&quot;</span> <span class="s2">&quot; cooldown will have no effect.&quot;</span>
            <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr_updates</span> <span class="o">=</span> <span class="n">lr_updates</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr_decay_factor</span> <span class="o">=</span> <span class="n">lr_decay_factor</span>

<div class="viewcode-block" id="StepLRCallback.perform_scheduling"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.StepLRCallback.perform_scheduling">[docs]</a>    <span class="k">def</span> <span class="nf">perform_scheduling</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">num_updates_passed</span> <span class="o">=</span> <span class="p">[</span><span class="n">x</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">lr_updates</span> <span class="k">if</span> <span class="n">x</span> <span class="o">&lt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">lr_decay_factor</span> <span class="o">**</span> <span class="nb">len</span><span class="p">(</span><span class="n">num_updates_passed</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">update_lr</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span></div>

<div class="viewcode-block" id="StepLRCallback.is_lr_scheduling_enabled"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.StepLRCallback.is_lr_scheduling_enabled">[docs]</a>    <span class="k">def</span> <span class="nf">is_lr_scheduling_enabled</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">&lt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span></div></div>


<div class="viewcode-block" id="ExponentialLRCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.ExponentialLRCallback">[docs]</a><span class="k">class</span> <span class="nc">ExponentialLRCallback</span><span class="p">(</span><span class="n">LRCallbackBase</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Exponential decay learning rate scheduling. Decays the learning rate by `lr_decay_factor` every epoch.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">lr_decay_factor</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">phase</span><span class="o">=</span><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_STEP</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr_decay_factor</span> <span class="o">=</span> <span class="n">lr_decay_factor</span>

<div class="viewcode-block" id="ExponentialLRCallback.perform_scheduling"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.ExponentialLRCallback.perform_scheduling">[docs]</a>    <span class="k">def</span> <span class="nf">perform_scheduling</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">effective_epoch</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span>
        <span class="n">current_iter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span> <span class="o">*</span> <span class="n">effective_epoch</span> <span class="o">+</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">lr_decay_factor</span> <span class="o">**</span> <span class="p">(</span><span class="n">current_iter</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">update_lr</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">)</span></div>

<div class="viewcode-block" id="ExponentialLRCallback.is_lr_scheduling_enabled"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.ExponentialLRCallback.is_lr_scheduling_enabled">[docs]</a>    <span class="k">def</span> <span class="nf">is_lr_scheduling_enabled</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">post_warmup_epochs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">&lt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">&lt;</span> <span class="n">post_warmup_epochs</span></div></div>


<div class="viewcode-block" id="PolyLRCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.PolyLRCallback">[docs]</a><span class="k">class</span> <span class="nc">PolyLRCallback</span><span class="p">(</span><span class="n">LRCallbackBase</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Hard coded polynomial decay learning rate scheduling (i.e at specific milestones).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">max_epochs</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">PolyLRCallback</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><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_STEP</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">=</span> <span class="n">max_epochs</span>

<div class="viewcode-block" id="PolyLRCallback.perform_scheduling"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.PolyLRCallback.perform_scheduling">[docs]</a>    <span class="k">def</span> <span class="nf">perform_scheduling</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">effective_epoch</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span>
        <span class="n">effective_max_epochs</span> <span class="o">=</span> <span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
        <span class="p">)</span>
        <span class="n">current_iter</span> <span class="o">=</span> <span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span> <span class="o">*</span> <span class="n">effective_epoch</span> <span class="o">+</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span>
        <span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">batch_accumulate</span>
        <span class="n">max_iter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span> <span class="o">*</span> <span class="n">effective_max_epochs</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">batch_accumulate</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">*</span> <span class="nb">pow</span><span class="p">((</span><span class="mf">1.0</span> <span class="o">-</span> <span class="p">(</span><span class="n">current_iter</span> <span class="o">/</span> <span class="n">max_iter</span><span class="p">)),</span> <span class="mf">0.9</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">update_lr</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">)</span></div>

<div class="viewcode-block" id="PolyLRCallback.is_lr_scheduling_enabled"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.PolyLRCallback.is_lr_scheduling_enabled">[docs]</a>    <span class="k">def</span> <span class="nf">is_lr_scheduling_enabled</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">post_warmup_epochs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">&lt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">&lt;</span> <span class="n">post_warmup_epochs</span></div></div>


<div class="viewcode-block" id="CosineLRCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.CosineLRCallback">[docs]</a><span class="k">class</span> <span class="nc">CosineLRCallback</span><span class="p">(</span><span class="n">LRCallbackBase</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Hard coded step Cosine anealing learning rate scheduling.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">max_epochs</span><span class="p">,</span> <span class="n">cosine_final_lr_ratio</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">CosineLRCallback</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><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_STEP</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">=</span> <span class="n">max_epochs</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cosine_final_lr_ratio</span> <span class="o">=</span> <span class="n">cosine_final_lr_ratio</span>

<div class="viewcode-block" id="CosineLRCallback.perform_scheduling"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.CosineLRCallback.perform_scheduling">[docs]</a>    <span class="k">def</span> <span class="nf">perform_scheduling</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">effective_epoch</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span>
        <span class="n">effective_max_epochs</span> <span class="o">=</span> <span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
        <span class="p">)</span>
        <span class="n">current_iter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span> <span class="o">*</span> <span class="n">effective_epoch</span> <span class="o">+</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span>
        <span class="n">max_iter</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span> <span class="o">*</span> <span class="n">effective_max_epochs</span>
        <span class="n">lr</span> <span class="o">=</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">*</span> <span class="p">(</span><span class="mf">1.0</span> <span class="o">+</span> <span class="n">math</span><span class="o">.</span><span class="n">cos</span><span class="p">(</span><span class="n">current_iter</span> <span class="o">/</span> <span class="p">(</span><span class="n">max_iter</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span> <span class="o">*</span> <span class="n">math</span><span class="o">.</span><span class="n">pi</span><span class="p">))</span>
        <span class="c1"># the cosine starts from initial_lr and reaches initial_lr * cosine_final_lr_ratio in last epoch</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr</span> <span class="o">=</span> <span class="n">lr</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">cosine_final_lr_ratio</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">cosine_final_lr_ratio</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">update_lr</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">)</span></div>

<div class="viewcode-block" id="CosineLRCallback.is_lr_scheduling_enabled"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.CosineLRCallback.is_lr_scheduling_enabled">[docs]</a>    <span class="k">def</span> <span class="nf">is_lr_scheduling_enabled</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">post_warmup_epochs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">&lt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">&lt;</span> <span class="n">post_warmup_epochs</span></div></div>


<div class="viewcode-block" id="FunctionLRCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.FunctionLRCallback">[docs]</a><span class="k">class</span> <span class="nc">FunctionLRCallback</span><span class="p">(</span><span class="n">LRCallbackBase</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Hard coded rate scheduling for user defined lr scheduling function.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">max_epochs</span><span class="p">,</span> <span class="n">lr_schedule_function</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">FunctionLRCallback</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><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_BATCH_STEP</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="k">assert</span> <span class="n">callable</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lr_schedule_function</span><span class="p">),</span> <span class="s2">&quot;self.lr_function must be callable&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr_schedule_function</span> <span class="o">=</span> <span class="n">lr_schedule_function</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">=</span> <span class="n">max_epochs</span>

<div class="viewcode-block" id="FunctionLRCallback.is_lr_scheduling_enabled"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.FunctionLRCallback.is_lr_scheduling_enabled">[docs]</a>    <span class="k">def</span> <span class="nf">is_lr_scheduling_enabled</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">post_warmup_epochs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">&lt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">&lt;</span> <span class="n">post_warmup_epochs</span></div>

<div class="viewcode-block" id="FunctionLRCallback.perform_scheduling"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.FunctionLRCallback.perform_scheduling">[docs]</a>    <span class="k">def</span> <span class="nf">perform_scheduling</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
        <span class="n">effective_epoch</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span>
        <span class="n">effective_max_epochs</span> <span class="o">=</span> <span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">max_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">training_params</span><span class="o">.</span><span class="n">lr_cooldown_epochs</span>
        <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">lr_schedule_function</span><span class="p">(</span>
            <span class="n">initial_lr</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">initial_lr</span><span class="p">,</span>
            <span class="n">epoch</span><span class="o">=</span><span class="n">effective_epoch</span><span class="p">,</span>
            <span class="nb">iter</span><span class="o">=</span><span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">,</span>
            <span class="n">max_epoch</span><span class="o">=</span><span class="n">effective_max_epochs</span><span class="p">,</span>
            <span class="n">iters_per_epoch</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">train_loader_len</span><span class="p">,</span>
        <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">update_lr</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="IllegalLRSchedulerMetric"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.IllegalLRSchedulerMetric">[docs]</a><span class="k">class</span> <span class="nc">IllegalLRSchedulerMetric</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Exception raised illegal combination of training parameters.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        message -- explanation of the error</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">metric_name</span><span class="p">,</span> <span class="n">metrics_dict</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">message</span> <span class="o">=</span> <span class="p">(</span>
            <span class="s2">&quot;Illegal metric name: &quot;</span> <span class="o">+</span> <span class="n">metric_name</span> <span class="o">+</span> <span class="s2">&quot;. Expected one of metics_dics keys: &quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">metrics_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
        <span class="p">)</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">message</span><span class="p">)</span></div>


<div class="viewcode-block" id="LRSchedulerCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.LRSchedulerCallback">[docs]</a><span class="k">class</span> <span class="nc">LRSchedulerCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Learning rate scheduler callback.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        scheduler: torch.optim._LRScheduler, the learning rate scheduler to be called step() with.</span>
<span class="sd">        metric_name: str, (default=None) the metric name for ReduceLROnPlateau learning rate scheduler.</span>

<span class="sd">        When passing __call__ a metrics_dict, with a key=self.metric_name, the value of that metric will monitored</span>
<span class="sd">         for ReduceLROnPlateau (i.e step(metrics_dict[self.metric_name]).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">scheduler</span><span class="p">,</span> <span class="n">phase</span><span class="p">,</span> <span class="n">metric_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">LRSchedulerCallback</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><span class="n">phase</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">scheduler</span> <span class="o">=</span> <span class="n">scheduler</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">metric_name</span> <span class="o">=</span> <span class="n">metric_name</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">lr_warmup_epochs</span> <span class="o">&lt;=</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">:</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">metric_name</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">metric_name</span> <span class="ow">in</span> <span class="n">context</span><span class="o">.</span><span class="n">metrics_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">scheduler</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">metrics_dict</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_name</span><span class="p">])</span>
            <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">metric_name</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">scheduler</span><span class="o">.</span><span class="n">step</span><span class="p">()</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">raise</span> <span class="n">IllegalLRSchedulerMetric</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_name</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">metrics_dict</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="s2">&quot;LRSchedulerCallback: &quot;</span> <span class="o">+</span> <span class="nb">repr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">scheduler</span><span class="p">)</span></div>


<div class="viewcode-block" id="MetricsUpdateCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.MetricsUpdateCallback">[docs]</a><span class="k">class</span> <span class="nc">MetricsUpdateCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">MetricsUpdateCallback</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><span class="n">phase</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="n">context</span><span class="o">.</span><span class="n">metrics_compute_fn</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="o">**</span><span class="n">context</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">criterion</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">context</span><span class="o">.</span><span class="n">loss_avg_meter</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">loss_log_items</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">inputs</span><span class="p">))</span></div>


<div class="viewcode-block" id="KDModelMetricsUpdateCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.KDModelMetricsUpdateCallback">[docs]</a><span class="k">class</span> <span class="nc">KDModelMetricsUpdateCallback</span><span class="p">(</span><span class="n">MetricsUpdateCallback</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">phase</span><span class="o">=</span><span class="n">phase</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="n">metrics_compute_fn_kwargs</span> <span class="o">=</span> <span class="p">{</span><span class="n">k</span><span class="p">:</span> <span class="n">v</span><span class="o">.</span><span class="n">student_output</span> <span class="k">if</span> <span class="n">k</span> <span class="o">==</span> <span class="s2">&quot;preds&quot;</span> <span class="k">else</span> <span class="n">v</span> <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">context</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">items</span><span class="p">()}</span>
        <span class="n">context</span><span class="o">.</span><span class="n">metrics_compute_fn</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="o">**</span><span class="n">metrics_compute_fn_kwargs</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">criterion</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">context</span><span class="o">.</span><span class="n">loss_avg_meter</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">loss_log_items</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">inputs</span><span class="p">))</span></div>


<div class="viewcode-block" id="PhaseContextTestCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.PhaseContextTestCallback">[docs]</a><span class="k">class</span> <span class="nc">PhaseContextTestCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A callback that saves the phase context the for testing.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">PhaseContextTestCallback</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><span class="n">phase</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">context</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">context</span> <span class="o">=</span> <span class="n">context</span></div>


<div class="viewcode-block" id="DetectionVisualizationCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.DetectionVisualizationCallback">[docs]</a><span class="k">class</span> <span class="nc">DetectionVisualizationCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A callback that adds a visualization of a batch of detection predictions to context.sg_logger</span>
<span class="sd">    Attributes:</span>
<span class="sd">        freq: frequency (in epochs) to perform this callback.</span>
<span class="sd">        batch_idx: batch index to perform visualization for.</span>
<span class="sd">        classes: class list of the dataset.</span>
<span class="sd">        last_img_idx_in_batch: Last image index to add to log. (default=-1, will take entire batch).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">,</span>
        <span class="n">freq</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
        <span class="n">post_prediction_callback</span><span class="p">:</span> <span class="n">DetectionPostPredictionCallback</span><span class="p">,</span>
        <span class="n">classes</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span>
        <span class="n">batch_idx</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
        <span class="n">last_img_idx_in_batch</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span>
    <span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DetectionVisualizationCallback</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><span class="n">phase</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">freq</span> <span class="o">=</span> <span class="n">freq</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">post_prediction_callback</span> <span class="o">=</span> <span class="n">post_prediction_callback</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span> <span class="o">=</span> <span class="n">batch_idx</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">classes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">last_img_idx_in_batch</span> <span class="o">=</span> <span class="n">last_img_idx_in_batch</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">freq</span> <span class="o">==</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span> <span class="o">==</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">:</span>
            <span class="c1"># SOME CALCULATIONS ARE IN-PLACE IN NMS, SO CLONE THE PREDICTIONS</span>
            <span class="n">preds</span> <span class="o">=</span> <span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">preds</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">(),</span> <span class="kc">None</span><span class="p">)</span>
            <span class="n">preds</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">post_prediction_callback</span><span class="p">(</span><span class="n">preds</span><span class="p">)</span>
            <span class="n">batch_imgs</span> <span class="o">=</span> <span class="n">DetectionVisualization</span><span class="o">.</span><span class="n">visualize_batch</span><span class="p">(</span>
                <span class="n">context</span><span class="o">.</span><span class="n">inputs</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">target</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">classes</span>
            <span class="p">)</span>
            <span class="n">batch_imgs</span> <span class="o">=</span> <span class="p">[</span><span class="n">cv2</span><span class="o">.</span><span class="n">cvtColor</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="n">cv2</span><span class="o">.</span><span class="n">COLOR_BGR2RGB</span><span class="p">)</span> <span class="k">for</span> <span class="n">image</span> <span class="ow">in</span> <span class="n">batch_imgs</span><span class="p">]</span>
            <span class="n">batch_imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">batch_imgs</span><span class="p">)</span>
            <span class="n">tag</span> <span class="o">=</span> <span class="s2">&quot;batch_&quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">)</span> <span class="o">+</span> <span class="s2">&quot;_images&quot;</span>
            <span class="n">context</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_images</span><span class="p">(</span>
                <span class="n">tag</span><span class="o">=</span><span class="n">tag</span><span class="p">,</span> <span class="n">images</span><span class="o">=</span><span class="n">batch_imgs</span><span class="p">[:</span> <span class="bp">self</span><span class="o">.</span><span class="n">last_img_idx_in_batch</span><span class="p">],</span> <span class="n">global_step</span><span class="o">=</span><span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="n">data_format</span><span class="o">=</span><span class="s2">&quot;NHWC&quot;</span>
            <span class="p">)</span></div>


<div class="viewcode-block" id="BinarySegmentationVisualizationCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.BinarySegmentationVisualizationCallback">[docs]</a><span class="k">class</span> <span class="nc">BinarySegmentationVisualizationCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A callback that adds a visualization of a batch of segmentation predictions to context.sg_logger</span>
<span class="sd">    Attributes:</span>
<span class="sd">        freq: frequency (in epochs) to perform this callback.</span>
<span class="sd">        batch_idx: batch index to perform visualization for.</span>
<span class="sd">        last_img_idx_in_batch: Last image index to add to log. (default=-1, will take entire batch).</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">,</span> <span class="n">freq</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">batch_idx</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="n">last_img_idx_in_batch</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">BinarySegmentationVisualizationCallback</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><span class="n">phase</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">freq</span> <span class="o">=</span> <span class="n">freq</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span> <span class="o">=</span> <span class="n">batch_idx</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">last_img_idx_in_batch</span> <span class="o">=</span> <span class="n">last_img_idx_in_batch</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">freq</span> <span class="o">==</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">context</span><span class="o">.</span><span class="n">batch_idx</span> <span class="o">==</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">preds</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
                <span class="n">preds</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">preds</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">preds</span> <span class="o">=</span> <span class="n">context</span><span class="o">.</span><span class="n">preds</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
            <span class="n">batch_imgs</span> <span class="o">=</span> <span class="n">BinarySegmentationVisualization</span><span class="o">.</span><span class="n">visualize_batch</span><span class="p">(</span>
                <span class="n">context</span><span class="o">.</span><span class="n">inputs</span><span class="p">,</span> <span class="n">preds</span><span class="p">,</span> <span class="n">context</span><span class="o">.</span><span class="n">target</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span>
            <span class="p">)</span>
            <span class="n">batch_imgs</span> <span class="o">=</span> <span class="p">[</span><span class="n">cv2</span><span class="o">.</span><span class="n">cvtColor</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="n">cv2</span><span class="o">.</span><span class="n">COLOR_BGR2RGB</span><span class="p">)</span> <span class="k">for</span> <span class="n">image</span> <span class="ow">in</span> <span class="n">batch_imgs</span><span class="p">]</span>
            <span class="n">batch_imgs</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span><span class="n">batch_imgs</span><span class="p">)</span>
            <span class="n">tag</span> <span class="o">=</span> <span class="s2">&quot;batch_&quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">batch_idx</span><span class="p">)</span> <span class="o">+</span> <span class="s2">&quot;_images&quot;</span>
            <span class="n">context</span><span class="o">.</span><span class="n">sg_logger</span><span class="o">.</span><span class="n">add_images</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="n">tag</span><span class="p">,</span> <span class="n">images</span><span class="o">=</span><span class="n">batch_imgs</span><span class="p">[:</span><span class="bp">self</span><span class="o">.</span><span class="n">last_img_idx_in_batch</span><span class="p">],</span>
                                         <span class="n">global_step</span><span class="o">=</span><span class="n">context</span><span class="o">.</span><span class="n">epoch</span><span class="p">,</span> <span class="n">data_format</span><span class="o">=</span><span class="s1">&#39;NHWC&#39;</span><span class="p">)</span></div>


<div class="viewcode-block" id="TrainingStageSwitchCallbackBase"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.TrainingStageSwitchCallbackBase">[docs]</a><span class="k">class</span> <span class="nc">TrainingStageSwitchCallbackBase</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    TrainingStageSwitchCallback</span>

<span class="sd">    A phase callback that is called at a specific epoch (epoch start) to support multi-stage training.</span>
<span class="sd">    It does so by manipulating the objects inside the context.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        next_stage_start_epoch: int, the epoch idx to apply the stage change.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">next_stage_start_epoch</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">TrainingStageSwitchCallbackBase</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><span class="n">phase</span><span class="o">=</span><span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_EPOCH_START</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">next_stage_start_epoch</span> <span class="o">=</span> <span class="n">next_stage_start_epoch</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">context</span><span class="o">.</span><span class="n">epoch</span> <span class="o">==</span> <span class="bp">self</span><span class="o">.</span><span class="n">next_stage_start_epoch</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">apply_stage_change</span><span class="p">(</span><span class="n">context</span><span class="p">)</span>

<div class="viewcode-block" id="TrainingStageSwitchCallbackBase.apply_stage_change"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.TrainingStageSwitchCallbackBase.apply_stage_change">[docs]</a>    <span class="k">def</span> <span class="nf">apply_stage_change</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This method is called when the callback is fired on the next_stage_start_epoch,</span>
<span class="sd">         and holds the stage change logic that should be applied to the context&#39;s objects.</span>

<span class="sd">        :param context: PhaseContext, context of current phase</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div></div>


<div class="viewcode-block" id="YoloXTrainingStageSwitchCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.YoloXTrainingStageSwitchCallback">[docs]</a><span class="k">class</span> <span class="nc">YoloXTrainingStageSwitchCallback</span><span class="p">(</span><span class="n">TrainingStageSwitchCallbackBase</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    YoloXTrainingStageSwitchCallback</span>

<span class="sd">    Training stage switch for YoloX training.</span>
<span class="sd">    Disables mosaic, and manipulates YoloX loss to use L1.</span>

<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">next_stage_start_epoch</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">285</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">YoloXTrainingStageSwitchCallback</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><span class="n">next_stage_start_epoch</span><span class="o">=</span><span class="n">next_stage_start_epoch</span><span class="p">)</span>

<div class="viewcode-block" id="YoloXTrainingStageSwitchCallback.apply_stage_change"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.YoloXTrainingStageSwitchCallback.apply_stage_change">[docs]</a>    <span class="k">def</span> <span class="nf">apply_stage_change</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">transform</span> <span class="ow">in</span> <span class="n">context</span><span class="o">.</span><span class="n">train_loader</span><span class="o">.</span><span class="n">dataset</span><span class="o">.</span><span class="n">transforms</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">transform</span><span class="p">,</span> <span class="s2">&quot;close&quot;</span><span class="p">):</span>
                <span class="n">transform</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
        <span class="nb">iter</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">train_loader</span><span class="p">)</span>
        <span class="n">context</span><span class="o">.</span><span class="n">criterion</span><span class="o">.</span><span class="n">use_l1</span> <span class="o">=</span> <span class="kc">True</span></div></div>


<div class="viewcode-block" id="CallbackHandler"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.CallbackHandler">[docs]</a><span class="k">class</span> <span class="nc">CallbackHandler</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Runs all callbacks who&#39;s phase attribute equals to the given phase.</span>

<span class="sd">    Attributes:</span>
<span class="sd">        callbacks: List[PhaseCallback]. Callbacks to be run.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">callbacks</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">callbacks</span> <span class="o">=</span> <span class="n">callbacks</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">callback</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">callbacks</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">callback</span><span class="o">.</span><span class="n">phase</span> <span class="o">==</span> <span class="n">phase</span><span class="p">:</span>
                <span class="n">callback</span><span class="p">(</span><span class="n">context</span><span class="p">)</span></div>


<span class="c1"># DICT FOR LEGACY LR HARD-CODED REGIMES, WILL BE DELETED IN THE FUTURE</span>
<span class="n">LR_SCHEDULERS_CLS_DICT</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;step&quot;</span><span class="p">:</span> <span class="n">StepLRCallback</span><span class="p">,</span>
    <span class="s2">&quot;poly&quot;</span><span class="p">:</span> <span class="n">PolyLRCallback</span><span class="p">,</span>
    <span class="s2">&quot;cosine&quot;</span><span class="p">:</span> <span class="n">CosineLRCallback</span><span class="p">,</span>
    <span class="s2">&quot;exp&quot;</span><span class="p">:</span> <span class="n">ExponentialLRCallback</span><span class="p">,</span>
    <span class="s2">&quot;function&quot;</span><span class="p">:</span> <span class="n">FunctionLRCallback</span><span class="p">,</span>
<span class="p">}</span>

<span class="n">LR_WARMUP_CLS_DICT</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;linear_step&quot;</span><span class="p">:</span> <span class="n">WarmupLRCallback</span><span class="p">}</span>


<div class="viewcode-block" id="TestLRCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.callbacks.TestLRCallback">[docs]</a><span class="k">class</span> <span class="nc">TestLRCallback</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Phase callback that collects the learning rates in lr_placeholder at the end of each epoch (used for testing). In</span>
<span class="sd">     the case of multiple parameter groups (i.e multiple learning rates) the learning rate is collected from the first</span>
<span class="sd">     one. The phase is VALIDATION_EPOCH_END to ensure all lr updates have been performed before calling this callback.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">lr_placeholder</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">TestLRCallback</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><span class="n">Phase</span><span class="o">.</span><span class="n">VALIDATION_EPOCH_END</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr_placeholder</span> <span class="o">=</span> <span class="n">lr_placeholder</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lr_placeholder</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">optimizer</span><span class="o">.</span><span class="n">param_groups</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="s2">&quot;lr&quot;</span><span class="p">])</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.utils.checkpoint_utils &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.utils.checkpoint_utils &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -89,17 +91,40 @@
 
                             <span></span><span class="kn">import</span> <span class="nn">os</span>
                
 
                             <span class="kn">import</span> <span class="nn">tempfile</span>
                
 
                             <span class="kn">import</span> <span class="nn">pkg_resources</span>
                
 
                            +
                
 
                             <span class="kn">import</span> <span class="nn">torch</span>
                
 
                            +
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.common</span> <span class="kn">import</span> <span class="n">explicit_params_validation</span><span class="p">,</span> <span class="n">ADNNModelRepositoryDataInterfaces</span>
                
 
                             <span class="kn">from</span> <span class="nn">super_gradients.training.pretrained_models</span> <span class="kn">import</span> <span class="n">MODEL_URLS</span>
                
 
                            +<span class="kn">from</span> <span class="nn">super_gradients.common.environment</span> <span class="kn">import</span> <span class="n">environment_config</span>
                
 
                            +
                
 
                             <span class="k">try</span><span class="p">:</span>
                
 
                                 <span class="kn">from</span> <span class="nn">torch.hub</span> <span class="kn">import</span> <span class="n">download_url_to_file</span><span class="p">,</span> <span class="n">load_state_dict_from_url</span>
                
 
                             <span class="k">except</span> <span class="p">(</span><span class="ne">ModuleNotFoundError</span><span class="p">,</span> <span class="ne">ImportError</span><span class="p">,</span> <span class="ne">NameError</span><span class="p">):</span>
                
 
                                 <span class="kn">from</span> <span class="nn">torch.hub</span> <span class="kn">import</span> <span class="n">_download_url_to_file</span> <span class="k">as</span> <span class="n">download_url_to_file</span>
                
 
                            -<div class="viewcode-block" id="get_ckpt_local_path"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.get_ckpt_local_path">[docs]</a><span class="k">def</span> <span class="nf">get_ckpt_local_path</span><span class="p">(</span><span class="n">source_ckpt_folder_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</spa
                
 
                            -                        <span class="n">overwrite_local_checkpoint</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">load_weights_only</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>
                
 
                            +<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">get_checkpoints_dir_path</span><span class="p">(</span><span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">ckpt_root_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Creating the checkpoint directory of a given experiment.</span>
                
 
                            +<span class="sd">    :param experiment_name:     Name of the experiment.</span>
                
 
                            +<span class="sd">    :param ckpt_root_dir:       Local root directory path where all experiment logging directories will</span>
                
 
                            +<span class="sd">                                reside. When none is give, it is assumed that pkg_resources.resource_filename(&#39;checkpoints&#39;, &quot;&quot;)</span>
                
 
                            +<span class="sd">                                exists and will be used.</span>
                
 
                            +<span class="sd">    :return:                    checkpoints_dir_path</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +    <span class="k">if</span> <span class="n">ckpt_root_dir</span><span class="p">:</span>
                
 
                            +        <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">ckpt_root_dir</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">)</span>
                
 
                            +    <span class="k">elif</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">environment_config</span><span class="o">.</span><span class="n">PKG_CHECKPOINTS_DIR</span><span class="p">):</span>
                
 
                            +        <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">environment_config</span><span class="o">.</span><span class="n">PKG_CHECKPOINTS_DIR</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">)</span>
                
 
                            +    <span class="k">else</span><span class="p">:</span>
                
 
                            +        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Illegal checkpoints directory: pass ckpt_root_dir that exists, or add &#39;checkpoints&#39; to resources.&quot;</span><span class="p">)</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">get_ckpt_local_path</span><span class="p">(</span><span class="n">source_ckpt_folder_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">experiment_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">ckpt_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">external_checkpoint_path</span><span class
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Gets the local path to the checkpoint file, which will be:</span>
                
 
                             <span class="sd">        - By default: YOUR_REPO_ROOT/super_gradients/checkpoints/experiment_name.</span>
                
@@ -110,37 +135,20 @@
 
                             <span class="sd">        - external_checkpoint_path when external_checkpoint_path != None</span>
                
 
                             <span class="sd">    @param source_ckpt_folder_name: The folder where the checkpoint is saved. When set to None- uses the experiment_name.</span>
                
 
                            -<span class="sd">    @param experiment_name: experiment name attr in sg_model</span>
                
 
                            +<span class="sd">    @param experiment_name: experiment name attr in trainer</span>
                
 
                             <span class="sd">    @param ckpt_name: checkpoint filename</span>
                
 
                            -<span class="sd">    @param model_checkpoints_location: S3, local ot URL</span>
                
 
                             <span class="sd">    @param external_checkpoint_path: full path to checkpoint file (that might be located outside of super_gradients/checkpoints directory)</span>
                
 
                            -<span class="sd">    @param overwrite_local_checkpoint: whether to overwrite the checkpoint file with the same name when downloading from S3.</span>
                
 
                            -<span class="sd">    @param load_weights_only: whether to load the network&#39;s state dict only.</span>
                
 
                             <span class="sd">    @return:</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="n">source_ckpt_folder_name</span> <span class="o">=</span> <span class="n">source_ckpt_folder_name</span> <span class="ow">or</span> <span class="n">experiment_name</span>
                
 
                            -    <span class="k">if</span> <span class="n">model_checkpoints_location</span> <span class="o">==</span> <span class="s1">&#39;local&#39;</span><span class="p">:</span>
                
 
                            -        <span class="n">ckpt_local_path</span> <span class="o">=</span> <span class="n">external_checkpoint_path</span> <span class="ow">or</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s1">&#39;checkpoints&#39;</span><span class="p">,</span> <span class="n">source_ckpt_folder_name</span> <span class="o">+</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span clas
                
 
                            -
                
 
                            -    <span class="c1"># COPY THE DATA FROM &#39;S3&#39;/&#39;URL&#39; INTO A LOCAL DIRECTORY</span>
                
 
                            -    <span class="k">elif</span> <span class="n">model_checkpoints_location</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;s3&#39;</span><span class="p">)</span> <span class="ow">or</span> <span class="n">model_checkpoints_location</span> <span class="o">==</span> <span class="s1">&#39;url&#39;</span><span class="p">:</span>
                
 
                            -        <span class="c1"># COPY REMOTE DATA TO A LOCAL DIRECTORY AND GET THAT DIRECTORYs NAME</span>
                
 
                            -        <span class="n">ckpt_local_path</span> <span class="o">=</span> <span class="n">copy_ckpt_to_local_folder</span><span class="p">(</span><span class="n">local_ckpt_destination_dir</span><span class="o">=</span><span class="n">experiment_name</span><span class="p">,</span>
                
 
                            -                                                    <span class="n">ckpt_filename</span><span class="o">=</span><span class="n">ckpt_name</span><span class="p">,</span>
                
 
                            -                                                    <span class="n">remote_ckpt_source_dir</span><span class="o">=</span><span class="n">source_ckpt_folder_name</span><span class="p">,</span>
                
 
                            -                                                    <span class="n">path_src</span><span class="o">=</span><span class="n">model_checkpoints_location</span><span class="p">,</span>
                
 
                            -                                                    <span class="n">overwrite_local_ckpt</span><span class="o">=</span><span class="n">overwrite_local_checkpoint</span><span class="p">,</span>
                
 
                            -                                                    <span class="n">load_weights_only</span><span class="o">=</span><span class="n">load_weights_only</span><span class="p">)</span>
                
 
                            -
                
 
                            +    <span class="k">if</span> <span class="n">external_checkpoint_path</span><span class="p">:</span>
                
 
                            +        <span class="k">return</span> <span class="n">external_checkpoint_path</span>
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
 
                            -        <span class="c1"># ERROR IN USER CODE FLOW - THIS WILL EVENTUALLY RAISE AN EXCEPTION</span>
                
 
                            -        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span>
                
 
                            -            <span class="s1">&#39;model_checkpoints_data_source: &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">model_checkpoints_location</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;not supported&#39;</span><span class="p">)</span>
                
 
                            +        <span class="n">checkpoints_folder_name</span> <span class="o">=</span> <span class="n">source_ckpt_folder_name</span> <span class="ow">or</span> <span class="n">experiment_name</span>
                
 
                            +        <span class="n">checkpoints_dir_path</span> <span class="o">=</span> <span class="n">get_checkpoints_dir_path</span><span class="p">(</span><span class="n">checkpoints_folder_name</span><span class="p">)</span>
                
 
                            +        <span class="k">return</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">checkpoints_dir_path</span><span class="p">,</span> <span class="n">ckpt_name</span><span class="p">)</span>
                
 
                            -    <span class="k">return</span> <span class="n">ckpt_local_path</span></div>
                
 
                            -
                
 
                            -<div class="viewcode-block" id="adaptive_load_state_dict"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.adaptive_load_state_dict">[docs]</a><span class="k">def</span> <span class="nf">adaptive_load_state_dict</span><span class="p">(</span><span class="n">net</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p"
                
 
                            +<span class="k">def</span> <span class="nf">adaptive_load_state_dict</span><span class="p">(</span><span class="n">net</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">state_dict</span><span class="p">:</span> <span class="nb">dict</span><span class="p">,</span> <span class="n">strict</span><span class="p">:</span> <span class="nb">str</span><s
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Adaptively loads state_dict to net, by adapting the state_dict to net&#39;s layer names first.</span>
                
@@ -150,19 +158,24 @@
 
                             <span class="sd">    @return:</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                                 <span class="k">try</span><span class="p">:</span>
                
 
                            -        <span class="n">net</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">state_dict</span><span class="p">[</span><span class="s1">&#39;net&#39;</span><span class="p">],</span> <span class="n">strict</span><span class="o">=</span><span class="n">strict</span><span class="p">)</span>
                
 
                            +        <span class="n">net</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">state_dict</span><span class="p">[</span><span class="s2">&quot;net&quot;</span><span class="p">]</span> <span class="k">if</span> <span class="s2">&quot;net&quot;</span> <span class="ow">in</span> <span class="n">state_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> <span class="k">else</span> <span class="n">state_dict<
                
 
                                 <span class="k">except</span> <span class="p">(</span><span class="ne">RuntimeError</span><span class="p">,</span> <span class="ne">ValueError</span><span class="p">,</span> <span class="ne">KeyError</span><span class="p">)</span> <span class="k">as</span> <span class="n">ex</span><span class="p">:</span>
                
 
                            -        <span class="k">if</span> <span class="n">strict</span> <span class="o">==</span> <span class="s1">&#39;no_key_matching&#39;</span><span class="p">:</span>
                
 
                            +        <span class="k">if</span> <span class="n">strict</span> <span class="o">==</span> <span class="s2">&quot;no_key_matching&quot;</span><span class="p">:</span>
                
 
                                         <span class="n">adapted_state_dict</span> <span class="o">=</span> <span class="n">adapt_state_dict_to_fit_model_layer_names</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">state_dict</span><span class="p">(),</span> <span class="n">state_dict</span><span class="p">)</span>
                
 
                            -            <span class="n">net</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">adapted_state_dict</span><span class="p">[</span><span class="s1">&#39;net&#39;</span><span class="p">],</span> <span class="n">strict</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                            +            <span class="n">net</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">adapted_state_dict</span><span class="p">[</span><span class="s2">&quot;net&quot;</span><span class="p">],</span> <span class="n">strict</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">raise_informative_runtime_error</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">state_dict</span><span class="p">(),</span> <span class="n">state_dict</span><span class="p">,</span> <span class="n">ex</span><span class="p">)</span></div>
                
 
                            -
                
 
                            -
                
 
                            -<span class="nd">@explicit_params_validation</span><span class="p">(</span><span class="n">validation_type</span><span class="o">=</span><span class="s1">&#39;None&#39;</span><span class="p">)</span>
                
 
                            -<span class="k">def</span> <span class="nf">copy_ckpt_to_local_folder</span><span class="p">(</span><span class="n">local_ckpt_destination_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">ckpt_filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">remote_ckpt_source_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span cla
                
 
                            -                              <span class="n">path_src</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;local&#39;</span><span class="p">,</span> <span class="n">overwrite_local_ckpt</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            -                              <span class="n">load_weights_only</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
                
 
                            +            <span class="n">raise_informative_runtime_error</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">state_dict</span><span class="p">(),</span> <span class="n">state_dict</span><span class="p">,</span> <span class="n">ex</span><span class="p">)</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="nd">@explicit_params_validation</span><span class="p">(</span><span class="n">validation_type</span><span class="o">=</span><span class="s2">&quot;None&quot;</span><span class="p">)</span>
                
 
                            +<span class="k">def</span> <span class="nf">copy_ckpt_to_local_folder</span><span class="p">(</span>
                
 
                            +    <span class="n">local_ckpt_destination_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                
 
                            +    <span class="n">ckpt_filename</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                
 
                            +    <span class="n">remote_ckpt_source_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                
 
                            +    <span class="n">path_src</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;local&quot;</span><span class="p">,</span>
                
 
                            +    <span class="n">overwrite_local_ckpt</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +    <span class="n">load_weights_only</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                
 
                            +<span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Copy the checkpoint from any supported source to a local destination path</span>
                
 
                             <span class="sd">        :param local_ckpt_destination_dir:  destination where the checkpoint will be saved to</span>
                
@@ -181,27 +194,31 @@
 
                                 <span class="k">if</span> <span class="ow">not</span> <span class="n">overwrite_local_ckpt</span><span class="p">:</span>
                
 
                                     <span class="c1"># CREATE A TEMP FOLDER TO SAVE THE CHECKPOINT TO</span>
                
 
                                     <span class="n">download_ckpt_destination_dir</span> <span class="o">=</span> <span class="n">tempfile</span><span class="o">.</span><span class="n">gettempdir</span><span class="p">()</span>
                
 
                            -        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;PLEASE NOTICE - YOU ARE IMPORTING A REMOTE CHECKPOINT WITH overwrite_local_checkpoint = False &#39;</span>
                
 
                            -              <span class="s1">&#39;-&gt; IT WILL BE REDIRECTED TO A TEMP FOLDER AND DELETED ON MACHINE RESTART&#39;</span><span class="p">)</span>
                
 
                            +        <span class="nb">print</span><span class="p">(</span>
                
 
                            +            <span class="s2">&quot;PLEASE NOTICE - YOU ARE IMPORTING A REMOTE CHECKPOINT WITH overwrite_local_checkpoint = False &quot;</span>
                
 
                            +            <span class="s2">&quot;-&gt; IT WILL BE REDIRECTED TO A TEMP FOLDER AND DELETED ON MACHINE RESTART&quot;</span>
                
 
                            +        <span class="p">)</span>
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
 
                                     <span class="c1"># SAVE THE CHECKPOINT TO MODEL&#39;s FOLDER</span>
                
 
                            -        <span class="n">download_ckpt_destination_dir</span> <span class="o">=</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s1">&#39;checkpoints&#39;</span><span class="p">,</span> <span class="n">local_ckpt_destination_dir</span><span class="p">)</span>
                
 
                            +        <span class="n">download_ckpt_destination_dir</span> <span class="o">=</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s2">&quot;checkpoints&quot;</span><span class="p">,</span> <span class="n">local_ckpt_destination_dir</span><span class="p">)</span>
                
 
                            -    <span class="k">if</span> <span class="n">path_src</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;s3&#39;</span><span class="p">):</span>
                
 
                            +    <span class="k">if</span> <span class="n">path_src</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;s3&quot;</span><span class="p">):</span>
                
 
                                     <span class="n">model_checkpoints_data_interface</span> <span class="o">=</span> <span class="n">ADNNModelRepositoryDataInterfaces</span><span class="p">(</span><span class="n">data_connection_location</span><span class="o">=</span><span class="n">path_src</span><span class="p">)</span>
                
 
                                     <span class="c1"># DOWNLOAD THE FILE FROM S3 TO THE DESTINATION FOLDER</span>
                
 
                                     <span class="n">ckpt_file_full_local_path</span> <span class="o">=</span> <span class="n">model_checkpoints_data_interface</span><span class="o">.</span><span class="n">load_remote_checkpoints_file</span><span class="p">(</span>
                
 
                                         <span class="n">ckpt_source_remote_dir</span><span class="o">=</span><span class="n">remote_ckpt_source_dir</span><span class="p">,</span>
                
 
                                         <span class="n">ckpt_destination_local_dir</span><span class="o">=</span><span class="n">download_ckpt_destination_dir</span><span class="p">,</span>
                
 
                                         <span class="n">ckpt_file_name</span><span class="o">=</span><span class="n">ckpt_filename</span><span class="p">,</span>
                
 
                            -            <span class="n">overwrite_local_checkpoints_file</span><span class="o">=</span><span class="n">overwrite_local_ckpt</span><span class="p">)</span>
                
 
                            +            <span class="n">overwrite_local_checkpoints_file</span><span class="o">=</span><span class="n">overwrite_local_ckpt</span><span class="p">,</span>
                
 
                            +        <span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="ow">not</span> <span class="n">load_weights_only</span><span class="p">:</span>
                
 
                                         <span class="c1"># COPY LOG FILES FROM THE REMOTE DIRECTORY TO THE LOCAL ONE ONLY IF LOADING THE CURRENT MODELs CKPT</span>
                
 
                            -            <span class="n">model_checkpoints_data_interface</span><span class="o">.</span><span class="n">load_all_remote_log_files</span><span class="p">(</span><span class="n">model_name</span><span class="o">=</span><span class="n">remote_ckpt_source_dir</span><span class="p">,</span>
                
 
                            -                                                                       <span class="n">model_checkpoint_local_dir</span><span class="o">=</span><span class="n">download_ckpt_destination_dir</span><span class="p">)</span>
                
 
                            +            <span class="n">model_checkpoints_data_interface</span><span class="o">.</span><span class="n">load_all_remote_log_files</span><span class="p">(</span>
                
 
                            +                <span class="n">model_name</span><span class="o">=</span><span class="n">remote_ckpt_source_dir</span><span class="p">,</span> <span class="n">model_checkpoint_local_dir</span><span class="o">=</span><span class="n">download_ckpt_destination_dir</span>
                
 
                            +            <span class="p">)</span>
                
 
                            -    <span class="k">if</span> <span class="n">path_src</span> <span class="o">==</span> <span class="s1">&#39;url&#39;</span><span class="p">:</span>
                
 
                            +    <span class="k">if</span> <span class="n">path_src</span> <span class="o">==</span> <span class="s2">&quot;url&quot;</span><span class="p">:</span>
                
 
                                     <span class="n">ckpt_file_full_local_path</span> <span class="o">=</span> <span class="n">download_ckpt_destination_dir</span> <span class="o">+</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">sep</span> <span class="o">+</span> <span class="n">ckpt_filename</span>
                
 
                                     <span class="c1"># DOWNLOAD THE FILE FROM URL TO THE DESTINATION FOLDER</span>
                
 
                                     <span class="n">download_url_to_file</span><span class="p">(</span><span class="n">remote_ckpt_source_dir</span><span class="p">,</span> <span class="n">ckpt_file_full_local_path</span><span class="p">,</span> <span class="n">progress</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
                
@@ -209,20 +226,19 @@
 
                                 <span class="k">return</span> <span class="n">ckpt_file_full_local_path</span>
                
 
                            -<div class="viewcode-block" id="read_ckpt_state_dict"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.read_ckpt_state_dict">[docs]</a><span class="k">def</span> <span class="nf">read_ckpt_state_dict</span><span class="p">(</span><span class="n">ckpt_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span><span c
                
 
                            +<span class="k">def</span> <span class="nf">read_ckpt_state_dict</span><span class="p">(</span><span class="n">ckpt_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span><span class="p">):</span>
                
 
                                 <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">ckpt_path</span><span class="p">):</span>
                
 
                            -        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Incorrect Checkpoint path&#39;</span><span class="p">)</span>
                
 
                            +        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Incorrect Checkpoint path&quot;</span><span class="p">)</span>
                
 
                                 <span class="k">if</span> <span class="n">device</span> <span class="o">==</span> <span class="s2">&quot;cuda&quot;</span><span class="p">:</span>
                
 
                                     <span class="n">state_dict</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">ckpt_path</span><span class="p">)</span>
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
 
                                     <span class="n">state_dict</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">ckpt_path</span><span class="p">,</span> <span class="n">map_location</span><span class="o">=</span><span class="k">lambda</span> <span class="n">storage</span><span class="p">,</span> <span class="n">loc</span><span class="p">:</span> <span class="n">storage</span><span class="p">)</span>
                
 
                            -    <span class="k">return</span> <span class="n">state_dict</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">state_dict</span>
                
 
                            -<div class="viewcode-block" id="adapt_state_dict_to_fit_model_layer_names"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.adapt_state_dict_to_fit_model_layer_names">[docs]</a><span class="k">def</span> <span class="nf">adapt_state_dict_to_fit_model_layer_names</span><span class="p">(</span><span class="n">model_state_dict</span><span class="p">:</span> <span class="nb">dict</span><span class="p">,</span> <span class="n">source_ckpt<
                
 
                            -                                              <span class="n">exclude</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="p">[],</span> <span class="n">solver</span><span class="p">:</span> <span class="n">callable</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="adapt_state_dict_to_fit_model_layer_names"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.adapt_state_dict_to_fit_model_layer_names">[docs]</a><span class="k">def</span> <span class="nf">adapt_state_dict_to_fit_model_layer_names</span><span class="p">(</span><span class="n">model_state_dict</span><span class="p">:</span> <span class="nb">dict</span><span class="p">,</span> <span class="n">source_ckpt</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Given a model state dict and source checkpoints, the method tries to correct the keys in the model_state_dict to fit</span>
                
 
                             <span class="sd">    the ckpt in order to properly load the weights into the model. If unsuccessful - returns None</span>
                
@@ -233,39 +249,41 @@
 
                             <span class="sd">                                        that returns a desired weight for ckpt_val.</span>
                
 
                             <span class="sd">        :return: renamed checkpoint dict (if possible)</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="k">if</span> <span class="s1">&#39;net&#39;</span> <span class="ow">in</span> <span class="n">source_ckpt</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                            -        <span class="n">source_ckpt</span> <span class="o">=</span> <span class="n">source_ckpt</span><span class="p">[</span><span class="s1">&#39;net&#39;</span><span class="p">]</span>
                
 
                            +    <span class="k">if</span> <span class="s2">&quot;net&quot;</span> <span class="ow">in</span> <span class="n">source_ckpt</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                            +        <span class="n">source_ckpt</span> <span class="o">=</span> <span class="n">source_ckpt</span><span class="p">[</span><span class="s2">&quot;net&quot;</span><span class="p">]</span>
                
 
                                 <span class="n">model_state_dict_excluded</span> <span class="o">=</span> <span class="p">{</span><span class="n">k</span><span class="p">:</span> <span class="n">v</span> <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">model_state_dict</span><span class="o">.</span><span class="n">items</span><span class="p">()</span> <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</sp
                
 
                                 <span class="n">new_ckpt_dict</span> <span class="o">=</span> <span class="p">{}</span>
                
 
                                 <span class="k">for</span> <span class="p">(</span><span class="n">ckpt_key</span><span class="p">,</span> <span class="n">ckpt_val</span><span class="p">),</span> <span class="p">(</span><span class="n">model_key</span><span class="p">,</span> <span class="n">model_val</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">source_ckpt</span><span class="o">.</span><span class="n">items</span><span class="p">(),</span> <s
                
 
                                     <span class="k">if</span> <span class="n">solver</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                         <span class="n">ckpt_val</span> <span class="o">=</span> <span class="n">solver</span><span class="p">(</span><span class="n">ckpt_key</span><span class="p">,</span> <span class="n">ckpt_val</span><span class="p">,</span> <span class="n">model_key</span><span class="p">,</span> <span class="n">model_val</span><span class="p">)</span>
                
 
                                     <span class="k">if</span> <span class="n">ckpt_val</span><span class="o">.</span><span class="n">shape</span> <span class="o">!=</span> <span class="n">model_val</span><span class="o">.</span><span class="n">shape</span><span class="p">:</span>
                
 
                            -            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;ckpt layer </span><span class="si">{</span><span class="n">ckpt_key</span><span class="si">}</span><span class="s1"> with shape </span><span class="si">{</span><span class="n">ckpt_val</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s1"> does not match </span><span class="si">{</span><span class="n">mode
                
 
                            -                             <span class="sa">f</span><span class="s1">&#39; with shape </span><span class="si">{</span><span class="n">model_val</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s1"> in the model&#39;</span><span class="p">)</span>
                
 
                            +            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;ckpt layer </span><span class="si">{</span><span class="n">ckpt_key</span><span class="si">}</span><span class="s2"> with shape </span><span class="si">{</span><span class="n">ckpt_val</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2"> does not match </span><span class="si">{</span><span class="n">mod
                
 
                                     <span class="n">new_ckpt_dict</span><span class="p">[</span><span class="n">model_key</span><span class="p">]</span> <span class="o">=</span> <span class="n">ckpt_val</span>
                
 
                            -    <span class="k">return</span> <span class="p">{</span><span class="s1">&#39;net&#39;</span><span class="p">:</span> <span class="n">new_ckpt_dict</span><span class="p">}</span></div>
                
 
                            +    <span class="k">return</span> <span class="p">{</span><span class="s2">&quot;net&quot;</span><span class="p">:</span> <span class="n">new_ckpt_dict</span><span class="p">}</span></div>
                
 
                            -<div class="viewcode-block" id="raise_informative_runtime_error"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.raise_informative_runtime_error">[docs]</a><span class="k">def</span> <span class="nf">raise_informative_runtime_error</span><span class="p">(</span><span class="n">state_dict</span><span class="p">,</span> <span class="n">checkpoint</span><span class="p">,</span> <span class="n">exception_msg</span><span class="p">):</spa
                
 
                            +<div class="viewcode-block" id="raise_informative_runtime_error"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.raise_informative_runtime_error">[docs]</a><span class="k">def</span> <span class="nf">raise_informative_runtime_error</span><span class="p">(</span><span class="n">state_dict</span><span class="p">,</span> <span class="n">checkpoint</span><span class="p">,</span> <span class="n">exception_msg</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Given a model state dict and source checkpoints, the method calls &quot;adapt_state_dict_to_fit_model_layer_names&quot;</span>
                
 
                             <span class="sd">    and enhances the exception_msg if loading the checkpoint_dict via the conversion method is possible</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                                 <span class="k">try</span><span class="p">:</span>
                
 
                                     <span class="n">new_ckpt_dict</span> <span class="o">=</span> <span class="n">adapt_state_dict_to_fit_model_layer_names</span><span class="p">(</span><span class="n">state_dict</span><span class="p">,</span> <span class="n">checkpoint</span><span class="p">)</span>
                
 
                            -        <span class="n">temp_file</span> <span class="o">=</span> <span class="n">tempfile</span><span class="o">.</span><span class="n">NamedTemporaryFile</span><span class="p">()</span><span class="o">.</span><span class="n">name</span> <span class="o">+</span> <span class="s1">&#39;.pt&#39;</span>
                
 
                            +        <span class="n">temp_file</span> <span class="o">=</span> <span class="n">tempfile</span><span class="o">.</span><span class="n">NamedTemporaryFile</span><span class="p">()</span><span class="o">.</span><span class="n">name</span> <span class="o">+</span> <span class="s2">&quot;.pt&quot;</span>
                
 
                                     <span class="n">torch</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">new_ckpt_dict</span><span class="p">,</span> <span class="n">temp_file</span><span class="p">)</span>
                
 
                            -        <span class="n">exception_msg</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="se">\n</span><span class="si">{</span><span class="s1">&#39;=&#39;</span> <span class="o">*</span> <span class="mi">200</span><span class="si">}</span><span class="se">\n</span><span class="si">{</span><span class="nb">str</span><span class="p">(</span><span class="n">exception_msg</span><span class="p">)</span><span class="si">}</span><span class="s2"> </spa
                
 
                            -                        <span class="sa">f</span><span class="s2">&quot;model_layer_names method</span><span class="se">\n</span><span class="s2">a converted checkpoint file was saved in the path </span><span class="si">{</span><span class="n">temp_file</span><span class="si">}</span><span class="se">\n</span><span class="si">{</span><span class="s1">&#39;=&#39;</span> <span class="o">*</span> <span class="mi">200</span><span class="si">}</span><span class="s2">&quot;</span>
                
 
                            +        <span class="n">exception_msg</span> <span class="o">=</span> <span class="p">(</span>
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;</span><span class="se">\n</span><span class="si">{</span><span class="s1">&#39;=&#39;</span> <span class="o">*</span> <span class="mi">200</span><span class="si">}</span><span class="se">\n</span><span class="si">{</span><span class="nb">str</span><span class="p">(</span><span class="n">exception_msg</span><span class="p">)</span><span class="si">}</span><span class="s2"> </span><span class="se">\n</span><span class="s2">convert ckpt 
                
 
                            +            <span class="sa">f</span><span class="s2">&quot;model_layer_names method</span><span class="se">\n</span><span class="s2">a converted checkpoint file was saved in the path </span><span class="si">{</span><span class="n">temp_file</span><span class="si">}</span><span class="se">\n</span><span class="si">{</span><span class="s1">&#39;=&#39;</span> <span class="o">*</span> <span class="mi">200</span><span class="si">}</span><span class="s2">&quot;</span>
                
 
                            +        <span class="p">)</span>
                
 
                                 <span class="k">except</span> <span class="ne">ValueError</span> <span class="k">as</span> <span class="n">ex</span><span class="p">:</span>  <span class="c1"># IN CASE adapt_state_dict_to_fit_model_layer_names WAS UNSUCCESSFUL</span>
                
 
                                     <span class="n">exception_msg</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="se">\n</span><span class="si">{</span><span class="s1">&#39;=&#39;</span> <span class="o">*</span> <span class="mi">200</span><span class="si">}</span><span class="s2"> </span><span class="se">\n</span><span class="s2">The checkpoint and model shapes do no fit, e.g.: </span><span class="si">{</span><span class="n">ex</span><span class="si">}</span><span class
                
 
                                 <span class="k">finally</span><span class="p">:</span>
                
 
                                     <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="n">exception_msg</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="load_checkpoint_to_model"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.load_checkpoint_to_model">[docs]</a><span class="k">def</span> <span class="nf">load_checkpoint_to_model</span><span class="p">(</span><span class="n">ckpt_local_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">load_backbone</span><span class="p">:</span> <span class="nb">boo
                
 
                            -                             <span class="n">load_weights_only</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">load_ema_as_net</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">load_checkpoint_to_model</span><span class="p">(</span>
                
 
                            +    <span class="n">ckpt_local_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">load_backbone</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">net</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">strict</span><span class="p">:</span> <span 
                
 
                            +<span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Loads the state dict in ckpt_local_path to net and returns the checkpoint&#39;s state dict.</span>
                
@@ -278,35 +296,39 @@
 
                             <span class="sd">    @return:</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                                 <span class="k">if</span> <span class="n">ckpt_local_path</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">ckpt_local_path</span><span class="p">):</span>
                
 
                            -        <span class="n">error_msg</span> <span class="o">=</span> <span class="s1">&#39;Error - loading Model Checkpoint: Path </span><span class="si">{}</span><span class="s1"> does not exist&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ckpt_local_path</span><span class="p">)</span>
                
 
                            +        <span class="n">error_msg</span> <span class="o">=</span> <span class="s2">&quot;Error - loading Model Checkpoint: Path </span><span class="si">{}</span><span class="s2"> does not exist&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ckpt_local_path</span><span class="p">)</span>
                
 
                                     <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="n">error_msg</span><span class="p">)</span>
                
 
                            -    <span class="k">if</span> <span class="n">load_backbone</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="s1">&#39;backbone&#39;</span><span class="p">):</span>
                
 
                            +    <span class="k">if</span> <span class="n">load_backbone</span> <span class="ow">and</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">net</span><span class="p">,</span> <span class="s2">&quot;backbone&quot;</span><span class="p">):</span>
                
 
                                     <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;No backbone attribute in net - Can&#39;t load backbone weights&quot;</span><span class="p">)</span>
                
 
                                 <span class="c1"># LOAD THE LOCAL CHECKPOINT PATH INTO A state_dict OBJECT</span>
                
 
                                 <span class="n">checkpoint</span> <span class="o">=</span> <span class="n">read_ckpt_state_dict</span><span class="p">(</span><span class="n">ckpt_path</span><span class="o">=</span><span class="n">ckpt_local_path</span><span class="p">)</span>
                
 
                                 <span class="k">if</span> <span class="n">load_ema_as_net</span><span class="p">:</span>
                
 
                            -        <span class="k">if</span> <span class="s1">&#39;ema_net&#39;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">checkpoint</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                            +        <span class="k">if</span> <span class="s2">&quot;ema_net&quot;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">checkpoint</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                                         <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Can&#39;t load ema network- no EMA network stored in checkpoint file&quot;</span><span class="p">)</span>
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">checkpoint</span><span class="p">[</span><span class="s1">&#39;net&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">checkpoint</span><span class="p">[</span><span class="s1">&#39;ema_net&#39;</span><span class="p">]</span>
                
 
                            +            <span class="n">checkpoint</span><span class="p">[</span><span class="s2">&quot;net&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">checkpoint</span><span class="p">[</span><span class="s2">&quot;ema_net&quot;</span><span class="p">]</span>
                
 
                                 <span class="c1"># LOAD THE CHECKPOINTS WEIGHTS TO THE MODEL</span>
                
 
                                 <span class="k">if</span> <span class="n">load_backbone</span><span class="p">:</span>
                
 
                            -        <span class="n">adaptive_load_state_dict</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">backbone</span><span class="p">,</span> <span class="n">checkpoint</span><span class="p">,</span> <span class="n">strict</span><span class="p">)</span>
                
 
                            +        <span class="n">adaptive_load_state_dict</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">backbone</span><span class="p">,</span> <span class="n">checkpoint</span><span class="p">,</span> <span class="n">strict</span><span class="p">)</span>
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
 
                                     <span class="n">adaptive_load_state_dict</span><span class="p">(</span><span class="n">net</span><span class="p">,</span> <span class="n">checkpoint</span><span class="p">,</span> <span class="n">strict</span><span class="p">)</span>
                
 
                            +    <span class="n">message_suffix</span> <span class="o">=</span> <span class="s2">&quot; checkpoint.&quot;</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">load_ema_as_net</span> <span class="k">else</span> <span class="s2">&quot; EMA checkpoint.&quot;</span>
                
 
                            +    <span class="n">message_model</span> <span class="o">=</span> <span class="s2">&quot;model&quot;</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">load_backbone</span> <span class="k">else</span> <span class="s2">&quot;model&#39;s backbone&quot;</span>
                
 
                            +    <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Successfully loaded &quot;</span> <span class="o">+</span> <span class="n">message_model</span> <span class="o">+</span> <span class="s2">&quot; weights from &quot;</span> <span class="o">+</span> <span class="n">ckpt_local_path</span> <span class="o">+</span> <span class="n">message_suffix</span><span class="p">)</span>
                
 
                            +
                
 
                                 <span class="k">if</span> <span class="n">load_weights_only</span> <span class="ow">or</span> <span class="n">load_backbone</span><span class="p">:</span>
                
 
                                     <span class="c1"># DISCARD ALL THE DATA STORED IN CHECKPOINT OTHER THAN THE WEIGHTS</span>
                
 
                            -        <span class="p">[</span><span class="n">checkpoint</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="n">key</span><span class="p">)</span> <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="n">checkpoint</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> <span class="k">if</span> <span class="n">key</span> <span class=
                
 
                            +        <span class="p">[</span><span class="n">checkpoint</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="n">key</span><span class="p">)</span> <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="n">checkpoint</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> <span class="k">if</span> <span class="n">key</span> <span class=
                
 
                            -    <span class="k">return</span> <span class="n">checkpoint</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">checkpoint</span>
                
 
                            -<div class="viewcode-block" id="MissingPretrainedWeightsException"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.MissingPretrainedWeightsException">[docs]</a><span class="k">class</span> <span class="nc">MissingPretrainedWeightsException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
                
 
                            +<span class="k">class</span> <span class="nc">MissingPretrainedWeightsException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Exception raised by unsupported pretrianed model.</span>
                
 
                             <span class="sd">    Attributes:</span>
                
@@ -315,7 +337,7 @@
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">desc</span><span class="p">):</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">message</span> <span class="o">=</span> <span class="s2">&quot;Missing pretrained wights: &quot;</span> <span class="o">+</span> <span class="n">desc</span>
                
 
                            -        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">message</span><span class="p">)</span></div>
                
 
                            +        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">message</span><span class="p">)</span>
                
 
                             <span class="k">def</span> <span class="nf">_yolox_ckpt_solver</span><span class="p">(</span><span class="n">ckpt_key</span><span class="p">,</span> <span class="n">ckpt_val</span><span class="p">,</span> <span class="n">model_key</span><span class="p">,</span> <span class="n">model_val</span><span class="p">):</span>
                
@@ -323,8 +345,11 @@
 
                             <span class="sd">    Helper method for reshaping old pretrained checkpoint&#39;s focus weights to 6x6 conv weights.</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="k">if</span> <span class="n">ckpt_val</span><span class="o">.</span><span class="n">shape</span> <span class="o">!=</span> <span class="n">model_val</span><span class="o">.</span><span class="n">shape</span> <span class="ow">and</span> <span class="n">ckpt_key</span> <span class="o">==</span> <span class="s1">&#39;module._backbone._modules_list.0.conv.conv.weight&#39;</span> <span class="ow">and</span> \
                
 
                            -            <span class="n">model_key</span> <span class="o">==</span> <span class="s1">&#39;_backbone._modules_list.0.conv.weight&#39;</span><span class="p">:</span>
                
 
                            +    <span class="k">if</span> <span class="p">(</span>
                
 
                            +        <span class="n">ckpt_val</span><span class="o">.</span><span class="n">shape</span> <span class="o">!=</span> <span class="n">model_val</span><span class="o">.</span><span class="n">shape</span>
                
 
                            +        <span class="ow">and</span> <span class="n">ckpt_key</span> <span class="o">==</span> <span class="s2">&quot;module._backbone._modules_list.0.conv.conv.weight&quot;</span>
                
 
                            +        <span class="ow">and</span> <span class="n">model_key</span> <span class="o">==</span> <span class="s2">&quot;_backbone._modules_list.0.conv.weight&quot;</span>
                
 
                            +    <span class="p">):</span>
                
 
                                     <span class="n">model_val</span><span class="o">.</span><span class="n">data</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">::</span><span class="mi">2</span><span class="p">,</span> <span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">ckpt_val</span><span class="o">.</span><span class="n">data</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">3</span><span class="p">]</spa
                
 
                                     <span class="n">model_val</span><span class="o">.</span><span class="n">data</span><span class="p">[:,</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="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">ckpt_val</span><span class="o">.</span><span class="n">data</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">:</sp
                
 
                                     <span class="n">model_val</span><span class="o">.</span><span class="n">data</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">::</span><span class="mi">2</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="o">=</span> <span class="n">ckpt_val</span><span class="o">.</span><span class="n">data</span><span class="p">[:,</span> <span class="mi">6</span><span class="p">:</sp
                
@@ -336,7 +361,7 @@
 
                                 <span class="k">return</span> <span class="n">replacement</span>
                
 
                            -<div class="viewcode-block" id="load_pretrained_weights"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.load_pretrained_weights">[docs]</a><span class="k">def</span> <span class="nf">load_pretrained_weights</span><span class="p">(</span><span class="n">model</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">
                
 
                            +<span class="k">def</span> <span class="nf">load_pretrained_weights</span><span class="p">(</span><span class="n">model</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">architecture</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">pretrained_weights</span><span class="p">:</span> <span class="nb
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Loads pretrained weights from the MODEL_URLS dictionary to model</span>
                
@@ -345,21 +370,41 @@
 
                             <span class="sd">    @param pretrained_weights: name for the pretrianed weights (i.e imagenet)</span>
                
 
                             <span class="sd">    @return: None</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                            -    <span class="n">model_url_key</span> <span class="o">=</span> <span class="n">architecture</span> <span class="o">+</span> <span class="s1">&#39;_&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">pretrained_weights</span><span class="p">)</span>
                
 
                            +    <span class="n">model_url_key</span> <span class="o">=</span> <span class="n">architecture</span> <span class="o">+</span> <span class="s2">&quot;_&quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">pretrained_weights</span><span class="p">)</span>
                
 
                                 <span class="k">if</span> <span class="n">model_url_key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">MODEL_URLS</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                                     <span class="k">raise</span> <span class="n">MissingPretrainedWeightsException</span><span class="p">(</span><span class="n">model_url_key</span><span class="p">)</span>
                
 
                                 <span class="n">url</span> <span class="o">=</span> <span class="n">MODEL_URLS</span><span class="p">[</span><span class="n">model_url_key</span><span class="p">]</span>
                
 
                            -    <span class="n">unique_filename</span> <span class="o">=</span> <span class="n">url</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;https://deci-pretrained-models.s3.amazonaws.com/&quot;</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s1">&#39;/&#39;</span><span class="p">,</span> <span class="s1">&#39;_&#39;</spa
                
 
                            -    <span class="n">map_location</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="s1">&#39;cpu&#39;</span><span class="p">)</span> <span class="k">if</span> <span class="ow">not</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="kc">None
                
 
                            +    <span class="n">unique_filename</span> <span class="o">=</span> <span class="n">url</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;https://deci-pretrained-models.s3.amazonaws.com/&quot;</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot;/&quot;</span><span class="p">,</span> <span class="s2">&quot;_&quot;<
                
 
                            +    <span class="n">map_location</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">&quot;cpu&quot;</span><span class="p">)</span>
                
 
                                 <span class="n">pretrained_state_dict</span> <span class="o">=</span> <span class="n">load_state_dict_from_url</span><span class="p">(</span><span class="n">url</span><span class="o">=</span><span class="n">url</span><span class="p">,</span> <span class="n">map_location</span><span class="o">=</span><span class="n">map_location</span><span class="p">,</span> <span class="n">file_name</span><span class="o">=</span><span class="n">unique_filename</span><span class="p">)</span>
                
 
                            -    <span class="k">if</span> <span class="s1">&#39;ema_net&#39;</span> <span class="ow">in</span> <span class="n">pretrained_state_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                            -        <span class="n">pretrained_state_dict</span><span class="p">[</span><span class="s1">&#39;net&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">pretrained_state_dict</span><span class="p">[</span><span class="s1">&#39;ema_net&#39;</span><span class="p">]</span>
                
 
                            +    <span class="n">_load_weights</span><span class="p">(</span><span class="n">architecture</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">pretrained_state_dict</span><span class="p">)</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">_load_weights</span><span class="p">(</span><span class="n">architecture</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">pretrained_state_dict</span><span class="p">):</span>
                
 
                            +    <span class="k">if</span> <span class="s2">&quot;ema_net&quot;</span> <span class="ow">in</span> <span class="n">pretrained_state_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                            +        <span class="n">pretrained_state_dict</span><span class="p">[</span><span class="s2">&quot;net&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">pretrained_state_dict</span><span class="p">[</span><span class="s2">&quot;ema_net&quot;</span><span class="p">]</span>
                
 
                                 <span class="n">solver</span> <span class="o">=</span> <span class="n">_yolox_ckpt_solver</span> <span class="k">if</span> <span class="s2">&quot;yolox&quot;</span> <span class="ow">in</span> <span class="n">architecture</span> <span class="k">else</span> <span class="kc">None</span>
                
 
                            -    <span class="n">adapted_pretrained_state_dict</span> <span class="o">=</span> <span class="n">adapt_state_dict_to_fit_model_layer_names</span><span class="p">(</span><span class="n">model_state_dict</span><span class="o">=</span><span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">(),</span>
                
 
                            -                                                                              <span class="n">source_ckpt</span><span class="o">=</span><span class="n">pretrained_state_dict</span><span class="p">,</span>
                
 
                            -                                                                              <span class="n">solver</span><span class="o">=</span><span class="n">solver</span><span class="p">)</span>
                
 
                            -    <span class="n">model</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">adapted_pretrained_state_dict</span><span class="p">[</span><span class="s1">&#39;net&#39;</span><span class="p">],</span> <span class="n">strict</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span></div>
                
 
                            +    <span class="n">adapted_pretrained_state_dict</span> <span class="o">=</span> <span class="n">adapt_state_dict_to_fit_model_layer_names</span><span class="p">(</span>
                
 
                            +        <span class="n">model_state_dict</span><span class="o">=</span><span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">(),</span> <span class="n">source_ckpt</span><span class="o">=</span><span class="n">pretrained_state_dict</span><span class="p">,</span> <span class="n">solver</span><span class="o">=</span><span class="n">solver</span>
                
 
                            +    <span class="p">)</span>
                
 
                            +    <span class="n">model</span><span class="o">.</span><span class="n">load_state_dict</span><span class="p">(</span><span class="n">adapted_pretrained_state_dict</span><span class="p">[</span><span class="s2">&quot;net&quot;</span><span class="p">],</span> <span class="n">strict</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">load_pretrained_weights_local</span><span class="p">(</span><span class="n">model</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">architecture</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">pretrained_weights</span><span class="p">:</span> <span cla
                
 
                            +
                
 
                            +    <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">    Loads pretrained weights from the MODEL_URLS dictionary to model</span>
                
 
                            +<span class="sd">    @param architecture: name of the model&#39;s architecture</span>
                
 
                            +<span class="sd">    @param model: model to load pretrinaed weights for</span>
                
 
                            +<span class="sd">    @param pretrained_weights: path tp pretrained weights</span>
                
 
                            +<span class="sd">    @return: None</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +
                
 
                            +    <span class="n">map_location</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">&quot;cpu&quot;</span><span class="p">)</span>
                
 
                            +
                
 
                            +    <span class="n">pretrained_state_dict</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">pretrained_weights</span><span class="p">,</span> <span class="n">map_location</span><span class="o">=</span><span class="n">map_location</span><span class="p">)</span>
                
 
                            +    <span class="n">_load_weights</span><span class="p">(</span><span class="n">architecture</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">pretrained_state_dict</span><span class="p">)</span>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -389,4 +434,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.detection_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.detection_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.detection_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">math</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">pathlib</span>
<span class="kn">from</span> <span class="nn">abc</span> <span class="kn">import</span> <span class="n">ABC</span><span class="p">,</span> <span class="n">abstractmethod</span>
<span class="kn">from</span> <span class="nn">enum</span> <span class="kn">import</span> <span class="n">Enum</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Dict</span>

<span class="kn">import</span> <span class="nn">cv2</span>
<span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="nn">plt</span>

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">torchvision</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>
<span class="kn">from</span> <span class="nn">torch.utils.data._utils.collate</span> <span class="kn">import</span> <span class="n">default_collate</span>
<span class="kn">from</span> <span class="nn">omegaconf</span> <span class="kn">import</span> <span class="n">ListConfig</span>


<div class="viewcode-block" id="DetectionTargetsFormat"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.DetectionTargetsFormat">[docs]</a><span class="k">class</span> <span class="nc">DetectionTargetsFormat</span><span class="p">(</span><span class="n">Enum</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Enum class for the different detection output formats</span>

<span class="sd">    When NORMALIZED is not specified- the type refers to unnormalized image coordinates (of the bboxes).</span>

<span class="sd">    For example:</span>
<span class="sd">    LABEL_NORMALIZED_XYXY means [class_idx,x1,y1,x2,y2]</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">LABEL_XYXY</span> <span class="o">=</span> <span class="s2">&quot;LABEL_XYXY&quot;</span>
    <span class="n">XYXY_LABEL</span> <span class="o">=</span> <span class="s2">&quot;XYXY_LABEL&quot;</span>
    <span class="n">LABEL_NORMALIZED_XYXY</span> <span class="o">=</span> <span class="s2">&quot;LABEL_NORMALIZED_XYXY&quot;</span>
    <span class="n">NORMALIZED_XYXY_LABEL</span> <span class="o">=</span> <span class="s2">&quot;NORMALIZED_XYXY_LABEL&quot;</span>
    <span class="n">LABEL_CXCYWH</span> <span class="o">=</span> <span class="s2">&quot;LABEL_CXCYWH&quot;</span>
    <span class="n">CXCYWH_LABEL</span> <span class="o">=</span> <span class="s2">&quot;CXCYWH_LABEL&quot;</span>
    <span class="n">LABEL_NORMALIZED_CXCYWH</span> <span class="o">=</span> <span class="s2">&quot;LABEL_NORMALIZED_CXCYWH&quot;</span>
    <span class="n">NORMALIZED_CXCYWH_LABEL</span> <span class="o">=</span> <span class="s2">&quot;NORMALIZED_CXCYWH_LABEL&quot;</span></div>


<div class="viewcode-block" id="get_cls_posx_in_target"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.get_cls_posx_in_target">[docs]</a><span class="k">def</span> <span class="nf">get_cls_posx_in_target</span><span class="p">(</span><span class="n">target_format</span><span class="p">:</span> <span class="n">DetectionTargetsFormat</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Get the label of a given target</span>
<span class="sd">    :param target_format:   Representation of the target (ex: LABEL_XYXY)</span>
<span class="sd">    :return:                Position of the class id in a bbox</span>
<span class="sd">                                ex: 0 if bbox of format label_xyxy | -1 if bbox of format xyxy_label</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">format_split</span> <span class="o">=</span> <span class="n">target_format</span><span class="o">.</span><span class="n">value</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;_&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">format_split</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;LABEL&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="mi">0</span>
    <span class="k">elif</span> <span class="n">format_split</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;LABEL&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="o">-</span><span class="mi">1</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;No implementation to find index of LABEL in </span><span class="si">{</span><span class="n">target_format</span><span class="o">.</span><span class="n">value</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_set_batch_labels_index</span><span class="p">(</span><span class="n">labels_batch</span><span class="p">):</span>
    <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">labels</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">labels_batch</span><span class="p">):</span>
        <span class="n">labels</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">i</span>
    <span class="k">return</span> <span class="n">labels_batch</span>


<div class="viewcode-block" id="convert_xywh_bbox_to_xyxy"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.convert_xywh_bbox_to_xyxy">[docs]</a><span class="k">def</span> <span class="nf">convert_xywh_bbox_to_xyxy</span><span class="p">(</span><span class="n">input_bbox</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Converts bounding box format from [x, y, w, h] to [x1, y1, x2, y2]</span>
<span class="sd">        :param input_bbox:  input bbox either 2-dimensional (for all boxes of a single image) or 3-dimensional (for</span>
<span class="sd">                            boxes of a batch of images)</span>
<span class="sd">        :return:            Converted bbox in same dimensions as the original</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">need_squeeze</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="c1"># the input is always processed as a batch. in case it not a batch, it is unsqueezed, process and than squeeze back.</span>
    <span class="k">if</span> <span class="n">input_bbox</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">&lt;</span> <span class="mi">3</span><span class="p">:</span>
        <span class="n">need_squeeze</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="n">input_bbox</span> <span class="o">=</span> <span class="n">input_bbox</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="n">converted_bbox</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">input_bbox</span><span class="p">)</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">input_bbox</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="k">else</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">input_bbox</span><span class="p">)</span>
    <span class="n">converted_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>
    <span class="n">converted_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>
    <span class="n">converted_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>
    <span class="n">converted_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">input_bbox</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>

    <span class="c1"># squeeze back if needed</span>
    <span class="k">if</span> <span class="n">need_squeeze</span><span class="p">:</span>
        <span class="n">converted_bbox</span> <span class="o">=</span> <span class="n">converted_bbox</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>

    <span class="k">return</span> <span class="n">converted_bbox</span></div>


<span class="k">def</span> <span class="nf">_iou</span><span class="p">(</span><span class="n">CIoU</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">DIoU</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">GIoU</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_x2</span><span class="p">,</span> <span class="n">b1_y1</span><span class="p">,</span> <span class="n">b1_y2</span><span class="p">,</span> <span class="n">b2_x1</span><span class="p">,</span> <span class="n">b2_x2</span><span class="p">,</span> <span class="n">b2_y1</span><span class="p">,</span> <span class="n">b2_y2</span><span class="p">,</span> <span class="n">eps</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Internal function for the use of calculate_bbox_iou_matrix and calculate_bbox_iou_elementwise functions</span>
<span class="sd">    DO NOT CALL THIS FUNCTIONS DIRECTLY - use one of the functions mentioned above</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># Intersection area</span>
    <span class="n">intersection_area</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">b1_x2</span><span class="p">,</span> <span class="n">b2_x2</span><span class="p">)</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">b1_x1</span><span class="p">,</span> <span class="n">b2_x1</span><span class="p">))</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">*</span> \
                        <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">b1_y2</span><span class="p">,</span> <span class="n">b2_y2</span><span class="p">)</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">b1_y1</span><span class="p">,</span> <span class="n">b2_y1</span><span class="p">))</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
    <span class="c1"># Union Area</span>
    <span class="n">w1</span><span class="p">,</span> <span class="n">h1</span> <span class="o">=</span> <span class="n">b1_x2</span> <span class="o">-</span> <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_y2</span> <span class="o">-</span> <span class="n">b1_y1</span>
    <span class="n">w2</span><span class="p">,</span> <span class="n">h2</span> <span class="o">=</span> <span class="n">b2_x2</span> <span class="o">-</span> <span class="n">b2_x1</span><span class="p">,</span> <span class="n">b2_y2</span> <span class="o">-</span> <span class="n">b2_y1</span>
    <span class="n">union_area</span> <span class="o">=</span> <span class="n">w1</span> <span class="o">*</span> <span class="n">h1</span> <span class="o">+</span> <span class="n">w2</span> <span class="o">*</span> <span class="n">h2</span> <span class="o">-</span> <span class="n">intersection_area</span> <span class="o">+</span> <span class="n">eps</span>
    <span class="n">iou</span> <span class="o">=</span> <span class="n">intersection_area</span> <span class="o">/</span> <span class="n">union_area</span>  <span class="c1"># iou</span>
    <span class="k">if</span> <span class="n">GIoU</span> <span class="ow">or</span> <span class="n">DIoU</span> <span class="ow">or</span> <span class="n">CIoU</span><span class="p">:</span>
        <span class="n">cw</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">b1_x2</span><span class="p">,</span> <span class="n">b2_x2</span><span class="p">)</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">b1_x1</span><span class="p">,</span> <span class="n">b2_x1</span><span class="p">)</span>  <span class="c1"># convex (smallest enclosing box) width</span>
        <span class="n">ch</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">b1_y2</span><span class="p">,</span> <span class="n">b2_y2</span><span class="p">)</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">b1_y1</span><span class="p">,</span> <span class="n">b2_y1</span><span class="p">)</span>  <span class="c1"># convex height</span>
        <span class="c1"># Generalized IoU https://arxiv.org/pdf/1902.09630.pdf</span>
        <span class="k">if</span> <span class="n">GIoU</span><span class="p">:</span>
            <span class="n">c_area</span> <span class="o">=</span> <span class="n">cw</span> <span class="o">*</span> <span class="n">ch</span> <span class="o">+</span> <span class="n">eps</span>  <span class="c1"># convex area</span>
            <span class="n">iou</span> <span class="o">-=</span> <span class="p">(</span><span class="n">c_area</span> <span class="o">-</span> <span class="n">union_area</span><span class="p">)</span> <span class="o">/</span> <span class="n">c_area</span>  <span class="c1"># GIoU</span>
        <span class="c1"># Distance or Complete IoU https://arxiv.org/abs/1911.08287v1</span>
        <span class="k">if</span> <span class="n">DIoU</span> <span class="ow">or</span> <span class="n">CIoU</span><span class="p">:</span>
            <span class="c1"># convex diagonal squared</span>
            <span class="n">c2</span> <span class="o">=</span> <span class="n">cw</span> <span class="o">**</span> <span class="mi">2</span> <span class="o">+</span> <span class="n">ch</span> <span class="o">**</span> <span class="mi">2</span> <span class="o">+</span> <span class="n">eps</span>
            <span class="c1"># centerpoint distance squared</span>
            <span class="n">rho2</span> <span class="o">=</span> <span class="p">((</span><span class="n">b2_x1</span> <span class="o">+</span> <span class="n">b2_x2</span> <span class="o">-</span> <span class="n">b1_x1</span> <span class="o">-</span> <span class="n">b1_x2</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span> <span class="o">+</span> <span class="p">(</span><span class="n">b2_y1</span> <span class="o">+</span> <span class="n">b2_y2</span> <span class="o">-</span> <span class="n">b1_y1</span> <span class="o">-</span> <span class="n">b1_y2</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span><span class="p">)</span> <span class="o">/</span> <span class="mi">4</span>
            <span class="k">if</span> <span class="n">DIoU</span><span class="p">:</span>
                <span class="n">iou</span> <span class="o">-=</span> <span class="n">rho2</span> <span class="o">/</span> <span class="n">c2</span>  <span class="c1"># DIoU</span>
            <span class="k">elif</span> <span class="n">CIoU</span><span class="p">:</span>  <span class="c1"># https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47</span>
                <span class="n">v</span> <span class="o">=</span> <span class="p">(</span><span class="mi">4</span> <span class="o">/</span> <span class="n">math</span><span class="o">.</span><span class="n">pi</span> <span class="o">**</span> <span class="mi">2</span><span class="p">)</span> <span class="o">*</span> <span class="n">torch</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">atan</span><span class="p">(</span><span class="n">w2</span> <span class="o">/</span> <span class="n">h2</span><span class="p">)</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">atan</span><span class="p">(</span><span class="n">w1</span> <span class="o">/</span> <span class="n">h1</span><span class="p">),</span> <span class="mi">2</span><span class="p">)</span>
                <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">():</span>
                    <span class="n">alpha</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="n">eps</span><span class="p">)</span> <span class="o">-</span> <span class="n">iou</span> <span class="o">+</span> <span class="n">v</span><span class="p">)</span>
                <span class="n">iou</span> <span class="o">-=</span> <span class="p">(</span><span class="n">rho2</span> <span class="o">/</span> <span class="n">c2</span> <span class="o">+</span> <span class="n">v</span> <span class="o">*</span> <span class="n">alpha</span><span class="p">)</span>  <span class="c1"># CIoU</span>
    <span class="k">return</span> <span class="n">iou</span>


<div class="viewcode-block" id="calculate_bbox_iou_matrix"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.calculate_bbox_iou_matrix">[docs]</a><span class="k">def</span> <span class="nf">calculate_bbox_iou_matrix</span><span class="p">(</span><span class="n">box1</span><span class="p">,</span> <span class="n">box2</span><span class="p">,</span> <span class="n">x1y1x2y2</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">GIoU</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">DIoU</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">CIoU</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="mf">1e-9</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    calculate iou matrix containing the iou of every couple iuo(i,j) where i is in box1 and j is in box2</span>
<span class="sd">        :param box1: a 2D tensor of boxes (shape N x 4)</span>
<span class="sd">        :param box2: a 2D tensor of boxes (shape M x 4)</span>
<span class="sd">        :param x1y1x2y2: boxes format is x1y1x2y2 (True) or xywh where xy is the center (False)</span>
<span class="sd">        :return: a 2D iou matrix (shape NxM)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">box1</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">box1</span> <span class="o">=</span> <span class="n">box1</span><span class="o">.</span><span class="n">T</span>

    <span class="c1"># Get the coordinates of bounding boxes</span>
    <span class="k">if</span> <span class="n">x1y1x2y2</span><span class="p">:</span>  <span class="c1"># x1, y1, x2, y2 = box1</span>
        <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_y1</span><span class="p">,</span> <span class="n">b1_x2</span><span class="p">,</span> <span class="n">b1_y2</span> <span class="o">=</span> <span class="n">box1</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">box1</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">box1</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span> <span class="n">box1</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
        <span class="n">b2_x1</span><span class="p">,</span> <span class="n">b2_y1</span><span class="p">,</span> <span class="n">b2_x2</span><span class="p">,</span> <span class="n">b2_y2</span> <span class="o">=</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">],</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>  <span class="c1"># x, y, w, h = box1</span>
        <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_x2</span> <span class="o">=</span> <span class="n">box1</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">box1</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">,</span> <span class="n">box1</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">box1</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>
        <span class="n">b1_y1</span><span class="p">,</span> <span class="n">b1_y2</span> <span class="o">=</span> <span class="n">box1</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">box1</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">,</span> <span class="n">box1</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">box1</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>
        <span class="n">b2_x1</span><span class="p">,</span> <span class="n">b2_x2</span> <span class="o">=</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">,</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>
        <span class="n">b2_y1</span><span class="p">,</span> <span class="n">b2_y2</span> <span class="o">=</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span><span class="p">,</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">/</span> <span class="mi">2</span>

    <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_y1</span><span class="p">,</span> <span class="n">b1_x2</span><span class="p">,</span> <span class="n">b1_y2</span> <span class="o">=</span> <span class="n">b1_x1</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="n">b1_y1</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="n">b1_x2</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="n">b1_y2</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">_iou</span><span class="p">(</span><span class="n">CIoU</span><span class="p">,</span> <span class="n">DIoU</span><span class="p">,</span> <span class="n">GIoU</span><span class="p">,</span> <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_x2</span><span class="p">,</span> <span class="n">b1_y1</span><span class="p">,</span> <span class="n">b1_y2</span><span class="p">,</span> <span class="n">b2_x1</span><span class="p">,</span> <span class="n">b2_x2</span><span class="p">,</span> <span class="n">b2_y1</span><span class="p">,</span> <span class="n">b2_y2</span><span class="p">,</span> <span class="n">eps</span><span class="p">)</span></div>


<div class="viewcode-block" id="calc_bbox_iou_matrix"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.calc_bbox_iou_matrix">[docs]</a><span class="k">def</span> <span class="nf">calc_bbox_iou_matrix</span><span class="p">(</span><span class="n">pred</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    calculate iou for every pair of boxes in the boxes vector</span>
<span class="sd">    :param pred: a 3-dimensional tensor containing all boxes for a batch of images [N, num_boxes, 4], where</span>
<span class="sd">                 each box format is [x1,y1,x2,y2]</span>
<span class="sd">    :return: a 3-dimensional matrix where M_i_j_k is the iou of box j and box k of the i&#39;th image in the batch</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">box</span> <span class="o">=</span> <span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span>  <span class="c1">#</span>
    <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_y1</span> <span class="o">=</span> <span class="n">box</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="n">box</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">b1_x2</span><span class="p">,</span> <span class="n">b1_y2</span> <span class="o">=</span> <span class="n">box</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">2</span><span class="p">]</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="n">box</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">3</span><span class="p">]</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="n">b2_x1</span> <span class="o">=</span> <span class="n">b1_x1</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">b2_x2</span> <span class="o">=</span> <span class="n">b1_x2</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">b2_y1</span> <span class="o">=</span> <span class="n">b1_y1</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">b2_y2</span> <span class="o">=</span> <span class="n">b1_y2</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">intersection_area</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">b1_x2</span><span class="p">,</span> <span class="n">b2_x2</span><span class="p">)</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">b1_x1</span><span class="p">,</span> <span class="n">b2_x1</span><span class="p">))</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">*</span> \
                        <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">b1_y2</span><span class="p">,</span> <span class="n">b2_y2</span><span class="p">)</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">b1_y1</span><span class="p">,</span> <span class="n">b2_y1</span><span class="p">))</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
    <span class="c1"># Union Area</span>
    <span class="n">w1</span><span class="p">,</span> <span class="n">h1</span> <span class="o">=</span> <span class="n">b1_x2</span> <span class="o">-</span> <span class="n">b1_x1</span><span class="p">,</span> <span class="n">b1_y2</span> <span class="o">-</span> <span class="n">b1_y1</span>
    <span class="n">w2</span><span class="p">,</span> <span class="n">h2</span> <span class="o">=</span> <span class="n">b2_x2</span> <span class="o">-</span> <span class="n">b2_x1</span><span class="p">,</span> <span class="n">b2_y2</span> <span class="o">-</span> <span class="n">b2_y1</span>
    <span class="n">union_area</span> <span class="o">=</span> <span class="p">(</span><span class="n">w1</span> <span class="o">*</span> <span class="n">h1</span> <span class="o">+</span> <span class="mf">1e-16</span><span class="p">)</span> <span class="o">+</span> <span class="n">w2</span> <span class="o">*</span> <span class="n">h2</span> <span class="o">-</span> <span class="n">intersection_area</span>
    <span class="n">ious</span> <span class="o">=</span> <span class="n">intersection_area</span> <span class="o">/</span> <span class="n">union_area</span>
    <span class="k">return</span> <span class="n">ious</span></div>


<div class="viewcode-block" id="change_bbox_bounds_for_image_size"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.change_bbox_bounds_for_image_size">[docs]</a><span class="k">def</span> <span class="nf">change_bbox_bounds_for_image_size</span><span class="p">(</span><span class="n">boxes</span><span class="p">,</span> <span class="n">img_shape</span><span class="p">):</span>
    <span class="c1"># CLIP BOUNDING XYXY BOUNDING BOXES TO IMAGE SHAPE (HEIGHT, WIDTH)</span>
    <span class="n">boxes</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="nb">min</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="nb">max</span><span class="o">=</span><span class="n">img_shape</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
    <span class="n">boxes</span><span class="p">[:,</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="o">=</span> <span class="n">boxes</span><span class="p">[:,</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="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="nb">min</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="nb">max</span><span class="o">=</span><span class="n">img_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
    <span class="k">return</span> <span class="n">boxes</span></div>


<div class="viewcode-block" id="DetectionPostPredictionCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.DetectionPostPredictionCallback">[docs]</a><span class="k">class</span> <span class="nc">DetectionPostPredictionCallback</span><span class="p">(</span><span class="n">ABC</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

<div class="viewcode-block" id="DetectionPostPredictionCallback.forward"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.DetectionPostPredictionCallback.forward">[docs]</a>    <span class="nd">@abstractmethod</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">device</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">        :param x:       the output of your model</span>
<span class="sd">        :param device:  the device to move all output tensors into</span>
<span class="sd">        :return:        a list with length batch_size, each item in the list is a detections</span>
<span class="sd">                        with shape: nx6 (x1, y1, x2, y2, confidence, class) where x and y are in range [0,1]</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div></div>


<div class="viewcode-block" id="IouThreshold"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.IouThreshold">[docs]</a><span class="k">class</span> <span class="nc">IouThreshold</span><span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
    <span class="n">MAP_05</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">)</span>
    <span class="n">MAP_05_TO_095</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">0.95</span><span class="p">)</span>

<div class="viewcode-block" id="IouThreshold.is_range"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.IouThreshold.is_range">[docs]</a>    <span class="k">def</span> <span class="nf">is_range</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="bp">self</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span></div>

<div class="viewcode-block" id="IouThreshold.to_tensor"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.IouThreshold.to_tensor">[docs]</a>    <span class="k">def</span> <span class="nf">to_tensor</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">is_range</span><span class="p">():</span>
            <span class="n">n_iou_thresh</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="nb">round</span><span class="p">((</span><span class="bp">self</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="bp">self</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">/</span> <span class="mf">0.05</span><span class="p">))</span> <span class="o">+</span> <span class="mi">1</span>
            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="bp">self</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="bp">self</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">n_iou_thresh</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">n_iou_thresh</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">([</span><span class="bp">self</span><span class="p">[</span><span class="mi">0</span><span class="p">]])</span></div></div>


<div class="viewcode-block" id="box_iou"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.box_iou">[docs]</a><span class="k">def</span> <span class="nf">box_iou</span><span class="p">(</span><span class="n">box1</span><span class="p">,</span> <span class="n">box2</span><span class="p">):</span>
    <span class="c1"># https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Return intersection-over-union (Jaccard index) of boxes.</span>
<span class="sd">    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.</span>
<span class="sd">    Arguments:</span>
<span class="sd">        box1 (Tensor[N, 4])</span>
<span class="sd">        box2 (Tensor[M, 4])</span>
<span class="sd">    Returns:</span>
<span class="sd">        iou (Tensor[N, M]): the NxM matrix containing the pairwise</span>
<span class="sd">            IoU values for every element in boxes1 and boxes2</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">box_area</span><span class="p">(</span><span class="n">box</span><span class="p">):</span>
        <span class="c1"># box = 4xn</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">box</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">box</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">*</span> <span class="p">(</span><span class="n">box</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">-</span> <span class="n">box</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>

    <span class="n">area1</span> <span class="o">=</span> <span class="n">box_area</span><span class="p">(</span><span class="n">box1</span><span class="o">.</span><span class="n">T</span><span class="p">)</span>
    <span class="n">area2</span> <span class="o">=</span> <span class="n">box_area</span><span class="p">(</span><span class="n">box2</span><span class="o">.</span><span class="n">T</span><span class="p">)</span>

    <span class="c1"># inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)</span>
    <span class="n">inter</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">box1</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">,</span> <span class="mi">2</span><span class="p">:],</span> <span class="n">box2</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:])</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">box1</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">,</span> <span class="p">:</span><span class="mi">2</span><span class="p">],</span> <span class="n">box2</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]))</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">inter</span> <span class="o">/</span> <span class="p">(</span><span class="n">area1</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">+</span> <span class="n">area2</span> <span class="o">-</span> <span class="n">inter</span><span class="p">)</span>  <span class="c1"># iou = inter / (area1 + area2 - inter)</span></div>


<div class="viewcode-block" id="non_max_suppression"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.non_max_suppression">[docs]</a><span class="k">def</span> <span class="nf">non_max_suppression</span><span class="p">(</span><span class="n">prediction</span><span class="p">,</span> <span class="n">conf_thres</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">iou_thres</span><span class="o">=</span><span class="mf">0.6</span><span class="p">,</span>
                        <span class="n">multi_label_per_box</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> <span class="n">with_confidence</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Performs Non-Maximum Suppression (NMS) on inference results</span>
<span class="sd">        :param prediction: raw model prediction</span>
<span class="sd">        :param conf_thres: below the confidence threshold - prediction are discarded</span>
<span class="sd">        :param iou_thres: IoU threshold for the nms algorithm</span>
<span class="sd">        :param multi_label_per_box: whether to use re-use each box with all possible labels</span>
<span class="sd">                                    (instead of the maximum confidence all confidences above threshold</span>
<span class="sd">                                    will be sent to NMS); by default is set to True</span>
<span class="sd">        :param with_confidence: whether to multiply objectness score with class score.</span>
<span class="sd">                                usually valid for Yolo models only.</span>
<span class="sd">        :return:  (x1, y1, x2, y2, object_conf, class_conf, class)</span>
<span class="sd">    Returns:</span>
<span class="sd">         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">candidates_above_thres</span> <span class="o">=</span> <span class="n">prediction</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">&gt;</span> <span class="n">conf_thres</span>  <span class="c1"># filter by confidence</span>
    <span class="n">output</span> <span class="o">=</span> <span class="p">[</span><span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="n">prediction</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>

    <span class="k">for</span> <span class="n">image_idx</span><span class="p">,</span> <span class="n">pred</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">prediction</span><span class="p">):</span>

        <span class="n">pred</span> <span class="o">=</span> <span class="n">pred</span><span class="p">[</span><span class="n">candidates_above_thres</span><span class="p">[</span><span class="n">image_idx</span><span class="p">]]</span>  <span class="c1"># confident</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="n">pred</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]:</span>  <span class="c1"># If none remain process next image</span>
            <span class="k">continue</span>

        <span class="k">if</span> <span class="n">with_confidence</span><span class="p">:</span>
            <span class="n">pred</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">:]</span> <span class="o">*=</span> <span class="n">pred</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">:</span><span class="mi">5</span><span class="p">]</span>  <span class="c1"># multiply objectness score with class score</span>

        <span class="n">box</span> <span class="o">=</span> <span class="n">convert_xywh_bbox_to_xyxy</span><span class="p">(</span><span class="n">pred</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">])</span>  <span class="c1"># xywh to xyxy</span>

        <span class="c1"># Detections matrix nx6 (xyxy, conf, cls)</span>
        <span class="k">if</span> <span class="n">multi_label_per_box</span><span class="p">:</span>  <span class="c1"># try for all good confidence classes</span>
            <span class="n">i</span><span class="p">,</span> <span class="n">j</span> <span class="o">=</span> <span class="p">(</span><span class="n">pred</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">:]</span> <span class="o">&gt;</span> <span class="n">conf_thres</span><span class="p">)</span><span class="o">.</span><span class="n">nonzero</span><span class="p">(</span><span class="n">as_tuple</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span><span class="o">.</span><span class="n">T</span>
            <span class="n">pred</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">box</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="n">pred</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span> <span class="o">+</span> <span class="mi">5</span><span class="p">,</span> <span class="kc">None</span><span class="p">],</span> <span class="n">j</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span><span class="o">.</span><span class="n">float</span><span class="p">()),</span> <span class="mi">1</span><span class="p">)</span>

        <span class="k">else</span><span class="p">:</span>  <span class="c1"># best class only</span>
            <span class="n">conf</span><span class="p">,</span> <span class="n">j</span> <span class="o">=</span> <span class="n">pred</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">:]</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
            <span class="n">pred</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">box</span><span class="p">,</span> <span class="n">conf</span><span class="p">,</span> <span class="n">j</span><span class="o">.</span><span class="n">float</span><span class="p">()),</span> <span class="mi">1</span><span class="p">)[</span><span class="n">conf</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="o">&gt;</span> <span class="n">conf_thres</span><span class="p">]</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="n">pred</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]:</span>  <span class="c1"># If none remain process next image</span>
            <span class="k">continue</span>

        <span class="c1"># Apply torch batched NMS algorithm</span>
        <span class="n">boxes</span><span class="p">,</span> <span class="n">scores</span><span class="p">,</span> <span class="n">cls_idx</span> <span class="o">=</span> <span class="n">pred</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">],</span> <span class="n">pred</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">],</span> <span class="n">pred</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">]</span>
        <span class="n">idx_to_keep</span> <span class="o">=</span> <span class="n">torchvision</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">boxes</span><span class="o">.</span><span class="n">batched_nms</span><span class="p">(</span><span class="n">boxes</span><span class="p">,</span> <span class="n">scores</span><span class="p">,</span> <span class="n">cls_idx</span><span class="p">,</span> <span class="n">iou_thres</span><span class="p">)</span>
        <span class="n">output</span><span class="p">[</span><span class="n">image_idx</span><span class="p">]</span> <span class="o">=</span> <span class="n">pred</span><span class="p">[</span><span class="n">idx_to_keep</span><span class="p">]</span>

    <span class="k">return</span> <span class="n">output</span></div>


<div class="viewcode-block" id="matrix_non_max_suppression"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.matrix_non_max_suppression">[docs]</a><span class="k">def</span> <span class="nf">matrix_non_max_suppression</span><span class="p">(</span><span class="n">pred</span><span class="p">,</span> <span class="n">conf_thres</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span> <span class="n">kernel</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;gaussian&#39;</span><span class="p">,</span>
                               <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">3.0</span><span class="p">,</span> <span class="n">max_num_of_detections</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">500</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Performs Matrix Non-Maximum Suppression (NMS) on inference results</span>
<span class="sd">        https://arxiv.org/pdf/1912.04488.pdf</span>
<span class="sd">        :param pred: raw model prediction (in test mode) - a Tensor of shape [batch, num_predictions, 85]</span>
<span class="sd">        where each item format is (x, y, w, h, object_conf, class_conf, ... 80 classes score ...)</span>
<span class="sd">        :param conf_thres: below the confidence threshold - prediction are discarded</span>
<span class="sd">        :param kernel: type of kernel to use [&#39;gaussian&#39;, &#39;linear&#39;]</span>
<span class="sd">        :param sigma: sigma for the gussian kernel</span>
<span class="sd">        :param max_num_of_detections: maximum number of boxes to output</span>
<span class="sd">        :return:  list of (x1, y1, x2, y2, object_conf, class_conf, class)</span>

<span class="sd">    Returns:</span>
<span class="sd">         detections list with shape: (x1, y1, x2, y2, conf, cls)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># MULTIPLY CONF BY CLASS CONF TO GET COMBINED CONFIDENCE</span>
    <span class="n">class_conf</span><span class="p">,</span> <span class="n">class_pred</span> <span class="o">=</span> <span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">5</span><span class="p">:]</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
    <span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">class_conf</span>

    <span class="c1"># BOX (CENTER X, CENTER Y, WIDTH, HEIGHT) TO (X1, Y1, X2, Y2)</span>
    <span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">=</span> <span class="n">convert_xywh_bbox_to_xyxy</span><span class="p">(</span><span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">])</span>

    <span class="c1"># DETECTIONS ORDERED AS (x1y1x2y2, obj_conf, class_conf, class_pred)</span>
    <span class="n">pred</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="p">:</span><span class="mi">5</span><span class="p">],</span> <span class="n">class_pred</span><span class="o">.</span><span class="n">unsqueeze</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="c1"># SORT DETECTIONS BY DECREASING CONFIDENCE SCORES</span>
    <span class="n">sort_ind</span> <span class="o">=</span> <span class="p">(</span><span class="o">-</span><span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">4</span><span class="p">])</span><span class="o">.</span><span class="n">argsort</span><span class="p">()</span>
    <span class="n">pred</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">pred</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">sort_ind</span><span class="p">[</span><span class="n">i</span><span class="p">]]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">pred</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])])[:,</span> <span class="mi">0</span><span class="p">:</span><span class="n">max_num_of_detections</span><span class="p">]</span>

    <span class="n">ious</span> <span class="o">=</span> <span class="n">calc_bbox_iou_matrix</span><span class="p">(</span><span class="n">pred</span><span class="p">)</span>

    <span class="n">ious</span> <span class="o">=</span> <span class="n">ious</span><span class="o">.</span><span class="n">triu</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="c1"># CREATE A LABELS MASK, WE WANT ONLY BOXES WITH THE SAME LABEL TO AFFECT EACH OTHER</span>
    <span class="n">labels</span> <span class="o">=</span> <span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">5</span><span class="p">:]</span>
    <span class="n">labeles_matrix</span> <span class="o">=</span> <span class="p">(</span><span class="n">labels</span> <span class="o">==</span> <span class="n">labels</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">triu</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="n">ious</span> <span class="o">*=</span> <span class="n">labeles_matrix</span>
    <span class="n">ious_cmax</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">ious</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">ious_cmax</span> <span class="o">=</span> <span class="n">ious_cmax</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">max_num_of_detections</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">kernel</span> <span class="o">==</span> <span class="s1">&#39;gaussian&#39;</span><span class="p">:</span>
        <span class="n">decay_matrix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span> <span class="o">*</span> <span class="n">sigma</span> <span class="o">*</span> <span class="p">(</span><span class="n">ious</span> <span class="o">**</span> <span class="mi">2</span><span class="p">))</span>
        <span class="n">compensate_matrix</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span> <span class="o">*</span> <span class="n">sigma</span> <span class="o">*</span> <span class="p">(</span><span class="n">ious_cmax</span> <span class="o">**</span> <span class="mi">2</span><span class="p">))</span>
        <span class="n">decay</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="p">(</span><span class="n">decay_matrix</span> <span class="o">/</span> <span class="n">compensate_matrix</span><span class="p">)</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">decay</span> <span class="o">=</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">ious</span><span class="p">)</span> <span class="o">/</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">ious_cmax</span><span class="p">)</span>
        <span class="n">decay</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">decay</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>

    <span class="n">pred</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">decay</span>

    <span class="n">output</span> <span class="o">=</span> <span class="p">[</span><span class="n">pred</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">pred</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">&gt;</span> <span class="n">conf_thres</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">pred</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])]</span>

    <span class="k">return</span> <span class="n">output</span></div>


<div class="viewcode-block" id="NMS_Type"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.NMS_Type">[docs]</a><span class="k">class</span> <span class="nc">NMS_Type</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Enum</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Type of non max suppression algorithm that can be used for post processing detection</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">ITERATIVE</span> <span class="o">=</span> <span class="s1">&#39;iterative&#39;</span>
    <span class="n">MATRIX</span> <span class="o">=</span> <span class="s1">&#39;matrix&#39;</span></div>


<div class="viewcode-block" id="undo_image_preprocessing"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.undo_image_preprocessing">[docs]</a><span class="k">def</span> <span class="nf">undo_image_preprocessing</span><span class="p">(</span><span class="n">im_tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    :param im_tensor: images in a batch after preprocessing for inference, RGB, (B, C, H, W)</span>
<span class="sd">    :return:          images in a batch in cv2 format, BGR, (B, H, W, C)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">im_np</span> <span class="o">=</span> <span class="n">im_tensor</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>
    <span class="n">im_np</span> <span class="o">=</span> <span class="n">im_np</span><span class="p">[:,</span> <span class="p">::</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:]</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">im_np</span> <span class="o">*=</span> <span class="mf">255.</span>
    <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">ascontiguousarray</span><span class="p">(</span><span class="n">im_np</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span></div>


<div class="viewcode-block" id="DetectionVisualization"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.DetectionVisualization">[docs]</a><span class="k">class</span> <span class="nc">DetectionVisualization</span><span class="p">:</span>
    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_generate_color_mapping</span><span class="p">(</span><span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">]]:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Generate a unique BGR color for each class</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">cmap</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">cm</span><span class="o">.</span><span class="n">get_cmap</span><span class="p">(</span><span class="s1">&#39;gist_rainbow&#39;</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">)</span>
        <span class="n">colors</span> <span class="o">=</span> <span class="p">[</span><span class="n">cmap</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="nb">bytes</span><span class="o">=</span><span class="kc">True</span><span class="p">)[:</span><span class="mi">3</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">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">num_classes</span><span class="p">)]</span>
        <span class="k">return</span> <span class="p">[</span><span class="nb">tuple</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">v</span><span class="p">)</span> <span class="k">for</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">c</span><span class="p">)</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="n">colors</span><span class="p">]</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_draw_box_title</span><span class="p">(</span><span class="n">color_mapping</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span> <span class="n">class_names</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="n">box_thickness</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                        <span class="n">image_np</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">x1</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">y1</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">x2</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">y2</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">class_id</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                        <span class="n">pred_conf</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">is_target</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
        <span class="n">color</span> <span class="o">=</span> <span class="n">color_mapping</span><span class="p">[</span><span class="n">class_id</span><span class="p">]</span>
        <span class="n">class_name</span> <span class="o">=</span> <span class="n">class_names</span><span class="p">[</span><span class="n">class_id</span><span class="p">]</span>

        <span class="c1"># Draw the box</span>
        <span class="n">image_np</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">rectangle</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">),</span> <span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="n">y2</span><span class="p">),</span> <span class="n">color</span><span class="p">,</span> <span class="n">box_thickness</span><span class="p">)</span>

        <span class="c1"># Caption with class name and confidence if given</span>
        <span class="n">text_color</span> <span class="o">=</span> <span class="p">(</span><span class="mi">255</span><span class="p">,</span> <span class="mi">255</span><span class="p">,</span> <span class="mi">255</span><span class="p">)</span>  <span class="c1"># white</span>

        <span class="k">if</span> <span class="n">is_target</span><span class="p">:</span>
            <span class="n">title</span> <span class="o">=</span> <span class="sa">f</span><span class="s1">&#39;[GT] </span><span class="si">{</span><span class="n">class_name</span><span class="si">}</span><span class="s1">&#39;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">is_target</span><span class="p">:</span>
            <span class="n">title</span> <span class="o">=</span> <span class="sa">f</span><span class="s1">&#39;[Pred] </span><span class="si">{</span><span class="n">class_name</span><span class="si">}</span><span class="s1">  </span><span class="si">{</span><span class="nb">str</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">pred_conf</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span> <span class="k">if</span> <span class="n">pred_conf</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="s2">&quot;&quot;</span><span class="si">}</span><span class="s1">&#39;</span>

        <span class="n">image_np</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">rectangle</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">y1</span> <span class="o">-</span> <span class="mi">15</span><span class="p">),</span> <span class="p">(</span><span class="n">x1</span> <span class="o">+</span> <span class="nb">len</span><span class="p">(</span><span class="n">title</span><span class="p">)</span> <span class="o">*</span> <span class="mi">10</span><span class="p">,</span> <span class="n">y1</span><span class="p">),</span> <span class="n">color</span><span class="p">,</span> <span class="n">cv2</span><span class="o">.</span><span class="n">FILLED</span><span class="p">)</span>
        <span class="n">image_np</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">putText</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="n">title</span><span class="p">,</span> <span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">y1</span> <span class="o">-</span> <span class="n">box_thickness</span><span class="p">),</span> <span class="mi">2</span><span class="p">,</span> <span class="mf">.5</span><span class="p">,</span> <span class="n">text_color</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">lineType</span><span class="o">=</span><span class="n">cv2</span><span class="o">.</span><span class="n">LINE_AA</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">image_np</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_visualize_image</span><span class="p">(</span><span class="n">image_np</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">pred_boxes</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">target_boxes</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span>
                         <span class="n">class_names</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="n">box_thickness</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">gt_alpha</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">image_scale</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
                         <span class="n">checkpoint_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">image_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">image_np</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">fx</span><span class="o">=</span><span class="n">image_scale</span><span class="p">,</span> <span class="n">fy</span><span class="o">=</span><span class="n">image_scale</span><span class="p">,</span> <span class="n">interpolation</span><span class="o">=</span><span class="n">cv2</span><span class="o">.</span><span class="n">INTER_NEAREST</span><span class="p">)</span>
        <span class="n">color_mapping</span> <span class="o">=</span> <span class="n">DetectionVisualization</span><span class="o">.</span><span class="n">_generate_color_mapping</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">class_names</span><span class="p">))</span>

        <span class="c1"># Draw predictions</span>
        <span class="n">pred_boxes</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">4</span><span class="p">]</span> <span class="o">*=</span> <span class="n">image_scale</span>
        <span class="k">for</span> <span class="n">box</span> <span class="ow">in</span> <span class="n">pred_boxes</span><span class="p">:</span>
            <span class="n">image_np</span> <span class="o">=</span> <span class="n">DetectionVisualization</span><span class="o">.</span><span class="n">_draw_box_title</span><span class="p">(</span><span class="n">color_mapping</span><span class="p">,</span> <span class="n">class_names</span><span class="p">,</span> <span class="n">box_thickness</span><span class="p">,</span>
                                                              <span class="n">image_np</span><span class="p">,</span> <span class="o">*</span><span class="n">box</span><span class="p">[:</span><span class="mi">4</span><span class="p">]</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">),</span>
                                                              <span class="n">class_id</span><span class="o">=</span><span class="nb">int</span><span class="p">(</span><span class="n">box</span><span class="p">[</span><span class="mi">5</span><span class="p">]),</span> <span class="n">pred_conf</span><span class="o">=</span><span class="n">box</span><span class="p">[</span><span class="mi">4</span><span class="p">])</span>

        <span class="c1"># Draw ground truths</span>
        <span class="n">target_boxes_image</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">box</span> <span class="ow">in</span> <span class="n">target_boxes</span><span class="p">:</span>
            <span class="n">target_boxes_image</span> <span class="o">=</span> <span class="n">DetectionVisualization</span><span class="o">.</span><span class="n">_draw_box_title</span><span class="p">(</span><span class="n">color_mapping</span><span class="p">,</span> <span class="n">class_names</span><span class="p">,</span> <span class="n">box_thickness</span><span class="p">,</span>
                                                                        <span class="n">target_boxes_image</span><span class="p">,</span> <span class="o">*</span><span class="n">box</span><span class="p">[</span><span class="mi">2</span><span class="p">:],</span>
                                                                        <span class="n">class_id</span><span class="o">=</span><span class="n">box</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">is_target</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

        <span class="c1"># Transparent overlay of ground truth boxes</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">target_boxes_image</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">bool</span><span class="p">)</span>
        <span class="n">image_np</span><span class="p">[</span><span class="n">mask</span><span class="p">]</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">addWeighted</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">gt_alpha</span><span class="p">,</span> <span class="n">target_boxes_image</span><span class="p">,</span> <span class="n">gt_alpha</span><span class="p">,</span> <span class="mi">0</span><span class="p">)[</span><span class="n">mask</span><span class="p">]</span>

        <span class="k">if</span> <span class="n">checkpoint_dir</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">image_np</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">pathlib</span><span class="o">.</span><span class="n">Path</span><span class="p">(</span><span class="n">checkpoint_dir</span><span class="p">)</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">parents</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
            <span class="n">cv2</span><span class="o">.</span><span class="n">imwrite</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">checkpoint_dir</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">image_name</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;.jpg&#39;</span><span class="p">),</span> <span class="n">image_np</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_scaled_ccwh_to_xyxy</span><span class="p">(</span><span class="n">target_boxes</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">h</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">w</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">image_scale</span><span class="p">:</span> <span class="nb">float</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Modifies target_boxes inplace</span>
<span class="sd">        :param target_boxes:    (c1, c2, w, h) boxes in [0, 1] range</span>
<span class="sd">        :param h:               image height</span>
<span class="sd">        :param w:               image width</span>
<span class="sd">        :param image_scale:     desired scale for the boxes w.r.t. w and h</span>
<span class="sd">        :return:                targets in (x1, y1, x2, y2) format</span>
<span class="sd">                                in range [0, w * self.image_scale] [0, h * self.image_scale]</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># unscale</span>
        <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</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="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">]])</span>

        <span class="c1"># x1 = c1 - w // 2; y1 = c2 - h // 2</span>
        <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">-=</span> <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">//</span> <span class="mi">2</span>
        <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">-=</span> <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">]</span> <span class="o">//</span> <span class="mi">2</span>
        <span class="c1"># x2 = w + x1; y2 = h + y1</span>
        <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">]</span> <span class="o">+=</span> <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span>
        <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">5</span><span class="p">]</span> <span class="o">+=</span> <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span>

        <span class="n">target_boxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:]</span> <span class="o">*=</span> <span class="n">image_scale</span>
        <span class="n">target_boxes</span> <span class="o">=</span> <span class="n">target_boxes</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">target_boxes</span>

<div class="viewcode-block" id="DetectionVisualization.visualize_batch"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.DetectionVisualization.visualize_batch">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">visualize_batch</span><span class="p">(</span><span class="n">image_tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">pred_boxes</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">target_boxes</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                        <span class="n">batch_name</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">str</span><span class="p">],</span> <span class="n">class_names</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="n">checkpoint_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                        <span class="n">undo_preprocessing_func</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="n">undo_image_preprocessing</span><span class="p">,</span>
                        <span class="n">box_thickness</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">2</span><span class="p">,</span> <span class="n">image_scale</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.</span><span class="p">,</span> <span class="n">gt_alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">.4</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        A helper function to visualize detections predicted by a network:</span>
<span class="sd">        saves images into a given path with a name that is {batch_name}_{imade_idx_in_the_batch}.jpg, one batch per call.</span>
<span class="sd">        Colors are generated on the fly: uniformly sampled from color wheel to support all given classes.</span>

<span class="sd">        Adjustable:</span>
<span class="sd">            * Ground truth box transparency;</span>
<span class="sd">            * Box width;</span>
<span class="sd">            * Image size (larger or smaller than what&#39;s provided)</span>

<span class="sd">        :param image_tensor:            rgb images, (B, H, W, 3)</span>
<span class="sd">        :param pred_boxes:              boxes after NMS for each image in a batch, each (Num_boxes, 6),</span>
<span class="sd">                                        values on dim 1 are: x1, y1, x2, y2, confidence, class</span>
<span class="sd">        :param target_boxes:            (Num_targets, 6), values on dim 1 are: image id in a batch, class, x y w h</span>
<span class="sd">                                        (coordinates scaled to [0, 1])</span>
<span class="sd">        :param batch_name:              id of the current batch to use for image naming</span>

<span class="sd">        :param class_names:             names of all classes, each on its own index</span>
<span class="sd">        :param checkpoint_dir:          a path where images with boxes will be saved. if None, the result images will</span>
<span class="sd">                                        be returns as a list of numpy image arrays</span>

<span class="sd">        :param undo_preprocessing_func: a function to convert preprocessed images tensor into a batch of cv2-like images</span>
<span class="sd">        :param box_thickness:           box line thickness in px</span>
<span class="sd">        :param image_scale:             scale of an image w.r.t. given image size,</span>
<span class="sd">                                        e.g. incoming images are (320x320), use scale = 2. to preview in (640x640)</span>
<span class="sd">        :param gt_alpha:                a value in [0., 1.] transparency on ground truth boxes,</span>
<span class="sd">                                        0 for invisible, 1 for fully opaque</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">image_np</span> <span class="o">=</span> <span class="n">undo_preprocessing_func</span><span class="p">(</span><span class="n">image_tensor</span><span class="o">.</span><span class="n">detach</span><span class="p">())</span>
        <span class="n">targets</span> <span class="o">=</span> <span class="n">DetectionVisualization</span><span class="o">.</span><span class="n">_scaled_ccwh_to_xyxy</span><span class="p">(</span><span class="n">target_boxes</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> <span class="o">*</span><span class="n">image_np</span><span class="o">.</span><span class="n">shape</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">image_scale</span><span class="p">)</span>

        <span class="n">out_images</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">image_np</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]):</span>
            <span class="n">preds</span> <span class="o">=</span> <span class="n">pred_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">]</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> <span class="k">if</span> <span class="n">pred_boxes</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">np</span><span class="o">.</span><span class="n">empty</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">6</span><span class="p">))</span>
            <span class="n">targets_cur</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="n">targets</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">i</span><span class="p">]</span>

            <span class="n">image_name</span> <span class="o">=</span> <span class="s1">&#39;_&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">str</span><span class="p">(</span><span class="n">batch_name</span><span class="p">),</span> <span class="nb">str</span><span class="p">(</span><span class="n">i</span><span class="p">)])</span>
            <span class="n">res_image</span> <span class="o">=</span> <span class="n">DetectionVisualization</span><span class="o">.</span><span class="n">_visualize_image</span><span class="p">(</span><span class="n">image_np</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="n">preds</span><span class="p">,</span> <span class="n">targets_cur</span><span class="p">,</span> <span class="n">class_names</span><span class="p">,</span> <span class="n">box_thickness</span><span class="p">,</span> <span class="n">gt_alpha</span><span class="p">,</span> <span class="n">image_scale</span><span class="p">,</span>
                                                                <span class="n">checkpoint_dir</span><span class="p">,</span> <span class="n">image_name</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">res_image</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">out_images</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">res_image</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">out_images</span></div></div>


<div class="viewcode-block" id="Anchors"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.Anchors">[docs]</a><span class="k">class</span> <span class="nc">Anchors</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A wrapper function to hold the anchors used by detection models such as Yolo</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">anchors_list</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">List</span><span class="p">],</span> <span class="n">strides</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param anchors_list: of the shape [[w1,h1,w2,h2,w3,h3], [w4,h4,w5,h5,w6,h6] .... where each sublist holds</span>
<span class="sd">            the width and height of the anchors of a specific detection layer.</span>
<span class="sd">            i.e. for a model with 3 detection layers, each containing 5 anchors the format will be a of 3 sublists of 10 numbers each</span>
<span class="sd">            The width and height are in pixels (not relative to image size)</span>
<span class="sd">        :param strides: a list containing the stride of the layers from which the detection heads are fed.</span>
<span class="sd">            i.e. if the firs detection head is connected to the backbone after the input dimensions were reduces by 8, the first number will be 8</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">__anchors_list</span> <span class="o">=</span> <span class="n">anchors_list</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">__strides</span> <span class="o">=</span> <span class="n">strides</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_check_all_lists</span><span class="p">(</span><span class="n">anchors_list</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_check_all_len_equal_and_even</span><span class="p">(</span><span class="n">anchors_list</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_stride</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">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">strides</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">(),</span> <span class="n">requires_grad</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
        <span class="n">anchors</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">anchors_list</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">anchors_list</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_anchors</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">anchors</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">_stride</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">requires_grad</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_anchor_grid</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">anchors</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">anchors_list</span><span class="p">),</span> <span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</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="n">requires_grad</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_check_all_lists</span><span class="p">(</span><span class="n">anchors</span><span class="p">:</span> <span class="nb">list</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">a</span> <span class="ow">in</span> <span class="n">anchors</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="n">ListConfig</span><span class="p">)):</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s1">&#39;All objects of anchors_list must be lists&#39;</span><span class="p">)</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_check_all_len_equal_and_even</span><span class="p">(</span><span class="n">anchors</span><span class="p">:</span> <span class="nb">list</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
        <span class="n">len_of_first</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">anchors</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
        <span class="k">for</span> <span class="n">a</span> <span class="ow">in</span> <span class="n">anchors</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">a</span><span class="p">)</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">a</span><span class="p">)</span> <span class="o">!=</span> <span class="n">len_of_first</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s1">&#39;All objects of anchors_list must be of the same even length&#39;</span><span class="p">)</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">stride</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">:</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_stride</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">anchors</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">:</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_anchors</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">anchor_grid</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">:</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_anchor_grid</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">detection_layers_num</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_anchors</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">num_anchors</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_anchors</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>

    <span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;anchors_list: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">__anchors_list</span><span class="si">}</span><span class="s2"> strides: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">__strides</span><span class="si">}</span><span class="s2">&quot;</span></div>


<div class="viewcode-block" id="xyxy2cxcywh"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.xyxy2cxcywh">[docs]</a><span class="k">def</span> <span class="nf">xyxy2cxcywh</span><span class="p">(</span><span class="n">bboxes</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Transforms bboxes from xyxy format to centerized xy wh format</span>
<span class="sd">    :param bboxes: array, shaped (nboxes, 4)</span>
<span class="sd">    :return: modified bboxes</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">-</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="mf">0.5</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">*</span> <span class="mf">0.5</span>
    <span class="k">return</span> <span class="n">bboxes</span></div>


<div class="viewcode-block" id="cxcywh2xyxy"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.cxcywh2xyxy">[docs]</a><span class="k">def</span> <span class="nf">cxcywh2xyxy</span><span class="p">(</span><span class="n">bboxes</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Transforms bboxes from centerized xy wh format to xyxy format</span>
<span class="sd">    :param bboxes: array, shaped (nboxes, 4)</span>
<span class="sd">    :return: modified bboxes</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">*</span> <span class="mf">0.5</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="mf">0.5</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">+</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
    <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">bboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">bboxes</span></div>


<div class="viewcode-block" id="get_mosaic_coordinate"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.get_mosaic_coordinate">[docs]</a><span class="k">def</span> <span class="nf">get_mosaic_coordinate</span><span class="p">(</span><span class="n">mosaic_index</span><span class="p">,</span> <span class="n">xc</span><span class="p">,</span> <span class="n">yc</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">input_h</span><span class="p">,</span> <span class="n">input_w</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns the mosaic coordinates of final mosaic image according to mosaic image index.</span>

<span class="sd">    :param mosaic_index: (int) mosaic image index</span>
<span class="sd">    :param xc: (int) center x coordinate of the entire mosaic grid.</span>
<span class="sd">    :param yc: (int) center y coordinate of the entire mosaic grid.</span>
<span class="sd">    :param w: (int) width of bbox</span>
<span class="sd">    :param h: (int) height of bbox</span>
<span class="sd">    :param input_h: (int) image input height (should be 1/2 of the final mosaic output image height).</span>
<span class="sd">    :param input_w: (int) image input width (should be 1/2 of the final mosaic output image width).</span>
<span class="sd">    :return: (x1, y1, x2, y2), (x1s, y1s, x2s, y2s) where (x1, y1, x2, y2) are the coordinates in the final mosaic</span>
<span class="sd">        output image, and (x1s, y1s, x2s, y2s) are the coordinates in the placed image.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># index0 to top left part of image</span>
    <span class="k">if</span> <span class="n">mosaic_index</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">y2</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">xc</span> <span class="o">-</span> <span class="n">w</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="nb">max</span><span class="p">(</span><span class="n">yc</span> <span class="o">-</span> <span class="n">h</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">xc</span><span class="p">,</span> <span class="n">yc</span>
        <span class="n">small_coord</span> <span class="o">=</span> <span class="n">w</span> <span class="o">-</span> <span class="p">(</span><span class="n">x2</span> <span class="o">-</span> <span class="n">x1</span><span class="p">),</span> <span class="n">h</span> <span class="o">-</span> <span class="p">(</span><span class="n">y2</span> <span class="o">-</span> <span class="n">y1</span><span class="p">),</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span>
    <span class="c1"># index1 to top right part of image</span>
    <span class="k">elif</span> <span class="n">mosaic_index</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">y2</span> <span class="o">=</span> <span class="n">xc</span><span class="p">,</span> <span class="nb">max</span><span class="p">(</span><span class="n">yc</span> <span class="o">-</span> <span class="n">h</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="nb">min</span><span class="p">(</span><span class="n">xc</span> <span class="o">+</span> <span class="n">w</span><span class="p">,</span> <span class="n">input_w</span> <span class="o">*</span> <span class="mi">2</span><span class="p">),</span> <span class="n">yc</span>
        <span class="n">small_coord</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="n">h</span> <span class="o">-</span> <span class="p">(</span><span class="n">y2</span> <span class="o">-</span> <span class="n">y1</span><span class="p">),</span> <span class="nb">min</span><span class="p">(</span><span class="n">w</span><span class="p">,</span> <span class="n">x2</span> <span class="o">-</span> <span class="n">x1</span><span class="p">),</span> <span class="n">h</span>
    <span class="c1"># index2 to bottom left part of image</span>
    <span class="k">elif</span> <span class="n">mosaic_index</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">y2</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">xc</span> <span class="o">-</span> <span class="n">w</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">yc</span><span class="p">,</span> <span class="n">xc</span><span class="p">,</span> <span class="nb">min</span><span class="p">(</span><span class="n">input_h</span> <span class="o">*</span> <span class="mi">2</span><span class="p">,</span> <span class="n">yc</span> <span class="o">+</span> <span class="n">h</span><span class="p">)</span>
        <span class="n">small_coord</span> <span class="o">=</span> <span class="n">w</span> <span class="o">-</span> <span class="p">(</span><span class="n">x2</span> <span class="o">-</span> <span class="n">x1</span><span class="p">),</span> <span class="mi">0</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="nb">min</span><span class="p">(</span><span class="n">y2</span> <span class="o">-</span> <span class="n">y1</span><span class="p">,</span> <span class="n">h</span><span class="p">)</span>
    <span class="c1"># index2 to bottom right part of image</span>
    <span class="k">elif</span> <span class="n">mosaic_index</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
        <span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">y2</span> <span class="o">=</span> <span class="n">xc</span><span class="p">,</span> <span class="n">yc</span><span class="p">,</span> <span class="nb">min</span><span class="p">(</span><span class="n">xc</span> <span class="o">+</span> <span class="n">w</span><span class="p">,</span> <span class="n">input_w</span> <span class="o">*</span> <span class="mi">2</span><span class="p">),</span> <span class="nb">min</span><span class="p">(</span><span class="n">input_h</span> <span class="o">*</span> <span class="mi">2</span><span class="p">,</span> <span class="n">yc</span> <span class="o">+</span> <span class="n">h</span><span class="p">)</span>  <span class="c1"># noqa</span>
        <span class="n">small_coord</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="nb">min</span><span class="p">(</span><span class="n">w</span><span class="p">,</span> <span class="n">x2</span> <span class="o">-</span> <span class="n">x1</span><span class="p">),</span> <span class="nb">min</span><span class="p">(</span><span class="n">y2</span> <span class="o">-</span> <span class="n">y1</span><span class="p">,</span> <span class="n">h</span><span class="p">)</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">y1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">y2</span><span class="p">),</span> <span class="n">small_coord</span></div>


<div class="viewcode-block" id="adjust_box_anns"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.adjust_box_anns">[docs]</a><span class="k">def</span> <span class="nf">adjust_box_anns</span><span class="p">(</span><span class="n">bbox</span><span class="p">,</span> <span class="n">scale_ratio</span><span class="p">,</span> <span class="n">padw</span><span class="p">,</span> <span class="n">padh</span><span class="p">,</span> <span class="n">w_max</span><span class="p">,</span> <span class="n">h_max</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Adjusts the bbox annotations of rescaled, padded image.</span>

<span class="sd">    :param bbox: (np.array) bbox to modify.</span>
<span class="sd">    :param scale_ratio: (float) scale ratio between rescale output image and original one.</span>
<span class="sd">    :param padw: (int) width padding size.</span>
<span class="sd">    :param padh: (int) height padding size.</span>
<span class="sd">    :param w_max: (int) width border.</span>
<span class="sd">    :param h_max: (int) height border</span>
<span class="sd">    :return: modified bbox (np.array)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">bbox</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">bbox</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">::</span><span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="n">scale_ratio</span> <span class="o">+</span> <span class="n">padw</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">w_max</span><span class="p">)</span>
    <span class="n">bbox</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="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">clip</span><span class="p">(</span><span class="n">bbox</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="o">*</span> <span class="n">scale_ratio</span> <span class="o">+</span> <span class="n">padh</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">h_max</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">bbox</span></div>


<div class="viewcode-block" id="DetectionCollateFN"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.DetectionCollateFN">[docs]</a><span class="k">class</span> <span class="nc">DetectionCollateFN</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Collate function for Yolox training</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
        <span class="n">batch</span> <span class="o">=</span> <span class="n">default_collate</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
        <span class="n">ims</span><span class="p">,</span> <span class="n">targets</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span><span class="mi">2</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">ims</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_format_targets</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_format_targets</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
        <span class="n">nlabel</span> <span class="o">=</span> <span class="p">(</span><span class="n">targets</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>  <span class="c1"># number of label per image</span>
        <span class="n">targets_merged</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">targets</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]):</span>
            <span class="n">targets_im</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="p">:</span><span class="n">nlabel</span><span class="p">[</span><span class="n">i</span><span class="p">]]</span>
            <span class="n">batch_column</span> <span class="o">=</span> <span class="n">targets</span><span class="o">.</span><span class="n">new_ones</span><span class="p">((</span><span class="n">targets_im</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="mi">1</span><span class="p">))</span> <span class="o">*</span> <span class="n">i</span>
            <span class="n">targets_merged</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">batch_column</span><span class="p">,</span> <span class="n">targets_im</span><span class="p">),</span> <span class="mi">1</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">targets_merged</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span></div>


<div class="viewcode-block" id="CrowdDetectionCollateFN"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.CrowdDetectionCollateFN">[docs]</a><span class="k">class</span> <span class="nc">CrowdDetectionCollateFN</span><span class="p">(</span><span class="n">DetectionCollateFN</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Collate function for Yolox training with additional_batch_items that includes crowd targets</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">data</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]]:</span>
        <span class="n">batch</span> <span class="o">=</span> <span class="n">default_collate</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
        <span class="n">ims</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">crowd_targets</span> <span class="o">=</span> <span class="n">batch</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span><span class="mi">3</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">ims</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_format_targets</span><span class="p">(</span><span class="n">targets</span><span class="p">),</span> <span class="p">{</span><span class="s2">&quot;crowd_targets&quot;</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">_format_targets</span><span class="p">(</span><span class="n">crowd_targets</span><span class="p">)}</span></div>


<div class="viewcode-block" id="compute_box_area"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.compute_box_area">[docs]</a><span class="k">def</span> <span class="nf">compute_box_area</span><span class="p">(</span><span class="n">box</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Compute the area of one or many boxes.</span>
<span class="sd">         :param box: One or many boxes, shape = (4, ?), each box in format (x1, y1, x2, y2)</span>
<span class="sd">    Returns:</span>
<span class="sd">        Area of every box, shape = (1, ?)</span>
<span class="sd">     &quot;&quot;&quot;</span>
    <span class="c1"># box = 4xn</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">box</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">-</span> <span class="n">box</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">*</span> <span class="p">(</span><span class="n">box</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">-</span> <span class="n">box</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span></div>


<div class="viewcode-block" id="crowd_ioa"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.crowd_ioa">[docs]</a><span class="k">def</span> <span class="nf">crowd_ioa</span><span class="p">(</span><span class="n">det_box</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">crowd_box</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Return intersection-over-detection_area of boxes, used for crowd ground truths.</span>
<span class="sd">    Both sets of boxes are expected to be in (x1, y1, x2, y2) format.</span>
<span class="sd">    Arguments:</span>
<span class="sd">        det_box (Tensor[N, 4])</span>
<span class="sd">        crowd_box (Tensor[M, 4])</span>
<span class="sd">    Returns:</span>
<span class="sd">        crowd_ioa (Tensor[N, M]): the NxM matrix containing the pairwise</span>
<span class="sd">            IoA values for every element in det_box and crowd_box</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">det_area</span> <span class="o">=</span> <span class="n">compute_box_area</span><span class="p">(</span><span class="n">det_box</span><span class="o">.</span><span class="n">T</span><span class="p">)</span>

    <span class="c1"># inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)</span>
    <span class="n">inter</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">det_box</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">,</span> <span class="mi">2</span><span class="p">:],</span> <span class="n">crowd_box</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">:])</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">det_box</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">,</span> <span class="p">:</span><span class="mi">2</span><span class="p">],</span> <span class="n">crowd_box</span><span class="p">[:,</span> <span class="p">:</span><span class="mi">2</span><span class="p">]))</span> \
        <span class="o">.</span><span class="n">clamp</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">prod</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">inter</span> <span class="o">/</span> <span class="n">det_area</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span>  <span class="c1"># crowd_ioa = inter / det_area</span></div>


<div class="viewcode-block" id="compute_detection_matching"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.compute_detection_matching">[docs]</a><span class="k">def</span> <span class="nf">compute_detection_matching</span><span class="p">(</span>
    <span class="n">output</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">height</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
    <span class="n">width</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
    <span class="n">iou_thresholds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">denormalize_targets</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
    <span class="n">device</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
    <span class="n">crowd_targets</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">top_k</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span>
    <span class="n">return_on_cpu</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="n">Tuple</span><span class="p">]:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Match predictions (NMS output) and the targets (ground truth) with respect to IoU and confidence score.</span>
<span class="sd">    :param output:          list (of length batch_size) of Tensors of shape (num_predictions, 6)</span>
<span class="sd">                            format:     (x1, y1, x2, y2, confidence, class_label) where x1,y1,x2,y2 are according to image size</span>
<span class="sd">    :param targets:         targets for all images of shape (total_num_targets, 6)</span>
<span class="sd">                            format:     (index, x, y, w, h, label) where x,y,w,h are in range [0,1]</span>
<span class="sd">    :param height:          dimensions of the image</span>
<span class="sd">    :param width:           dimensions of the image</span>
<span class="sd">    :param iou_thresholds:  Threshold to compute the mAP</span>
<span class="sd">    :param device:          Device</span>
<span class="sd">    :param crowd_targets:   crowd targets for all images of shape (total_num_crowd_targets, 6)</span>
<span class="sd">                            format:     (index, x, y, w, h, label) where x,y,w,h are in range [0,1]</span>
<span class="sd">    :param top_k:           Number of predictions to keep per class, ordered by confidence score</span>
<span class="sd">    :param denormalize_targets: If True, denormalize the targets and crowd_targets</span>
<span class="sd">    :param return_on_cpu:   If True, the output will be returned on &quot;CPU&quot;, otherwise it will be returned on &quot;device&quot;</span>

<span class="sd">    :return:                list of the following tensors, for every image:</span>
<span class="sd">        :preds_matched:     Tensor of shape (num_img_predictions, n_iou_thresholds)</span>
<span class="sd">                            True when prediction (i) is matched with a target with respect to the (j)th IoU threshold</span>
<span class="sd">        :preds_to_ignore:   Tensor of shape (num_img_predictions, n_iou_thresholds)</span>
<span class="sd">                            True when prediction (i) is matched with a crowd target with respect to the (j)th IoU threshold</span>
<span class="sd">        :preds_scores:      Tensor of shape (num_img_predictions), confidence score for every prediction</span>
<span class="sd">        :preds_cls:         Tensor of shape (num_img_predictions), predicted class for every prediction</span>
<span class="sd">        :targets_cls:       Tensor of shape (num_img_targets), ground truth class for every target</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">output</span> <span class="o">=</span> <span class="nb">map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">tensor</span><span class="p">:</span> <span class="kc">None</span> <span class="k">if</span> <span class="n">tensor</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">tensor</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">),</span> <span class="n">output</span><span class="p">)</span>
    <span class="n">targets</span><span class="p">,</span> <span class="n">iou_thresholds</span> <span class="o">=</span> <span class="n">targets</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">),</span> <span class="n">iou_thresholds</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="c1"># If crowd_targets is not provided, we patch it with an empty tensor</span>
    <span class="n">crowd_targets</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">6</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="k">if</span> <span class="n">crowd_targets</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">crowd_targets</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="n">batch_metrics</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">img_i</span><span class="p">,</span> <span class="n">img_preds</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">output</span><span class="p">):</span>
        <span class="c1"># If img_preds is None (not prediction for this image), we patch it with an empty tensor</span>
        <span class="n">img_preds</span> <span class="o">=</span> <span class="n">img_preds</span> <span class="k">if</span> <span class="n">img_preds</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">6</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">img_targets</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[</span><span class="n">targets</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">img_i</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span>
        <span class="n">img_crowd_targets</span> <span class="o">=</span> <span class="n">crowd_targets</span><span class="p">[</span><span class="n">crowd_targets</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">img_i</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span>

        <span class="n">img_matching_tensors</span> <span class="o">=</span> <span class="n">compute_img_detection_matching</span><span class="p">(</span>
            <span class="n">preds</span><span class="o">=</span><span class="n">img_preds</span><span class="p">,</span>
            <span class="n">targets</span><span class="o">=</span><span class="n">img_targets</span><span class="p">,</span>
            <span class="n">crowd_targets</span><span class="o">=</span><span class="n">img_crowd_targets</span><span class="p">,</span>
            <span class="n">denormalize_targets</span><span class="o">=</span><span class="n">denormalize_targets</span><span class="p">,</span>
            <span class="n">height</span><span class="o">=</span><span class="n">height</span><span class="p">,</span>
            <span class="n">width</span><span class="o">=</span><span class="n">width</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">iou_thresholds</span><span class="o">=</span><span class="n">iou_thresholds</span><span class="p">,</span>
            <span class="n">top_k</span><span class="o">=</span><span class="n">top_k</span><span class="p">,</span>
            <span class="n">return_on_cpu</span><span class="o">=</span><span class="n">return_on_cpu</span>
        <span class="p">)</span>
        <span class="n">batch_metrics</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">img_matching_tensors</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">batch_metrics</span></div>


<div class="viewcode-block" id="compute_img_detection_matching"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.compute_img_detection_matching">[docs]</a><span class="k">def</span> <span class="nf">compute_img_detection_matching</span><span class="p">(</span>
        <span class="n">preds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">crowd_targets</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">height</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
        <span class="n">width</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
        <span class="n">iou_thresholds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">device</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
        <span class="n">denormalize_targets</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
        <span class="n">top_k</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span>
        <span class="n">return_on_cpu</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Match predictions (NMS output) and the targets (ground truth) with respect to IoU and confidence score</span>
<span class="sd">    for a given image.</span>
<span class="sd">    :param preds:           Tensor of shape (num_img_predictions, 6)</span>
<span class="sd">                            format:     (x1, y1, x2, y2, confidence, class_label) where x1,y1,x2,y2 are according to image size</span>
<span class="sd">    :param targets:         targets for this image of shape (num_img_targets, 6)</span>
<span class="sd">                            format:     (index, x, y, w, h, label) where x,y,w,h are in range [0,1]</span>
<span class="sd">    :param height:          dimensions of the image</span>
<span class="sd">    :param width:           dimensions of the image</span>
<span class="sd">    :param iou_thresholds:  Threshold to compute the mAP</span>
<span class="sd">    :param device:</span>
<span class="sd">    :param crowd_targets:   crowd targets for all images of shape (total_num_crowd_targets, 6)</span>
<span class="sd">                            format:     (index, x, y, w, h, label) where x,y,w,h are in range [0,1]</span>
<span class="sd">    :param top_k:           Number of predictions to keep per class, ordered by confidence score</span>
<span class="sd">    :param device:          Device</span>
<span class="sd">    :param denormalize_targets: If True, denormalize the targets and crowd_targets</span>
<span class="sd">    :param return_on_cpu:   If True, the output will be returned on &quot;CPU&quot;, otherwise it will be returned on &quot;device&quot;</span>

<span class="sd">    :return:</span>
<span class="sd">        :preds_matched:     Tensor of shape (num_img_predictions, n_iou_thresholds)</span>
<span class="sd">                                True when prediction (i) is matched with a target with respect to the (j)th IoU threshold</span>
<span class="sd">        :preds_to_ignore:   Tensor of shape (num_img_predictions, n_iou_thresholds)</span>
<span class="sd">                                True when prediction (i) is matched with a crowd target with respect to the (j)th IoU threshold</span>
<span class="sd">        :preds_scores:      Tensor of shape (num_img_predictions), confidence score for every prediction</span>
<span class="sd">        :preds_cls:         Tensor of shape (num_img_predictions), predicted class for every prediction</span>
<span class="sd">        :targets_cls:       Tensor of shape (num_img_targets), ground truth class for every target</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">num_iou_thresholds</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">iou_thresholds</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">preds</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">preds</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">return_on_cpu</span><span class="p">:</span>
            <span class="n">device</span> <span class="o">=</span> <span class="s2">&quot;cpu&quot;</span>
        <span class="n">preds_matched</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">num_iou_thresholds</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">bool</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">preds_to_ignore</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">num_iou_thresholds</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">bool</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">preds_scores</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">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</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">preds_cls</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">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</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">targets_cls</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[:,</span> <span class="mi">0</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="o">=</span><span class="n">device</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">preds_matched</span><span class="p">,</span> <span class="n">preds_to_ignore</span><span class="p">,</span> <span class="n">preds_scores</span><span class="p">,</span> <span class="n">preds_cls</span><span class="p">,</span> <span class="n">targets_cls</span>

    <span class="n">preds_matched</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">preds</span><span class="p">),</span> <span class="n">num_iou_thresholds</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">bool</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">targets_matched</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">targets</span><span class="p">),</span> <span class="n">num_iou_thresholds</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">bool</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">preds_to_ignore</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">preds</span><span class="p">),</span> <span class="n">num_iou_thresholds</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">bool</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">preds_cls</span><span class="p">,</span> <span class="n">preds_box</span><span class="p">,</span> <span class="n">preds_scores</span> <span class="o">=</span> <span class="n">preds</span><span class="p">[:,</span> <span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">preds</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">:</span><span class="mi">4</span><span class="p">],</span> <span class="n">preds</span><span class="p">[:,</span> <span class="mi">4</span><span class="p">]</span>
    <span class="n">targets_cls</span><span class="p">,</span> <span class="n">targets_box</span> <span class="o">=</span> <span class="n">targets</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">targets</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">:</span><span class="mi">5</span><span class="p">]</span>
    <span class="n">crowd_targets_cls</span><span class="p">,</span> <span class="n">crowd_target_box</span> <span class="o">=</span> <span class="n">crowd_targets</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">crowd_targets</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">:</span><span class="mi">5</span><span class="p">]</span>

    <span class="c1"># Ignore all but the predictions that were top_k for their class</span>
    <span class="n">preds_idx_to_use</span> <span class="o">=</span> <span class="n">get_top_k_idx_per_cls</span><span class="p">(</span><span class="n">preds_scores</span><span class="p">,</span> <span class="n">preds_cls</span><span class="p">,</span> <span class="n">top_k</span><span class="p">)</span>
    <span class="n">preds_to_ignore</span><span class="p">[:,</span> <span class="p">:]</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="n">preds_to_ignore</span><span class="p">[</span><span class="n">preds_idx_to_use</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">crowd_targets</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>

        <span class="c1"># CHANGE bboxes TO FIT THE IMAGE SIZE</span>
        <span class="n">change_bbox_bounds_for_image_size</span><span class="p">(</span><span class="n">preds</span><span class="p">,</span> <span class="p">(</span><span class="n">height</span><span class="p">,</span> <span class="n">width</span><span class="p">))</span>

        <span class="c1"># if target_format == &quot;xywh&quot;:</span>
        <span class="n">targets_box</span> <span class="o">=</span> <span class="n">convert_xywh_bbox_to_xyxy</span><span class="p">(</span><span class="n">targets_box</span><span class="p">)</span>  <span class="c1"># cxcywh2xyxy</span>
        <span class="n">crowd_target_box</span> <span class="o">=</span> <span class="n">convert_xywh_bbox_to_xyxy</span><span class="p">(</span><span class="n">crowd_target_box</span><span class="p">)</span>  <span class="c1"># convert_xywh_bbox_to_xyxy</span>

        <span class="k">if</span> <span class="n">denormalize_targets</span><span class="p">:</span>
            <span class="n">targets_box</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span> <span class="o">*=</span> <span class="n">width</span>
            <span class="n">targets_box</span><span class="p">[:,</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="o">*=</span> <span class="n">height</span>
            <span class="n">crowd_target_box</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span> <span class="o">*=</span> <span class="n">width</span>
            <span class="n">crowd_target_box</span><span class="p">[:,</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="o">*=</span> <span class="n">height</span>

    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">targets</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>

        <span class="c1"># shape = (n_preds x n_targets)</span>
        <span class="n">iou</span> <span class="o">=</span> <span class="n">box_iou</span><span class="p">(</span><span class="n">preds_box</span><span class="p">[</span><span class="n">preds_idx_to_use</span><span class="p">],</span> <span class="n">targets_box</span><span class="p">)</span>

        <span class="c1"># Fill IoU values at index (i, j) with 0 when the prediction (i) and target(j) are of different class</span>
        <span class="c1"># Filling with 0 is equivalent to ignore these values since with want IoU &gt; iou_threshold &gt; 0</span>
        <span class="n">cls_mismatch</span> <span class="o">=</span> <span class="p">(</span><span class="n">preds_cls</span><span class="p">[</span><span class="n">preds_idx_to_use</span><span class="p">]</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="o">!=</span> <span class="n">targets_cls</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>
        <span class="n">iou</span><span class="p">[</span><span class="n">cls_mismatch</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>

        <span class="c1"># The matching priority is first detection confidence and then IoU value.</span>
        <span class="c1"># The detection is already sorted by confidence in NMS, so here for each prediction we order the targets by iou.</span>
        <span class="n">sorted_iou</span><span class="p">,</span> <span class="n">target_sorted</span> <span class="o">=</span> <span class="n">iou</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="n">descending</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">stable</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

        <span class="c1"># Only iterate over IoU values higher than min threshold to speed up the process</span>
        <span class="k">for</span> <span class="n">pred_selected_i</span><span class="p">,</span> <span class="n">target_sorted_i</span> <span class="ow">in</span> <span class="p">(</span><span class="n">sorted_iou</span> <span class="o">&gt;</span> <span class="n">iou_thresholds</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span><span class="o">.</span><span class="n">nonzero</span><span class="p">(</span><span class="n">as_tuple</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>

            <span class="c1"># pred_selected_i and target_sorted_i are relative to filters/sorting, so we extract their absolute indexes</span>
            <span class="n">pred_i</span> <span class="o">=</span> <span class="n">preds_idx_to_use</span><span class="p">[</span><span class="n">pred_selected_i</span><span class="p">]</span>
            <span class="n">target_i</span> <span class="o">=</span> <span class="n">target_sorted</span><span class="p">[</span><span class="n">pred_selected_i</span><span class="p">,</span> <span class="n">target_sorted_i</span><span class="p">]</span>

            <span class="c1"># Vector[j], True when IoU(pred_i, target_i) is above the (j)th threshold</span>
            <span class="n">is_iou_above_threshold</span> <span class="o">=</span> <span class="n">sorted_iou</span><span class="p">[</span><span class="n">pred_selected_i</span><span class="p">,</span> <span class="n">target_sorted_i</span><span class="p">]</span> <span class="o">&gt;</span> <span class="n">iou_thresholds</span>

            <span class="c1"># Vector[j], True when both pred_i and target_i are not matched yet for the (j)th threshold</span>
            <span class="n">are_candidates_free</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="o">~</span><span class="n">preds_matched</span><span class="p">[</span><span class="n">pred_i</span><span class="p">,</span> <span class="p">:],</span> <span class="o">~</span><span class="n">targets_matched</span><span class="p">[</span><span class="n">target_i</span><span class="p">,</span> <span class="p">:])</span>

            <span class="c1"># Vector[j], True when (pred_i, target_i) can be matched for the (j)th threshold</span>
            <span class="n">are_candidates_good</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="n">is_iou_above_threshold</span><span class="p">,</span> <span class="n">are_candidates_free</span><span class="p">)</span>

            <span class="c1"># For every threshold (j) where target_i and pred_i can be matched together ( are_candidates_good[j]==True )</span>
            <span class="c1"># fill the matching placeholders with True</span>
            <span class="n">targets_matched</span><span class="p">[</span><span class="n">target_i</span><span class="p">,</span> <span class="n">are_candidates_good</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="n">preds_matched</span><span class="p">[</span><span class="n">pred_i</span><span class="p">,</span> <span class="n">are_candidates_good</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>

            <span class="c1"># When all the targets are matched with a prediction for every IoU Threshold, stop.</span>
            <span class="k">if</span> <span class="n">targets_matched</span><span class="o">.</span><span class="n">all</span><span class="p">():</span>
                <span class="k">break</span>

    <span class="c1"># Crowd targets can be matched with many predictions.</span>
    <span class="c1"># Therefore, for every prediction we just need to check if it has IoA large enough with any crowd target.</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">crowd_targets</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>

        <span class="c1"># shape = (n_preds_to_use x n_crowd_targets)</span>
        <span class="n">ioa</span> <span class="o">=</span> <span class="n">crowd_ioa</span><span class="p">(</span><span class="n">preds_box</span><span class="p">[</span><span class="n">preds_idx_to_use</span><span class="p">],</span> <span class="n">crowd_target_box</span><span class="p">)</span>

        <span class="c1"># Fill IoA values at index (i, j) with 0 when the prediction (i) and target(j) are of different class</span>
        <span class="c1"># Filling with 0 is equivalent to ignore these values since with want IoA &gt; threshold &gt; 0</span>
        <span class="n">cls_mismatch</span> <span class="o">=</span> <span class="p">(</span><span class="n">preds_cls</span><span class="p">[</span><span class="n">preds_idx_to_use</span><span class="p">]</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="o">!=</span> <span class="n">crowd_targets_cls</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>
        <span class="n">ioa</span><span class="p">[</span><span class="n">cls_mismatch</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>

        <span class="c1"># For each prediction, we keep it&#39;s highest score with any crowd target (of same class)</span>
        <span class="c1"># shape = (n_preds_to_use)</span>
        <span class="n">best_ioa</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">ioa</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

        <span class="c1"># If a prediction has IoA higher than threshold (with any target of same class), then there is a match</span>
        <span class="c1"># shape = (n_preds_to_use x iou_thresholds)</span>
        <span class="n">is_matching_with_crowd</span> <span class="o">=</span> <span class="p">(</span><span class="n">best_ioa</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="o">&gt;</span> <span class="n">iou_thresholds</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>

        <span class="n">preds_to_ignore</span><span class="p">[</span><span class="n">preds_idx_to_use</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">logical_or</span><span class="p">(</span><span class="n">preds_to_ignore</span><span class="p">[</span><span class="n">preds_idx_to_use</span><span class="p">],</span> <span class="n">is_matching_with_crowd</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">return_on_cpu</span><span class="p">:</span>
        <span class="n">preds_matched</span> <span class="o">=</span> <span class="n">preds_matched</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="s2">&quot;cpu&quot;</span><span class="p">)</span>
        <span class="n">preds_to_ignore</span> <span class="o">=</span> <span class="n">preds_to_ignore</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="s2">&quot;cpu&quot;</span><span class="p">)</span>
        <span class="n">preds_scores</span> <span class="o">=</span> <span class="n">preds_scores</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="s2">&quot;cpu&quot;</span><span class="p">)</span>
        <span class="n">preds_cls</span> <span class="o">=</span> <span class="n">preds_cls</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="s2">&quot;cpu&quot;</span><span class="p">)</span>
        <span class="n">targets_cls</span> <span class="o">=</span> <span class="n">targets_cls</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="s2">&quot;cpu&quot;</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">preds_matched</span><span class="p">,</span> <span class="n">preds_to_ignore</span><span class="p">,</span> <span class="n">preds_scores</span><span class="p">,</span> <span class="n">preds_cls</span><span class="p">,</span> <span class="n">targets_cls</span></div>


<div class="viewcode-block" id="get_top_k_idx_per_cls"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.get_top_k_idx_per_cls">[docs]</a><span class="k">def</span> <span class="nf">get_top_k_idx_per_cls</span><span class="p">(</span><span class="n">preds_scores</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">preds_cls</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">top_k</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Get the indexes of all the top k predictions for every class</span>

<span class="sd">    :param preds_scores:   The confidence scores, vector of shape (n_pred)</span>
<span class="sd">    :param preds_cls:      The predicted class, vector of shape (n_pred)</span>
<span class="sd">    :param top_k:          Number of predictions to keep per class, ordered by confidence score</span>

<span class="sd">    :return top_k_idx:     Indexes of the top k predictions. length &lt;= (k * n_unique_class)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">n_unique_cls</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">preds_cls</span><span class="p">)</span>
    <span class="n">mask</span> <span class="o">=</span> <span class="p">(</span><span class="n">preds_cls</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">n_unique_cls</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">preds_scores</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>
    <span class="n">preds_scores_per_cls</span> <span class="o">=</span> <span class="n">preds_scores</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="o">*</span> <span class="n">mask</span>

    <span class="n">sorted_scores_per_cls</span><span class="p">,</span> <span class="n">sorting_idx</span> <span class="o">=</span> <span class="n">preds_scores_per_cls</span><span class="o">.</span><span class="n">sort</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">descending</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="n">idx_with_satisfying_scores</span> <span class="o">=</span> <span class="n">sorted_scores_per_cls</span><span class="p">[:</span><span class="n">top_k</span><span class="p">,</span> <span class="p">:]</span><span class="o">.</span><span class="n">nonzero</span><span class="p">(</span><span class="n">as_tuple</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
    <span class="n">top_k_idx</span> <span class="o">=</span> <span class="n">sorting_idx</span><span class="p">[</span><span class="n">idx_with_satisfying_scores</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)]</span>
    <span class="k">return</span> <span class="n">top_k_idx</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span></div>


<div class="viewcode-block" id="compute_detection_metrics"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.compute_detection_metrics">[docs]</a><span class="k">def</span> <span class="nf">compute_detection_metrics</span><span class="p">(</span>
    <span class="n">preds_matched</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">preds_to_ignore</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">preds_scores</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">preds_cls</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">targets_cls</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">device</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
    <span class="n">recall_thresholds</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">score_threshold</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Compute the list of precision, recall, MaP and f1 for every recall IoU threshold and for every class.</span>

<span class="sd">    :param preds_matched:      Tensor of shape (num_predictions, n_iou_thresholds)</span>
<span class="sd">                                    True when prediction (i) is matched with a target with respect to the (j)th IoU threshold</span>
<span class="sd">    :param preds_to_ignore     Tensor of shape (num_predictions, n_iou_thresholds)</span>
<span class="sd">                                    True when prediction (i) is matched with a crowd target with respect to the (j)th IoU threshold</span>
<span class="sd">    :param preds_scores:       Tensor of shape (num_predictions), confidence score for every prediction</span>
<span class="sd">    :param preds_cls:          Tensor of shape (num_predictions), predicted class for every prediction</span>
<span class="sd">    :param targets_cls:        Tensor of shape (num_targets), ground truth class for every target box to be detected</span>
<span class="sd">    :param recall_thresholds:   Recall thresholds used to compute MaP.</span>
<span class="sd">    :param score_threshold:    Minimum confidence score to consider a prediction for the computation of</span>
<span class="sd">                                    precision, recall and f1 (not MaP)</span>
<span class="sd">    :param device:             Device</span>

<span class="sd">    :return:</span>
<span class="sd">        :ap, precision, recall, f1: Tensors of shape (n_class, nb_iou_thrs)</span>
<span class="sd">        :unique_classes:            Vector with all unique target classes</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">preds_matched</span><span class="p">,</span> <span class="n">preds_to_ignore</span> <span class="o">=</span> <span class="n">preds_matched</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">),</span> <span class="n">preds_to_ignore</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
    <span class="n">preds_scores</span><span class="p">,</span> <span class="n">preds_cls</span><span class="p">,</span> <span class="n">targets_cls</span> <span class="o">=</span> <span class="n">preds_scores</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">),</span> <span class="n">preds_cls</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">),</span> <span class="n">targets_cls</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="n">recall_thresholds</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">101</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="k">if</span> <span class="n">recall_thresholds</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">recall_thresholds</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>

    <span class="n">unique_classes</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">targets_cls</span><span class="p">)</span>
    <span class="n">n_class</span><span class="p">,</span> <span class="n">nb_iou_thrs</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">unique_classes</span><span class="p">),</span> <span class="n">preds_matched</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>

    <span class="n">ap</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_class</span><span class="p">,</span> <span class="n">nb_iou_thrs</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">precision</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_class</span><span class="p">,</span> <span class="n">nb_iou_thrs</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">recall</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_class</span><span class="p">,</span> <span class="n">nb_iou_thrs</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="k">for</span> <span class="n">cls_i</span><span class="p">,</span> <span class="bp">cls</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">unique_classes</span><span class="p">):</span>
        <span class="n">cls_preds_idx</span><span class="p">,</span> <span class="n">cls_targets_idx</span> <span class="o">=</span> <span class="p">(</span><span class="n">preds_cls</span> <span class="o">==</span> <span class="bp">cls</span><span class="p">),</span> <span class="p">(</span><span class="n">targets_cls</span> <span class="o">==</span> <span class="bp">cls</span><span class="p">)</span>
        <span class="n">cls_ap</span><span class="p">,</span> <span class="n">cls_precision</span><span class="p">,</span> <span class="n">cls_recall</span> <span class="o">=</span> <span class="n">compute_detection_metrics_per_cls</span><span class="p">(</span>
            <span class="n">preds_matched</span><span class="o">=</span><span class="n">preds_matched</span><span class="p">[</span><span class="n">cls_preds_idx</span><span class="p">],</span>
            <span class="n">preds_to_ignore</span><span class="o">=</span><span class="n">preds_to_ignore</span><span class="p">[</span><span class="n">cls_preds_idx</span><span class="p">],</span>
            <span class="n">preds_scores</span><span class="o">=</span><span class="n">preds_scores</span><span class="p">[</span><span class="n">cls_preds_idx</span><span class="p">],</span>
            <span class="n">n_targets</span><span class="o">=</span><span class="n">cls_targets_idx</span><span class="o">.</span><span class="n">sum</span><span class="p">(),</span>
            <span class="n">recall_thresholds</span><span class="o">=</span><span class="n">recall_thresholds</span><span class="p">,</span>
            <span class="n">score_threshold</span><span class="o">=</span><span class="n">score_threshold</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">ap</span><span class="p">[</span><span class="n">cls_i</span><span class="p">,</span> <span class="p">:]</span> <span class="o">=</span> <span class="n">cls_ap</span>
        <span class="n">precision</span><span class="p">[</span><span class="n">cls_i</span><span class="p">,</span> <span class="p">:]</span> <span class="o">=</span> <span class="n">cls_precision</span>
        <span class="n">recall</span><span class="p">[</span><span class="n">cls_i</span><span class="p">,</span> <span class="p">:]</span> <span class="o">=</span> <span class="n">cls_recall</span>

    <span class="n">f1</span> <span class="o">=</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">precision</span> <span class="o">*</span> <span class="n">recall</span> <span class="o">/</span> <span class="p">(</span><span class="n">precision</span> <span class="o">+</span> <span class="n">recall</span> <span class="o">+</span> <span class="mf">1e-16</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">ap</span><span class="p">,</span> <span class="n">precision</span><span class="p">,</span> <span class="n">recall</span><span class="p">,</span> <span class="n">f1</span><span class="p">,</span> <span class="n">unique_classes</span></div>


<div class="viewcode-block" id="compute_detection_metrics_per_cls"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.detection_utils.compute_detection_metrics_per_cls">[docs]</a><span class="k">def</span> <span class="nf">compute_detection_metrics_per_cls</span><span class="p">(</span>
        <span class="n">preds_matched</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">preds_to_ignore</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">preds_scores</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">n_targets</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
        <span class="n">recall_thresholds</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
        <span class="n">score_threshold</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
        <span class="n">device</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
<span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Compute the list of precision, recall and MaP of a given class for every recall IoU threshold.</span>

<span class="sd">        :param preds_matched:      Tensor of shape (num_predictions, n_iou_thresholds)</span>
<span class="sd">                                        True when prediction (i) is matched with a target</span>
<span class="sd">                                        with respect to the(j)th IoU threshold</span>
<span class="sd">        :param preds_to_ignore     Tensor of shape (num_predictions, n_iou_thresholds)</span>
<span class="sd">                                        True when prediction (i) is matched with a crowd target</span>
<span class="sd">                                        with respect to the (j)th IoU threshold</span>
<span class="sd">        :param preds_scores:       Tensor of shape (num_predictions), confidence score for every prediction</span>
<span class="sd">        :param n_targets:          Number of target boxes of this class</span>
<span class="sd">        :param recall_thresholds:  Tensor of shape (max_n_rec_thresh) list of recall thresholds used to compute MaP</span>
<span class="sd">        :param score_threshold:    Minimum confidence score to consider a prediction for the computation of</span>
<span class="sd">                                        precision and recall (not MaP)</span>
<span class="sd">        :param device:             Device</span>

<span class="sd">        :return ap, precision, recall:  Tensors of shape (nb_iou_thrs)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">nb_iou_thrs</span> <span class="o">=</span> <span class="n">preds_matched</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>

    <span class="n">tps</span> <span class="o">=</span> <span class="n">preds_matched</span>
    <span class="n">fps</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">logical_not</span><span class="p">(</span><span class="n">preds_matched</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">logical_not</span><span class="p">(</span><span class="n">preds_to_ignore</span><span class="p">))</span>

    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tps</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">return</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">nb_iou_thrs</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="c1"># Sort by decreasing score</span>
    <span class="n">dtype</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">uint8</span> <span class="k">if</span> <span class="n">preds_scores</span><span class="o">.</span><span class="n">is_cuda</span> <span class="ow">and</span> <span class="n">preds_scores</span><span class="o">.</span><span class="n">dtype</span> <span class="ow">is</span> <span class="n">torch</span><span class="o">.</span><span class="n">bool</span> <span class="k">else</span> <span class="n">preds_scores</span><span class="o">.</span><span class="n">dtype</span>
    <span class="n">sort_ind</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">argsort</span><span class="p">(</span><span class="n">preds_scores</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">dtype</span><span class="p">),</span> <span class="n">descending</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="n">tps</span> <span class="o">=</span> <span class="n">tps</span><span class="p">[</span><span class="n">sort_ind</span><span class="p">,</span> <span class="p">:]</span>
    <span class="n">fps</span> <span class="o">=</span> <span class="n">fps</span><span class="p">[</span><span class="n">sort_ind</span><span class="p">,</span> <span class="p">:]</span>
    <span class="n">preds_scores</span> <span class="o">=</span> <span class="n">preds_scores</span><span class="p">[</span><span class="n">sort_ind</span><span class="p">]</span>

    <span class="c1"># Rolling sum over the predictions</span>
    <span class="n">rolling_tps</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cumsum</span><span class="p">(</span><span class="n">tps</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</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">float</span><span class="p">)</span>
    <span class="n">rolling_fps</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cumsum</span><span class="p">(</span><span class="n">fps</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</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">float</span><span class="p">)</span>

    <span class="n">rolling_recalls</span> <span class="o">=</span> <span class="n">rolling_tps</span> <span class="o">/</span> <span class="n">n_targets</span>
    <span class="n">rolling_precisions</span> <span class="o">=</span> <span class="n">rolling_tps</span> <span class="o">/</span> <span class="p">(</span><span class="n">rolling_tps</span> <span class="o">+</span> <span class="n">rolling_fps</span> <span class="o">+</span> <span class="n">torch</span><span class="o">.</span><span class="n">finfo</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">float64</span><span class="p">)</span><span class="o">.</span><span class="n">eps</span><span class="p">)</span>

    <span class="c1"># Reversed cummax to only have decreasing values</span>
    <span class="n">rolling_precisions</span> <span class="o">=</span> <span class="n">rolling_precisions</span><span class="o">.</span><span class="n">flip</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">cummax</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">values</span><span class="o">.</span><span class="n">flip</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="c1"># ==================</span>
    <span class="c1"># RECALL &amp; PRECISION</span>

    <span class="c1"># We want the rolling precision/recall at index i so that: preds_scores[i-1] &gt;= score_threshold &gt; preds_scores[i]</span>
    <span class="c1"># Note: torch.searchsorted works on increasing sequence and preds_scores is decreasing, so we work with &quot;-&quot;</span>
    <span class="n">lowest_score_above_threshold</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">searchsorted</span><span class="p">(</span><span class="o">-</span><span class="n">preds_scores</span><span class="p">,</span> <span class="o">-</span><span class="n">score_threshold</span><span class="p">,</span> <span class="n">right</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">lowest_score_above_threshold</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>  <span class="c1"># Here score_threshold &gt; preds_scores[0], so no pred is above the threshold</span>
        <span class="n">recall</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="n">precision</span> <span class="o">=</span> <span class="mi">0</span>  <span class="c1"># the precision is not really defined when no pred but we need to give it a value</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">recall</span> <span class="o">=</span> <span class="n">rolling_recalls</span><span class="p">[</span><span class="n">lowest_score_above_threshold</span> <span class="o">-</span> <span class="mi">1</span><span class="p">]</span>
        <span class="n">precision</span> <span class="o">=</span> <span class="n">rolling_precisions</span><span class="p">[</span><span class="n">lowest_score_above_threshold</span> <span class="o">-</span> <span class="mi">1</span><span class="p">]</span>

    <span class="c1"># ==================</span>
    <span class="c1"># AVERAGE PRECISION</span>

    <span class="c1"># shape = (nb_iou_thrs, n_recall_thresholds)</span>
    <span class="n">recall_thresholds</span> <span class="o">=</span> <span class="n">recall_thresholds</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="n">nb_iou_thrs</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>

    <span class="c1"># We want the index i so that: rolling_recalls[i-1] &lt; recall_thresholds[k] &lt;= rolling_recalls[i]</span>
    <span class="c1"># Note:  when recall_thresholds[k] &gt; max(rolling_recalls), i = len(rolling_recalls)</span>
    <span class="c1"># Note2: we work with transpose (.T) to apply torch.searchsorted on first dim instead of the last one</span>
    <span class="n">recall_threshold_idx</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">searchsorted</span><span class="p">(</span><span class="n">rolling_recalls</span><span class="o">.</span><span class="n">T</span><span class="p">,</span> <span class="n">recall_thresholds</span><span class="p">,</span> <span class="n">right</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span><span class="o">.</span><span class="n">T</span>

    <span class="c1"># When recall_thresholds[k] &gt; max(rolling_recalls), rolling_precisions[i] is not defined, and we want precision = 0</span>
    <span class="n">rolling_precisions</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">((</span><span class="n">rolling_precisions</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">nb_iou_thrs</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">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>

    <span class="c1"># shape = (n_recall_thresholds, nb_iou_thrs)</span>
    <span class="n">sampled_precision_points</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">gather</span><span class="p">(</span><span class="nb">input</span><span class="o">=</span><span class="n">rolling_precisions</span><span class="p">,</span> <span class="n">index</span><span class="o">=</span><span class="n">recall_threshold_idx</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>

    <span class="c1"># Average over the recall_thresholds</span>
    <span class="n">ap</span> <span class="o">=</span> <span class="n">sampled_precision_points</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">ap</span><span class="p">,</span> <span class="n">precision</span><span class="p">,</span> <span class="n">recall</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.distributed_training_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.distributed_training_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.distributed_training_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">distributed</span> <span class="k">as</span> <span class="n">dist</span>
<span class="kn">from</span> <span class="nn">torch.cuda.amp</span> <span class="kn">import</span> <span class="n">autocast</span>
<span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
<span class="kn">import</span> <span class="nn">itertools</span>
<span class="kn">from</span> <span class="nn">contextlib</span> <span class="kn">import</span> <span class="n">contextmanager</span>


<div class="viewcode-block" id="distributed_all_reduce_tensor_average"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.distributed_all_reduce_tensor_average">[docs]</a><span class="k">def</span> <span class="nf">distributed_all_reduce_tensor_average</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">n</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This method performs a reduce operation on multiple nodes running distributed training</span>
<span class="sd">    It first sums all of the results and then divides the summation</span>
<span class="sd">    :param tensor:  The tensor to perform the reduce operation for</span>
<span class="sd">    :param n:  Number of nodes</span>
<span class="sd">    :return:   Averaged tensor from all of the nodes</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">rt</span> <span class="o">=</span> <span class="n">tensor</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">distributed</span><span class="o">.</span><span class="n">all_reduce</span><span class="p">(</span><span class="n">rt</span><span class="p">,</span> <span class="n">op</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">distributed</span><span class="o">.</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">)</span>
    <span class="n">rt</span> <span class="o">/=</span> <span class="n">n</span>
    <span class="k">return</span> <span class="n">rt</span></div>


<div class="viewcode-block" id="reduce_results_tuple_for_ddp"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.reduce_results_tuple_for_ddp">[docs]</a><span class="k">def</span> <span class="nf">reduce_results_tuple_for_ddp</span><span class="p">(</span><span class="n">validation_results_tuple</span><span class="p">,</span> <span class="n">device</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Gather all validation tuples from the various devices and average them&quot;&quot;&quot;</span>
    <span class="n">validation_results_list</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">validation_results_tuple</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">validation_result</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">validation_results_list</span><span class="p">):</span>
        <span class="n">validation_results_list</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">distributed_all_reduce_tensor_average</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">validation_result</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>
                                                                           <span class="n">torch</span><span class="o">.</span><span class="n">distributed</span><span class="o">.</span><span class="n">get_world_size</span><span class="p">())</span>
    <span class="n">validation_results_tuple</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">validation_results_list</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">validation_results_tuple</span></div>


<div class="viewcode-block" id="MultiGPUModeAutocastWrapper"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.MultiGPUModeAutocastWrapper">[docs]</a><span class="k">class</span> <span class="nc">MultiGPUModeAutocastWrapper</span><span class="p">():</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">func</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">func</span> <span class="o">=</span> <span class="n">func</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
        <span class="k">with</span> <span class="n">autocast</span><span class="p">():</span>
            <span class="n">out</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">func</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">out</span></div>


<div class="viewcode-block" id="scaled_all_reduce"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.scaled_all_reduce">[docs]</a><span class="k">def</span> <span class="nf">scaled_all_reduce</span><span class="p">(</span><span class="n">tensors</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">num_gpus</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Performs the scaled all_reduce operation on the provided tensors.</span>
<span class="sd">    The input tensors are modified in-place.</span>
<span class="sd">    Currently supports only the sum</span>
<span class="sd">    reduction operator.</span>
<span class="sd">    The reduced values are scaled by the inverse size of the</span>
<span class="sd">    process group (equivalent to num_gpus).</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># There is no need for reduction in the single-proc case</span>
    <span class="k">if</span> <span class="n">num_gpus</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">tensors</span>

    <span class="c1"># Queue the reductions</span>
    <span class="n">reductions</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="n">tensors</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">distributed</span><span class="o">.</span><span class="n">all_reduce</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">async_op</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="n">reductions</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>

    <span class="c1"># Wait for reductions to finish</span>
    <span class="k">for</span> <span class="n">reduction</span> <span class="ow">in</span> <span class="n">reductions</span><span class="p">:</span>
        <span class="n">reduction</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>

    <span class="c1"># Scale the results</span>
    <span class="k">for</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="n">tensors</span><span class="p">:</span>
        <span class="n">tensor</span><span class="o">.</span><span class="n">mul_</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">/</span> <span class="n">num_gpus</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">tensors</span></div>


<div class="viewcode-block" id="compute_precise_bn_stats"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.compute_precise_bn_stats">[docs]</a><span class="nd">@torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">()</span>
<span class="k">def</span> <span class="nf">compute_precise_bn_stats</span><span class="p">(</span><span class="n">model</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">loader</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">DataLoader</span><span class="p">,</span> <span class="n">precise_bn_batch_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">num_gpus</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    :param model:                   The model being trained (ie: SgModel.net)</span>
<span class="sd">    :param loader:                  Training dataloader (ie: SgModel.train_loader)</span>
<span class="sd">    :param precise_bn_batch_size:   The effective batch size we want to calculate the batchnorm on. For example, if we are training a model</span>
<span class="sd">                                    on 8 gpus, with a batch of 128 on each gpu, a good rule of thumb would be to give it 8192</span>
<span class="sd">                                    (ie: effective_batch_size * num_gpus = batch_per_gpu * num_gpus * num_gpus).</span>
<span class="sd">                                    If precise_bn_batch_size is not provided in the training_params, the latter heuristic</span>
<span class="sd">                                    will be taken.</span>
<span class="sd">    param num_gpus:                 The number of gpus we are training on</span>
<span class="sd">    &#39;&#39;&#39;</span>

    <span class="c1"># Compute the number of minibatches to use</span>
    <span class="n">num_iter</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">precise_bn_batch_size</span> <span class="o">/</span> <span class="p">(</span><span class="n">loader</span><span class="o">.</span><span class="n">batch_size</span> <span class="o">*</span> <span class="n">num_gpus</span><span class="p">))</span> <span class="k">if</span> <span class="n">precise_bn_batch_size</span> <span class="k">else</span> <span class="n">num_gpus</span>
    <span class="n">num_iter</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">num_iter</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">loader</span><span class="p">))</span>

    <span class="c1"># Retrieve the BN layers</span>
    <span class="n">bns</span> <span class="o">=</span> <span class="p">[</span><span class="n">m</span> <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">modules</span><span class="p">()</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">BatchNorm2d</span><span class="p">)]</span>

    <span class="c1"># Initialize BN stats storage for computing mean(mean(batch)) and mean(var(batch))</span>
    <span class="n">running_means</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">bn</span><span class="o">.</span><span class="n">running_mean</span><span class="p">)</span> <span class="k">for</span> <span class="n">bn</span> <span class="ow">in</span> <span class="n">bns</span><span class="p">]</span>
    <span class="n">running_vars</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">bn</span><span class="o">.</span><span class="n">running_var</span><span class="p">)</span> <span class="k">for</span> <span class="n">bn</span> <span class="ow">in</span> <span class="n">bns</span><span class="p">]</span>

    <span class="c1"># Remember momentum values</span>
    <span class="n">momentums</span> <span class="o">=</span> <span class="p">[</span><span class="n">bn</span><span class="o">.</span><span class="n">momentum</span> <span class="k">for</span> <span class="n">bn</span> <span class="ow">in</span> <span class="n">bns</span><span class="p">]</span>

    <span class="c1"># Set momentum to 1.0 to compute BN stats that only reflect the current batch</span>
    <span class="k">for</span> <span class="n">bn</span> <span class="ow">in</span> <span class="n">bns</span><span class="p">:</span>
        <span class="n">bn</span><span class="o">.</span><span class="n">momentum</span> <span class="o">=</span> <span class="mf">1.0</span>

    <span class="c1"># Average the BN stats for each BN layer over the batches</span>
    <span class="k">for</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">_labels</span> <span class="ow">in</span> <span class="n">itertools</span><span class="o">.</span><span class="n">islice</span><span class="p">(</span><span class="n">loader</span><span class="p">,</span> <span class="n">num_iter</span><span class="p">):</span>
        <span class="n">model</span><span class="p">(</span><span class="n">inputs</span><span class="o">.</span><span class="n">cuda</span><span class="p">())</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">bn</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">bns</span><span class="p">):</span>
            <span class="n">running_means</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">bn</span><span class="o">.</span><span class="n">running_mean</span> <span class="o">/</span> <span class="n">num_iter</span>
            <span class="n">running_vars</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">+=</span> <span class="n">bn</span><span class="o">.</span><span class="n">running_var</span> <span class="o">/</span> <span class="n">num_iter</span>

    <span class="c1"># Sync BN stats across GPUs (no reduction if 1 GPU used)</span>
    <span class="n">running_means</span> <span class="o">=</span> <span class="n">scaled_all_reduce</span><span class="p">(</span><span class="n">running_means</span><span class="p">,</span> <span class="n">num_gpus</span><span class="o">=</span><span class="n">num_gpus</span><span class="p">)</span>
    <span class="n">running_vars</span> <span class="o">=</span> <span class="n">scaled_all_reduce</span><span class="p">(</span><span class="n">running_vars</span><span class="p">,</span> <span class="n">num_gpus</span><span class="o">=</span><span class="n">num_gpus</span><span class="p">)</span>

    <span class="c1"># Set BN stats and restore original momentum values</span>
    <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">bn</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">bns</span><span class="p">):</span>
        <span class="n">bn</span><span class="o">.</span><span class="n">running_mean</span> <span class="o">=</span> <span class="n">running_means</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
        <span class="n">bn</span><span class="o">.</span><span class="n">running_var</span> <span class="o">=</span> <span class="n">running_vars</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
        <span class="n">bn</span><span class="o">.</span><span class="n">momentum</span> <span class="o">=</span> <span class="n">momentums</span><span class="p">[</span><span class="n">i</span><span class="p">]</span></div>


<div class="viewcode-block" id="get_local_rank"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.get_local_rank">[docs]</a><span class="k">def</span> <span class="nf">get_local_rank</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns the local rank if running in DDP, and 0 otherwise</span>
<span class="sd">    :return: local rank</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">dist</span><span class="o">.</span><span class="n">get_rank</span><span class="p">()</span> <span class="k">if</span> <span class="n">dist</span><span class="o">.</span><span class="n">is_initialized</span><span class="p">()</span> <span class="k">else</span> <span class="mi">0</span></div>


<div class="viewcode-block" id="get_world_size"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.get_world_size">[docs]</a><span class="k">def</span> <span class="nf">get_world_size</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns the world size if running in DDP, and 1 otherwise</span>
<span class="sd">    :return: world size</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">dist</span><span class="o">.</span><span class="n">is_available</span><span class="p">():</span>
        <span class="k">return</span> <span class="mi">1</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">dist</span><span class="o">.</span><span class="n">is_initialized</span><span class="p">():</span>
        <span class="k">return</span> <span class="mi">1</span>
    <span class="k">return</span> <span class="n">dist</span><span class="o">.</span><span class="n">get_world_size</span><span class="p">()</span></div>


<div class="viewcode-block" id="wait_for_the_master"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.distributed_training_utils.wait_for_the_master">[docs]</a><span class="nd">@contextmanager</span>
<span class="k">def</span> <span class="nf">wait_for_the_master</span><span class="p">(</span><span class="n">local_rank</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Make all processes waiting for the master to do some task.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">local_rank</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">dist</span><span class="o">.</span><span class="n">barrier</span><span class="p">()</span>
    <span class="k">yield</span>
    <span class="k">if</span> <span class="n">local_rank</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">dist</span><span class="o">.</span><span class="n">is_available</span><span class="p">():</span>
            <span class="k">return</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">dist</span><span class="o">.</span><span class="n">is_initialized</span><span class="p">():</span>
            <span class="k">return</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">dist</span><span class="o">.</span><span class="n">barrier</span><span class="p">()</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.early_stopping &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.early_stopping</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.early_stopping</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">super_gradients.training.utils.callbacks</span> <span class="kn">import</span> <span class="n">PhaseCallback</span><span class="p">,</span> <span class="n">Phase</span><span class="p">,</span> <span class="n">PhaseContext</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Optional</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>


<div class="viewcode-block" id="EarlyStop"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.early_stopping.EarlyStop">[docs]</a><span class="k">class</span> <span class="nc">EarlyStop</span><span class="p">(</span><span class="n">PhaseCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Callback to monitor a metric and stop training when it stops improving.</span>
<span class="sd">    Inspired by pytorch_lightning.callbacks.early_stopping and tf.keras.callbacks.EarlyStopping</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">mode_dict</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;min&quot;</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">lt</span><span class="p">,</span> <span class="s2">&quot;max&quot;</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">gt</span><span class="p">}</span>
    <span class="n">supported_phases</span> <span class="o">=</span> <span class="p">(</span><span class="n">Phase</span><span class="o">.</span><span class="n">VALIDATION_EPOCH_END</span><span class="p">,</span> <span class="n">Phase</span><span class="o">.</span><span class="n">TRAIN_EPOCH_END</span><span class="p">)</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                 <span class="n">phase</span><span class="p">:</span> <span class="n">Phase</span><span class="p">,</span>
                 <span class="n">monitor</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
                 <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;min&quot;</span><span class="p">,</span>
                 <span class="n">min_delta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
                 <span class="n">patience</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span>
                 <span class="n">check_finite</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                 <span class="n">threshold</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                 <span class="n">verbose</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
                 <span class="n">strict</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span>
                 <span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">        :param phase: Callback phase event.</span>
<span class="sd">        :param monitor: name of the metric to be monitored.</span>
<span class="sd">        :param mode: one of &#39;min&#39;, &#39;max&#39;. In &#39;min&#39; mode, training will stop when the quantity</span>
<span class="sd">           monitored has stopped decreasing and in &#39;max&#39; mode it will stop when the quantity</span>
<span class="sd">           monitored has stopped increasing.</span>
<span class="sd">        :param min_delta: minimum change in the monitored quantity to qualify as an improvement, i.e. an absolute</span>
<span class="sd">           change of less than `min_delta`, will count as no improvement.</span>
<span class="sd">        :param patience: number of checks with no improvement after which training will be stopped.</span>
<span class="sd">            One check happens after every phase event.</span>
<span class="sd">        :param check_finite: When set ``True``, stops training when the monitor becomes NaN or infinite.</span>
<span class="sd">        :param threshold: Stop training immediately once the monitored quantity reaches this threshold. For mode &#39;min&#39;</span>
<span class="sd">            stops training when below threshold, For mode &#39;max&#39; stops training when above threshold.</span>
<span class="sd">        :param verbose: If `True` print logs.</span>
<span class="sd">        :param strict: whether to crash the training if `monitor` is not found in the metrics.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">EarlyStop</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><span class="n">phase</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">phase</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">supported_phases</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;EarlyStop doesn&#39;t support phase: </span><span class="si">{</span><span class="n">phase</span><span class="si">}</span><span class="s2">, &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;excepted </span><span class="si">{</span><span class="s1">&#39;, &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">str</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">supported_phases</span><span class="p">])</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">phase</span> <span class="o">=</span> <span class="n">phase</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span> <span class="o">=</span> <span class="n">monitor</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_delta</span> <span class="o">=</span> <span class="n">min_delta</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">patience</span> <span class="o">=</span> <span class="n">patience</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="o">=</span> <span class="n">mode</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">check_finite</span> <span class="o">=</span> <span class="n">check_finite</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">threshold</span> <span class="o">=</span> <span class="n">threshold</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">verbose</span> <span class="o">=</span> <span class="n">verbose</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">strict</span> <span class="o">=</span> <span class="n">strict</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">wait_count</span> <span class="o">=</span> <span class="mi">0</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode</span> <span class="ow">not</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode_dict</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;`mode` can be </span><span class="si">{</span><span class="s1">&#39;, &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">mode_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span><span class="si">}</span><span class="s2">, got </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">mode</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">monitor_op</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">mode_dict</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">mode</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_delta</span> <span class="o">*=</span> <span class="mi">1</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">monitor_op</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">gt</span> <span class="k">else</span> <span class="o">-</span><span class="mi">1</span>

        <span class="n">torch_inf</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">np</span><span class="o">.</span><span class="n">Inf</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">best_score</span> <span class="o">=</span> <span class="n">torch_inf</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">monitor_op</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">lt</span> <span class="k">else</span> <span class="o">-</span><span class="n">torch_inf</span>

    <span class="k">def</span> <span class="nf">_get_metric_value</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">metrics_dict</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">metrics_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
            <span class="n">msg</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;Can&#39;t find EarlyStop monitor </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span><span class="si">}</span><span class="s2"> in metrics_dict: </span><span class="si">{</span><span class="n">metrics_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span>
            <span class="n">exception_cls</span> <span class="o">=</span> <span class="ne">RuntimeError</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">strict</span> <span class="k">else</span> <span class="n">MissingMonitorKeyException</span>
            <span class="k">raise</span> <span class="n">exception_cls</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">metrics_dict</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span><span class="p">]</span>

    <span class="k">def</span> <span class="nf">_check_for_early_stop</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">current</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
        <span class="n">should_stop</span> <span class="o">=</span> <span class="kc">False</span>

        <span class="c1"># check if current value is Nan or inf</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">check_finite</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">torch</span><span class="o">.</span><span class="n">isfinite</span><span class="p">(</span><span class="n">current</span><span class="p">):</span>
            <span class="n">should_stop</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="n">reason</span> <span class="o">=</span> <span class="p">(</span>
                <span class="sa">f</span><span class="s2">&quot;Monitored metric </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span><span class="si">}</span><span class="s2"> = </span><span class="si">{</span><span class="n">current</span><span class="si">}</span><span class="s2"> is not finite.&quot;</span>
                <span class="sa">f</span><span class="s2">&quot; Previous best value was </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">best_score</span><span class="si">:</span><span class="s2">.3f</span><span class="si">}</span><span class="s2">. Signaling Trainer to stop.&quot;</span>
            <span class="p">)</span>

        <span class="c1"># check if current value reached threshold value</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">threshold</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">monitor_op</span><span class="p">(</span><span class="n">current</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">threshold</span><span class="p">):</span>
            <span class="n">should_stop</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="n">reason</span> <span class="o">=</span> <span class="p">(</span>
                <span class="s2">&quot;Stopping threshold reached:&quot;</span>
                <span class="sa">f</span><span class="s2">&quot; </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span><span class="si">}</span><span class="s2"> = </span><span class="si">{</span><span class="n">current</span><span class="si">}</span><span class="s2"> </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_op</span><span class="si">}</span><span class="s2"> </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">threshold</span><span class="si">}</span><span class="s2">.&quot;</span>
                <span class="s2">&quot; Signaling Trainer to stop.&quot;</span>
            <span class="p">)</span>

        <span class="c1"># check if current is an improvement of monitor_key metric.</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">monitor_op</span><span class="p">(</span><span class="n">current</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">min_delta</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">best_score</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">current</span><span class="o">.</span><span class="n">device</span><span class="p">)):</span>
            <span class="n">should_stop</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">isfinite</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">best_score</span><span class="p">):</span>
                <span class="n">reason</span> <span class="o">=</span> <span class="p">(</span>
                    <span class="sa">f</span><span class="s2">&quot;Metric </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span><span class="si">}</span><span class="s2"> improved by </span><span class="si">{</span><span class="nb">abs</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">best_score</span> <span class="o">-</span> <span class="n">current</span><span class="p">)</span><span class="si">:</span><span class="s2">.3f</span><span class="si">}</span><span class="s2"> &gt;=&quot;</span>
                    <span class="sa">f</span><span class="s2">&quot; min_delta = </span><span class="si">{</span><span class="nb">abs</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">min_delta</span><span class="p">)</span><span class="si">}</span><span class="s2">. New best score: </span><span class="si">{</span><span class="n">current</span><span class="si">:</span><span class="s2">.3f</span><span class="si">}</span><span class="s2">&quot;</span>
                <span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">reason</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;Metric </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span><span class="si">}</span><span class="s2"> improved. New best score: </span><span class="si">{</span><span class="n">current</span><span class="si">:</span><span class="s2">.3f</span><span class="si">}</span><span class="s2">&quot;</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">best_score</span> <span class="o">=</span> <span class="n">current</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">wait_count</span> <span class="o">=</span> <span class="mi">0</span>

        <span class="c1"># no improvement in monitor_key metric, check if wait_count is bigger than patience.</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">wait_count</span> <span class="o">+=</span> <span class="mi">1</span>
            <span class="n">reason</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;Monitored metric </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">monitor_key</span><span class="si">}</span><span class="s2"> did not improve in the last </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">wait_count</span><span class="si">}</span><span class="s2"> records.&quot;</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">wait_count</span> <span class="o">&gt;=</span> <span class="bp">self</span><span class="o">.</span><span class="n">patience</span><span class="p">:</span>
                <span class="n">should_stop</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="n">reason</span> <span class="o">+=</span> <span class="sa">f</span><span class="s2">&quot; Best score: </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">best_score</span><span class="si">:</span><span class="s2">.3f</span><span class="si">}</span><span class="s2">. Signaling Trainer to stop.&quot;</span>

        <span class="k">return</span> <span class="n">reason</span><span class="p">,</span> <span class="n">should_stop</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">):</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">current</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_metric_value</span><span class="p">(</span><span class="n">context</span><span class="o">.</span><span class="n">metrics_dict</span><span class="p">)</span>
        <span class="k">except</span> <span class="n">MissingMonitorKeyException</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
            <span class="k">return</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">current</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
            <span class="n">current</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">current</span><span class="p">)</span>

        <span class="n">reason</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">should_stop</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_check_for_early_stop</span><span class="p">(</span><span class="n">current</span><span class="p">)</span>

        <span class="c1"># log reason message, and signal early stop if should_stop=True.</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">should_stop</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_signal_early_stop</span><span class="p">(</span><span class="n">context</span><span class="p">,</span> <span class="n">reason</span><span class="p">)</span>

        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">verbose</span><span class="p">:</span>
            <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="n">reason</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_signal_early_stop</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">:</span> <span class="n">PhaseContext</span><span class="p">,</span> <span class="n">reason</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="n">reason</span><span class="p">)</span>
        <span class="n">context</span><span class="o">.</span><span class="n">update_context</span><span class="p">(</span><span class="n">stop_training</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span></div>


<div class="viewcode-block" id="MissingMonitorKeyException"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.early_stopping.MissingMonitorKeyException">[docs]</a><span class="k">class</span> <span class="nc">MissingMonitorKeyException</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Exception raised for missing monitor key in metrics_dict.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">pass</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.ema &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.ema</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.ema</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">math</span>
<span class="kn">import</span> <span class="nn">warnings</span>
<span class="kn">from</span> <span class="nn">copy</span> <span class="kn">import</span> <span class="n">deepcopy</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>

<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.models</span> <span class="kn">import</span> <span class="n">SgModule</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.models.kd_modules.kd_module</span> <span class="kn">import</span> <span class="n">KDModule</span>


<div class="viewcode-block" id="copy_attr"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ema.copy_attr">[docs]</a><span class="k">def</span> <span class="nf">copy_attr</span><span class="p">(</span><span class="n">a</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">b</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">include</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">]</span> <span class="o">=</span> <span class="p">(),</span> <span class="n">exclude</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">]</span> <span class="o">=</span> <span class="p">()):</span>
    <span class="c1"># Copy attributes from b to a, options to only include [...] and to exclude [...]</span>
    <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">b</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">include</span><span class="p">)</span> <span class="ow">and</span> <span class="n">k</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">include</span><span class="p">)</span> <span class="ow">or</span> <span class="n">k</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;_&#39;</span><span class="p">)</span> <span class="ow">or</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">exclude</span><span class="p">:</span>
            <span class="k">continue</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span></div>


<div class="viewcode-block" id="ModelEMA"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ema.ModelEMA">[docs]</a><span class="k">class</span> <span class="nc">ModelEMA</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot; Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models</span>
<span class="sd">    Keep a moving average of everything in the model state_dict (parameters and buffers).</span>
<span class="sd">    This is intended to allow functionality like</span>
<span class="sd">    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage</span>
<span class="sd">    A smoothed version of the weights is necessary for some training schemes to perform well.</span>
<span class="sd">    This class is sensitive where it is initialized in the sequence of model init,</span>
<span class="sd">    GPU assignment and distributed training wrappers.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">decay</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.9999</span><span class="p">,</span> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">15</span><span class="p">,</span> <span class="n">exp_activation</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Init the EMA</span>
<span class="sd">        :param model: Union[SgModule, nn.Module], the training model to construct the EMA model by</span>
<span class="sd">                    IMPORTANT: WHEN THE APPLICATION OF EMA ONLY ON A SUBSET OF ATTRIBUTES IS DESIRED, WRAP THE NN.MODULE</span>
<span class="sd">                    AS SgModule AND OVERWRITE get_include_attributes() AND get_exclude_attributes() AS DESIRED.</span>
<span class="sd">        :param decay: the maximum decay value. as the training process advances, the decay will climb towards this value</span>
<span class="sd">                      until the EMA_t+1 = EMA_t * decay + TRAINING_MODEL * (1- decay)</span>
<span class="sd">        :param beta: the exponent coefficient. The higher the beta, the sooner in the training the decay will saturate to</span>
<span class="sd">                     its final value. beta=15 is ~40% of the training process.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># Create EMA</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ema</span> <span class="o">=</span> <span class="n">deepcopy</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ema</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">exp_activation</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">decay_function</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">decay</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">math</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="o">-</span><span class="n">x</span> <span class="o">*</span> <span class="n">beta</span><span class="p">))</span>  <span class="c1"># decay exponential ramp (to help early epochs)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">decay_function</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">decay</span>  <span class="c1"># always return the same decay factor</span>

        <span class="sd">&quot;&quot;&quot;&quot;</span>
<span class="sd">        we hold a list of model attributes (not wights and biases) which we would like to include in each</span>
<span class="sd">        attribute update or exclude from each update. a SgModule declare these attribute using</span>
<span class="sd">        get_include_attributes and get_exclude_attributes functions. for a nn.Module which is not a SgModule</span>
<span class="sd">        all non-private (not starting with &#39;_&#39;) attributes will be updated (and only them).</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="n">SgModule</span><span class="p">):</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">include_attributes</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">get_include_attributes</span><span class="p">()</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">exclude_attributes</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">get_exclude_attributes</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;Warning: EMA should be used with SgModule instance. All attributes of the model will be &quot;</span>
                          <span class="s2">&quot;included in EMA&quot;</span><span class="p">)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">include_attributes</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">exclude_attributes</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">ema</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">parameters</span><span class="p">():</span>
            <span class="n">p</span><span class="o">.</span><span class="n">requires_grad_</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

<div class="viewcode-block" id="ModelEMA.update"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ema.ModelEMA.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">training_percent</span><span class="p">:</span> <span class="nb">float</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Update the state of the EMA model.</span>
<span class="sd">        :param model: current training model</span>
<span class="sd">        :param training_percent: the percentage of the training process [0,1]. i.e 0.4 means 40% of the training have passed</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># Update EMA parameters</span>
        <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">():</span>
            <span class="n">decay</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">decay_function</span><span class="p">(</span><span class="n">training_percent</span><span class="p">)</span>

            <span class="k">for</span> <span class="n">ema_v</span><span class="p">,</span> <span class="n">model_v</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ema</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span><span class="o">.</span><span class="n">values</span><span class="p">(),</span> <span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span><span class="o">.</span><span class="n">values</span><span class="p">()):</span>
                <span class="k">if</span> <span class="n">ema_v</span><span class="o">.</span><span class="n">dtype</span><span class="o">.</span><span class="n">is_floating_point</span><span class="p">:</span>
                    <span class="n">ema_v</span><span class="o">.</span><span class="n">copy_</span><span class="p">(</span><span class="n">ema_v</span> <span class="o">*</span> <span class="n">decay</span> <span class="o">+</span> <span class="p">(</span><span class="mf">1.</span> <span class="o">-</span> <span class="n">decay</span><span class="p">)</span> <span class="o">*</span> <span class="n">model_v</span><span class="o">.</span><span class="n">detach</span><span class="p">())</span></div>

<div class="viewcode-block" id="ModelEMA.update_attr"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ema.ModelEMA.update_attr">[docs]</a>    <span class="k">def</span> <span class="nf">update_attr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This function updates model attributes (not weight and biases) from original model to the ema model.</span>
<span class="sd">        attributes of the original model, such as anchors and grids (of detection models), may be crucial to the</span>
<span class="sd">        model operation and need to be updated.</span>
<span class="sd">        If include_attributes and exclude_attributes lists were not defined, all non-private (not starting with &#39;_&#39;)</span>
<span class="sd">        attributes will be updated (and only them).</span>
<span class="sd">        :param model: the source model</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">copy_attr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">ema</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="n">model</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">include_attributes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">exclude_attributes</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="KDModelEMA"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ema.KDModelEMA">[docs]</a><span class="k">class</span> <span class="nc">KDModelEMA</span><span class="p">(</span><span class="n">ModelEMA</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot; Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models</span>
<span class="sd">    Keep a moving average of everything in the model state_dict (parameters and buffers).</span>
<span class="sd">    This is intended to allow functionality like</span>
<span class="sd">    https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage</span>
<span class="sd">    A smoothed version of the weights is necessary for some training schemes to perform well.</span>
<span class="sd">    This class is sensitive where it is initialized in the sequence of model init,</span>
<span class="sd">    GPU assignment and distributed training wrappers.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">kd_model</span><span class="p">:</span> <span class="n">KDModule</span><span class="p">,</span> <span class="n">decay</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.9999</span><span class="p">,</span> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">15</span><span class="p">,</span> <span class="n">exp_activation</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Init the EMA</span>
<span class="sd">        :param kd_model: KDModule, the training Knowledge distillation model to construct the EMA model by</span>
<span class="sd">                    IMPORTANT: WHEN THE APPLICATION OF EMA ONLY ON A SUBSET OF ATTRIBUTES IS DESIRED, WRAP THE NN.MODULE</span>
<span class="sd">                    AS SgModule AND OVERWRITE get_include_attributes() AND get_exclude_attributes() AS DESIRED.</span>
<span class="sd">        :param decay: the maximum decay value. as the training process advances, the decay will climb towards this value</span>
<span class="sd">                      until the EMA_t+1 = EMA_t * decay + TRAINING_MODEL * (1- decay)</span>
<span class="sd">        :param beta: the exponent coefficient. The higher the beta, the sooner in the training the decay will saturate to</span>
<span class="sd">                     its final value. beta=15 is ~40% of the training process.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># Only work on the student (we don&#39;t want to update and to have a duplicate of the teacher)</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">model</span><span class="o">=</span><span class="n">core_utils</span><span class="o">.</span><span class="n">WrappedModel</span><span class="p">(</span><span class="n">kd_model</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">student</span><span class="p">),</span>
                         <span class="n">decay</span><span class="o">=</span><span class="n">decay</span><span class="p">,</span>
                         <span class="n">beta</span><span class="o">=</span><span class="n">beta</span><span class="p">,</span>
                         <span class="n">exp_activation</span><span class="o">=</span><span class="n">exp_activation</span><span class="p">)</span>

        <span class="c1"># Overwrite current ema attribute with combination of the student model EMA (current self.ema)</span>
        <span class="c1"># with already the instantiated teacher, to have the final KD EMA</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ema</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">WrappedModel</span><span class="p">(</span><span class="n">KDModule</span><span class="p">(</span><span class="n">arch_params</span><span class="o">=</span><span class="n">kd_model</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">arch_params</span><span class="p">,</span>
                                                    <span class="n">student</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ema</span><span class="o">.</span><span class="n">module</span><span class="p">,</span>
                                                    <span class="n">teacher</span><span class="o">=</span><span class="n">kd_model</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">teacher</span><span class="p">,</span>
                                                    <span class="n">run_teacher_on_eval</span><span class="o">=</span><span class="n">kd_model</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">run_teacher_on_eval</span><span class="p">))</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.export_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.export_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.export_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>


<div class="viewcode-block" id="fuse_conv_bn"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.export_utils.fuse_conv_bn">[docs]</a><span class="k">def</span> <span class="nf">fuse_conv_bn</span><span class="p">(</span><span class="n">model</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">replace_bn_with_identity</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Fuses consecutive nn.Conv2d and nn.BatchNorm2d layers recursively inplace in all of the model</span>
<span class="sd">    :param replace_bn_with_identity: if set to true, bn will be replaced with identity. otherwise, bn will be removed</span>
<span class="sd">    :param model: the target model</span>
<span class="sd">    :return: the number of fuses executed</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">children</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">named_children</span><span class="p">())</span>
    <span class="n">counter</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">for</span> <span class="n">i</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">children</span><span class="p">)</span> <span class="o">-</span> <span class="mi">1</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">BatchNorm2d</span><span class="p">):</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">children</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">fuse_conv_bn_eval</span><span class="p">(</span><span class="n">children</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span> <span class="n">children</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">][</span><span class="mi">1</span><span class="p">]))</span>
            <span class="k">if</span> <span class="n">replace_bn_with_identity</span><span class="p">:</span>
                <span class="nb">setattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">children</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span> <span class="n">nn</span><span class="o">.</span><span class="n">Identity</span><span class="p">())</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="nb">delattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">children</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">][</span><span class="mi">0</span><span class="p">])</span>
            <span class="n">counter</span> <span class="o">+=</span> <span class="mi">1</span>
    <span class="k">for</span> <span class="n">child_name</span><span class="p">,</span> <span class="n">child</span> <span class="ow">in</span> <span class="n">children</span><span class="p">:</span>
        <span class="n">counter</span> <span class="o">+=</span> <span class="n">fuse_conv_bn</span><span class="p">(</span><span class="n">child</span><span class="p">,</span> <span class="n">replace_bn_with_identity</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">counter</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.module_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.module_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.module_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">collections</span> <span class="kn">import</span> <span class="n">OrderedDict</span>
<span class="kn">import</span> <span class="nn">copy</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Tuple</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>


<div class="viewcode-block" id="MultiOutputModule"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.MultiOutputModule">[docs]</a><span class="k">class</span> <span class="nc">MultiOutputModule</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This module wraps around a container nn.Module (such as Module, Sequential and ModuleList) and allows to extract</span>
<span class="sd">    multiple output from its inner modules on each forward call() (as a list of output tensors)</span>
<span class="sd">    note: the default output of the wrapped module will not be added to the output list by default. To get</span>
<span class="sd">    the default output in the outputs list, explicitly include its path in the @output_paths parameter</span>

<span class="sd">    i.e.</span>
<span class="sd">    for module:</span>
<span class="sd">        Sequential(</span>
<span class="sd">          (0): Sequential(</span>
<span class="sd">            (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)</span>
<span class="sd">            (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)</span>
<span class="sd">            (2): ReLU6(inplace=True)</span>
<span class="sd">          )                                         ===================================&gt;&gt;</span>
<span class="sd">          (1): InvertedResidual(</span>
<span class="sd">            (conv): Sequential(</span>
<span class="sd">              (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)</span>
<span class="sd">              (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)</span>
<span class="sd">              (2): ReLU6(inplace=True)              ===================================&gt;&gt;</span>
<span class="sd">              (3): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)</span>
<span class="sd">              (4): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)</span>
<span class="sd">            )</span>
<span class="sd">          )</span>
<span class="sd">        )</span>
<span class="sd">    and paths:</span>
<span class="sd">        [0, [1, &#39;conv&#39;, 2]]</span>
<span class="sd">    the output are marked with arrows</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">output_paths</span><span class="p">:</span> <span class="nb">list</span><span class="p">,</span> <span class="n">prune</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        :param module: The wrapped container module</span>
<span class="sd">        :param output_paths: a list of lists or keys containing the canonical paths to the outputs</span>
<span class="sd">        i.e. [3, [4, &#39;conv&#39;, 5], 7] will extract outputs of layers 3, 7 and 4-&gt;conv-&gt;5</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">output_paths</span> <span class="o">=</span> <span class="n">output_paths</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_modules</span><span class="p">[</span><span class="s1">&#39;0&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">module</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_outputs_lists</span> <span class="o">=</span> <span class="p">{}</span>

        <span class="k">for</span> <span class="n">path</span> <span class="ow">in</span> <span class="n">output_paths</span><span class="p">:</span>
            <span class="n">child</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_recursive</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span>
            <span class="n">child</span><span class="o">.</span><span class="n">register_forward_hook</span><span class="p">(</span><span class="n">hook</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">save_output_hook</span><span class="p">)</span>

        <span class="c1"># PRUNE THE MODULE TO SUPPORT ALL PROVIDED OUTPUT_PATHS BUT REMOVE ALL REDUNDANT LAYERS</span>
        <span class="k">if</span> <span class="n">prune</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_prune</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">output_paths</span><span class="p">)</span>

<div class="viewcode-block" id="MultiOutputModule.save_output_hook"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.MultiOutputModule.save_output_hook">[docs]</a>    <span class="k">def</span> <span class="nf">save_output_hook</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_outputs_lists</span><span class="p">[</span><span class="nb">input</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">device</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">output</span><span class="p">)</span></div>

<div class="viewcode-block" id="MultiOutputModule.forward"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.MultiOutputModule.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_outputs_lists</span><span class="p">[</span><span class="n">x</span><span class="o">.</span><span class="n">device</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_modules</span><span class="p">[</span><span class="s1">&#39;0&#39;</span><span class="p">](</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_outputs_lists</span><span class="p">[</span><span class="n">x</span><span class="o">.</span><span class="n">device</span><span class="p">]</span></div>

    <span class="k">def</span> <span class="nf">_get_recursive</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;recursively look for a module using a path&quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="k">return</span> <span class="n">module</span><span class="o">.</span><span class="n">_modules</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="n">path</span><span class="p">)]</span>
        <span class="k">elif</span> <span class="nb">len</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">module</span><span class="o">.</span><span class="n">_modules</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="n">path</span><span class="p">[</span><span class="mi">0</span><span class="p">])]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_recursive</span><span class="p">(</span><span class="n">module</span><span class="o">.</span><span class="n">_modules</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="n">path</span><span class="p">[</span><span class="mi">0</span><span class="p">])],</span> <span class="n">path</span><span class="p">[</span><span class="mi">1</span><span class="p">:])</span>

    <span class="k">def</span> <span class="nf">_prune</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">output_paths</span><span class="p">:</span> <span class="nb">list</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Recursively prune the module to support all provided output_paths but remove all redundant layers</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">last_index</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
        <span class="n">last_key</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="c1"># look for the last key from all paths</span>
        <span class="k">for</span> <span class="n">path</span> <span class="ow">in</span> <span class="n">output_paths</span><span class="p">:</span>
            <span class="n">key</span> <span class="o">=</span> <span class="n">path</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="nb">list</span><span class="p">)</span> <span class="k">else</span> <span class="n">path</span>
            <span class="n">index</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">module</span><span class="o">.</span><span class="n">_modules</span><span class="p">)</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">key</span><span class="p">))</span>
            <span class="k">if</span> <span class="n">index</span> <span class="o">&gt;</span> <span class="n">last_index</span><span class="p">:</span>
                <span class="n">last_index</span> <span class="o">=</span> <span class="n">index</span>
                <span class="n">last_key</span> <span class="o">=</span> <span class="n">key</span>

        <span class="n">module</span><span class="o">.</span><span class="n">_modules</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_slice_odict</span><span class="p">(</span><span class="n">module</span><span class="o">.</span><span class="n">_modules</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">last_index</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>

        <span class="n">next_level_paths</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">path</span> <span class="ow">in</span> <span class="n">output_paths</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="nb">list</span><span class="p">)</span> <span class="ow">and</span> <span class="n">path</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">last_key</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">path</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
                <span class="n">next_level_paths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">path</span><span class="p">[</span><span class="mi">1</span><span class="p">:])</span>

        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">next_level_paths</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_prune</span><span class="p">(</span><span class="n">module</span><span class="o">.</span><span class="n">_modules</span><span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="n">last_key</span><span class="p">)],</span> <span class="n">next_level_paths</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_slice_odict</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">odict</span><span class="p">:</span> <span class="n">OrderedDict</span><span class="p">,</span> <span class="n">start</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">end</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Slice an OrderedDict in the same logic list,tuple... are sliced&quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">OrderedDict</span><span class="p">([</span>
            <span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span> <span class="k">for</span> <span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span> <span class="ow">in</span> <span class="n">odict</span><span class="o">.</span><span class="n">items</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">k</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="n">odict</span><span class="o">.</span><span class="n">keys</span><span class="p">())[</span><span class="n">start</span><span class="p">:</span><span class="n">end</span><span class="p">]</span>
        <span class="p">])</span></div>


<span class="k">def</span> <span class="nf">_replace_activations_recursive</span><span class="p">(</span><span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">new_activation</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">activations_to_replace</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">type</span><span class="p">]):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    A helper called in replace_activations(...)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">for</span> <span class="n">n</span><span class="p">,</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">module</span><span class="o">.</span><span class="n">named_children</span><span class="p">():</span>
        <span class="k">if</span> <span class="nb">type</span><span class="p">(</span><span class="n">m</span><span class="p">)</span> <span class="ow">in</span> <span class="n">activations_to_replace</span><span class="p">:</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">n</span><span class="p">,</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">new_activation</span><span class="p">))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">_replace_activations_recursive</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">new_activation</span><span class="p">,</span> <span class="n">activations_to_replace</span><span class="p">)</span>


<div class="viewcode-block" id="replace_activations"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.replace_activations">[docs]</a><span class="k">def</span> <span class="nf">replace_activations</span><span class="p">(</span><span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">new_activation</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">activations_to_replace</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">type</span><span class="p">]):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Recursively go through module and replaces each activation in activations_to_replace with a copy of new_activation</span>
<span class="sd">    :param module:                  a module that will be changed inplace</span>
<span class="sd">    :param new_activation:          a sample of a new activation (will be copied)</span>
<span class="sd">    :param activations_to_replace:  types of activations to replace, each must be a subclass of nn.Module</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># check arguments once before the recursion</span>
    <span class="k">assert</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">new_activation</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">),</span> <span class="s1">&#39;new_activation should be nn.Module&#39;</span>
    <span class="k">assert</span> <span class="nb">all</span><span class="p">([</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="nb">type</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">)</span> <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">activations_to_replace</span><span class="p">]),</span> \
        <span class="s1">&#39;activations_to_replace should be types that are subclasses of nn.Module&#39;</span>

    <span class="c1"># do the replacement</span>
    <span class="n">_replace_activations_recursive</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">new_activation</span><span class="p">,</span> <span class="n">activations_to_replace</span><span class="p">)</span></div>


<div class="viewcode-block" id="fuse_repvgg_blocks_residual_branches"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.fuse_repvgg_blocks_residual_branches">[docs]</a><span class="k">def</span> <span class="nf">fuse_repvgg_blocks_residual_branches</span><span class="p">(</span><span class="n">model</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="sd">&#39;&#39;&#39;</span>
<span class="sd">    Call fuse_block_residual_branches for all repvgg blocks in the model</span>
<span class="sd">    :param model: torch.nn.Module with repvgg blocks. Doesn&#39;t have to be entirely consists of repvgg.</span>
<span class="sd">    :type model: torch.nn.Module</span>
<span class="sd">    &#39;&#39;&#39;</span>
    <span class="k">assert</span> <span class="ow">not</span> <span class="n">model</span><span class="o">.</span><span class="n">training</span><span class="p">,</span> <span class="s2">&quot;To fuse RepVGG block residual branches, model must be on eval mode&quot;</span>
    <span class="k">for</span> <span class="n">module</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">modules</span><span class="p">():</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="s1">&#39;fuse_block_residual_branches&#39;</span><span class="p">):</span>
            <span class="n">module</span><span class="o">.</span><span class="n">fuse_block_residual_branches</span><span class="p">()</span>
    <span class="n">model</span><span class="o">.</span><span class="n">build_residual_branches</span> <span class="o">=</span> <span class="kc">False</span></div>


<div class="viewcode-block" id="ConvBNReLU"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.ConvBNReLU">[docs]</a><span class="k">class</span> <span class="nc">ConvBNReLU</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Class for Convolution2d-Batchnorm2d-Relu layer. Default behaviour is Conv-BN-Relu. To exclude Batchnorm module use</span>
<span class="sd">        `use_normalization=False`, to exclude Relu activation use `use_activation=False`.</span>
<span class="sd">    For convolution arguments documentation see `nn.Conv2d`.</span>
<span class="sd">    For batchnorm arguments documentation see `nn.BatchNorm2d`.</span>
<span class="sd">    For relu arguments documentation see `nn.Relu`.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span>
                 <span class="n">in_channels</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                 <span class="n">out_channels</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                 <span class="n">kernel_size</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]],</span>
                 <span class="n">stride</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
                 <span class="n">padding</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
                 <span class="n">dilation</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
                 <span class="n">groups</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
                 <span class="n">bias</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                 <span class="n">padding_mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;zeros&#39;</span><span class="p">,</span>
                 <span class="n">use_normalization</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                 <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span>
                 <span class="n">momentum</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
                 <span class="n">affine</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                 <span class="n">track_running_stats</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                 <span class="n">device</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">dtype</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                 <span class="n">use_activation</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
                 <span class="n">inplace</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>

        <span class="nb">super</span><span class="p">(</span><span class="n">ConvBNReLU</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">seq</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">seq</span><span class="o">.</span><span class="n">add_module</span><span class="p">(</span><span class="s2">&quot;conv&quot;</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span><span class="p">(</span><span class="n">in_channels</span><span class="p">,</span>
                                              <span class="n">out_channels</span><span class="p">,</span>
                                              <span class="n">kernel_size</span><span class="o">=</span><span class="n">kernel_size</span><span class="p">,</span>
                                              <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
                                              <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
                                              <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span>
                                              <span class="n">groups</span><span class="o">=</span><span class="n">groups</span><span class="p">,</span>
                                              <span class="n">bias</span><span class="o">=</span><span class="n">bias</span><span class="p">,</span>
                                              <span class="n">padding_mode</span><span class="o">=</span><span class="n">padding_mode</span><span class="p">))</span>

        <span class="k">if</span> <span class="n">use_normalization</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">seq</span><span class="o">.</span><span class="n">add_module</span><span class="p">(</span><span class="s2">&quot;bn&quot;</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">BatchNorm2d</span><span class="p">(</span><span class="n">out_channels</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span> <span class="n">momentum</span><span class="o">=</span><span class="n">momentum</span><span class="p">,</span> <span class="n">affine</span><span class="o">=</span><span class="n">affine</span><span class="p">,</span>
                                                     <span class="n">track_running_stats</span><span class="o">=</span><span class="n">track_running_stats</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">dtype</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">use_activation</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">seq</span><span class="o">.</span><span class="n">add_module</span><span class="p">(</span><span class="s2">&quot;relu&quot;</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">ReLU</span><span class="p">(</span><span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">))</span>

<div class="viewcode-block" id="ConvBNReLU.forward"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.ConvBNReLU.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">seq</span><span class="p">(</span><span class="n">x</span><span class="p">)</span></div></div>


<div class="viewcode-block" id="NormalizationAdapter"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.NormalizationAdapter">[docs]</a><span class="k">class</span> <span class="nc">NormalizationAdapter</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Denormalizes input by mean_original, std_original, then normalizes by mean_required, std_required.</span>

<span class="sd">    Used in KD training where teacher expects data normalized by mean_required, std_required.</span>

<span class="sd">    mean_original, std_original, mean_required, std_required are all list-like objects of length that&#39;s equal to the</span>
<span class="sd">     number of input channels.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mean_original</span><span class="p">,</span> <span class="n">std_original</span><span class="p">,</span> <span class="n">mean_required</span><span class="p">,</span> <span class="n">std_required</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">NormalizationAdapter</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>
        <span class="n">mean_original</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">mean_original</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">std_original</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">std_original</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">mean_required</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">mean_required</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">std_required</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">std_required</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">additive</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">mean_original</span> <span class="o">-</span> <span class="n">mean_required</span><span class="p">)</span> <span class="o">/</span> <span class="n">std_original</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">multiplier</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">std_original</span> <span class="o">/</span> <span class="n">std_required</span><span class="p">)</span>

<div class="viewcode-block" id="NormalizationAdapter.forward"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.module_utils.NormalizationAdapter.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="n">x</span> <span class="o">=</span> <span class="p">(</span><span class="n">x</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">additive</span><span class="p">)</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">multiplier</span>
        <span class="k">return</span> <span class="n">x</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.optimizer_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.optimizer_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.optimizer_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch.optim</span> <span class="k">as</span> <span class="nn">optim</span>
<span class="kn">import</span> <span class="nn">torch.nn</span> <span class="k">as</span> <span class="nn">nn</span>
<span class="kn">from</span> <span class="nn">torch.nn.modules.batchnorm</span> <span class="kn">import</span> <span class="n">_BatchNorm</span>
<span class="kn">from</span> <span class="nn">torch.nn.modules.conv</span> <span class="kn">import</span> <span class="n">_ConvNd</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.abstractions.abstract_logger</span> <span class="kn">import</span> <span class="n">get_logger</span>
<span class="kn">from</span> <span class="nn">super_gradients.common.factories.optimizers_type_factory</span> <span class="kn">import</span> <span class="n">OptimizersTypeFactory</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.params</span> <span class="kn">import</span> <span class="n">DEFAULT_OPTIMIZER_PARAMS_SGD</span><span class="p">,</span> <span class="n">DEFAULT_OPTIMIZER_PARAMS_ADAM</span><span class="p">,</span> \
    <span class="n">DEFAULT_OPTIMIZER_PARAMS_RMSPROP</span><span class="p">,</span> <span class="n">DEFAULT_OPTIMIZER_PARAMS_RMSPROPTF</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils</span> <span class="kn">import</span> <span class="n">get_param</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.optimizers.rmsprop_tf</span> <span class="kn">import</span> <span class="n">RMSpropTF</span>

<span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>

<span class="n">OPTIMIZERS_DEFAULT_PARAMS</span> <span class="o">=</span> <span class="p">{</span><span class="n">optim</span><span class="o">.</span><span class="n">SGD</span><span class="p">:</span> <span class="n">DEFAULT_OPTIMIZER_PARAMS_SGD</span><span class="p">,</span>
                             <span class="n">optim</span><span class="o">.</span><span class="n">Adam</span><span class="p">:</span> <span class="n">DEFAULT_OPTIMIZER_PARAMS_ADAM</span><span class="p">,</span>
                             <span class="n">optim</span><span class="o">.</span><span class="n">RMSprop</span><span class="p">:</span> <span class="n">DEFAULT_OPTIMIZER_PARAMS_RMSPROP</span><span class="p">,</span>
                             <span class="n">RMSpropTF</span><span class="p">:</span> <span class="n">DEFAULT_OPTIMIZER_PARAMS_RMSPROPTF</span><span class="p">}</span>


<div class="viewcode-block" id="separate_zero_wd_params_groups_for_optimizer"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.optimizer_utils.separate_zero_wd_params_groups_for_optimizer">[docs]</a><span class="k">def</span> <span class="nf">separate_zero_wd_params_groups_for_optimizer</span><span class="p">(</span><span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">net_named_params</span><span class="p">,</span> <span class="n">weight_decay</span><span class="p">:</span> <span class="nb">float</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    separate param groups for batchnorm and biases and others with weight decay. return list of param groups in format</span>
<span class="sd">     required by torch Optimizer classes.</span>
<span class="sd">    bias + BN with weight decay=0 and the rest with the given weight decay</span>
<span class="sd">        :param module: train net module.</span>
<span class="sd">        :param net_named_params: list of params groups, output of SgModule.initialize_param_groups</span>
<span class="sd">        :param weight_decay: value to set for the non BN and bias parameters</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># FIXME - replace usage of ids addresses to find batchnorm and biases params.</span>
    <span class="c1">#  This solution iterate 2 times over module parameters, find a way to iterate only one time.</span>
    <span class="n">no_decay_ids</span> <span class="o">=</span> <span class="n">_get_no_decay_param_ids</span><span class="p">(</span><span class="n">module</span><span class="p">)</span>
    <span class="c1"># split param groups for optimizer</span>
    <span class="n">optimizer_param_groups</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">param_group</span> <span class="ow">in</span> <span class="n">net_named_params</span><span class="p">:</span>
        <span class="n">no_decay_params</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">decay_params</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">param</span> <span class="ow">in</span> <span class="n">param_group</span><span class="p">[</span><span class="s2">&quot;named_params&quot;</span><span class="p">]:</span>
            <span class="k">if</span> <span class="nb">id</span><span class="p">(</span><span class="n">param</span><span class="p">)</span> <span class="ow">in</span> <span class="n">no_decay_ids</span><span class="p">:</span>
                <span class="n">no_decay_params</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">param</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">decay_params</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">param</span><span class="p">)</span>
        <span class="c1"># append two param groups from the original param group, with and without weight decay.</span>
        <span class="n">extra_optim_params</span> <span class="o">=</span> <span class="p">{</span><span class="n">key</span><span class="p">:</span> <span class="n">param_group</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="n">param_group</span>
                              <span class="k">if</span> <span class="n">key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;named_params&quot;</span><span class="p">,</span> <span class="s2">&quot;weight_decay&quot;</span><span class="p">]}</span>
        <span class="n">optimizer_param_groups</span><span class="o">.</span><span class="n">append</span><span class="p">({</span><span class="s2">&quot;params&quot;</span><span class="p">:</span> <span class="n">no_decay_params</span><span class="p">,</span> <span class="s2">&quot;weight_decay&quot;</span><span class="p">:</span> <span class="mf">0.0</span><span class="p">,</span> <span class="o">**</span><span class="n">extra_optim_params</span><span class="p">})</span>
        <span class="n">optimizer_param_groups</span><span class="o">.</span><span class="n">append</span><span class="p">({</span><span class="s2">&quot;params&quot;</span><span class="p">:</span> <span class="n">decay_params</span><span class="p">,</span> <span class="s2">&quot;weight_decay&quot;</span><span class="p">:</span> <span class="n">weight_decay</span><span class="p">,</span> <span class="o">**</span><span class="n">extra_optim_params</span><span class="p">})</span>

    <span class="k">return</span> <span class="n">optimizer_param_groups</span></div>


<span class="k">def</span> <span class="nf">_get_no_decay_param_ids</span><span class="p">(</span><span class="n">module</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="c1"># FIXME - replace usage of ids addresses to find batchnorm and biases params.</span>
    <span class="c1">#  Use other common way to identify torch parameters other than id or layer names</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Iterate over module.modules() and returns params id addresses of batch-norm and biases params.</span>
<span class="sd">    NOTE - ALL MODULES WITH ATTRIBUTES NAMED BIAS AND ARE INSTANCE OF nn.Parameter WILL BE CONSIDERED A BIAS PARAM FOR</span>
<span class="sd">        ZERO WEIGHT DECAY.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">batchnorm_types</span> <span class="o">=</span> <span class="p">(</span><span class="n">_BatchNorm</span><span class="p">,)</span>
    <span class="n">torch_weight_with_bias_types</span> <span class="o">=</span> <span class="p">(</span><span class="n">_ConvNd</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">)</span>
    <span class="n">no_decay_ids</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">module</span><span class="o">.</span><span class="n">named_modules</span><span class="p">():</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">batchnorm_types</span><span class="p">):</span>
            <span class="n">no_decay_ids</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">id</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">weight</span><span class="p">))</span>
            <span class="n">no_decay_ids</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">id</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">bias</span><span class="p">))</span>
        <span class="k">elif</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="s2">&quot;bias&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">bias</span><span class="p">,</span> <span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span><span class="p">):</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">torch_weight_with_bias_types</span><span class="p">):</span>
                <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Module class: </span><span class="si">{</span><span class="n">m</span><span class="o">.</span><span class="vm">__class__</span><span class="si">}</span><span class="s2">, have a `bias` parameter attribute but is not instance of&quot;</span>
                               <span class="sa">f</span><span class="s2">&quot; torch primitive modules, this bias parameter will be part of param group with zero&quot;</span>
                               <span class="sa">f</span><span class="s2">&quot; weight decay.&quot;</span><span class="p">)</span>
            <span class="n">no_decay_ids</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">id</span><span class="p">(</span><span class="n">m</span><span class="o">.</span><span class="n">bias</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">no_decay_ids</span>


<div class="viewcode-block" id="build_optimizer"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.optimizer_utils.build_optimizer">[docs]</a><span class="k">def</span> <span class="nf">build_optimizer</span><span class="p">(</span><span class="n">net</span><span class="p">,</span> <span class="n">lr</span><span class="p">,</span> <span class="n">training_params</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Wrapper function for initializing the optimizer</span>
<span class="sd">        :param net: the nn_module to build the optimizer for</span>
<span class="sd">        :param lr: initial learning rate</span>
<span class="sd">        :param training_params: training_parameters</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">training_params</span><span class="o">.</span><span class="n">optimizer</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">optimizer_cls</span> <span class="o">=</span> <span class="n">OptimizersTypeFactory</span><span class="p">()</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">training_params</span><span class="o">.</span><span class="n">optimizer</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">optimizer_cls</span> <span class="o">=</span> <span class="n">training_params</span><span class="o">.</span><span class="n">optimizer</span>
    <span class="n">default_optimizer_params</span> <span class="o">=</span> <span class="n">OPTIMIZERS_DEFAULT_PARAMS</span><span class="p">[</span><span class="n">optimizer_cls</span><span class="p">]</span> <span class="k">if</span> <span class="n">optimizer_cls</span> <span class="ow">in</span> <span class="n">OPTIMIZERS_DEFAULT_PARAMS</span> <span class="k">else</span> <span class="p">{}</span>
    <span class="n">training_params</span><span class="o">.</span><span class="n">optimizer_params</span> <span class="o">=</span> <span class="n">get_param</span><span class="p">(</span><span class="n">training_params</span><span class="p">,</span> <span class="s1">&#39;optimizer_params&#39;</span><span class="p">,</span> <span class="n">default_optimizer_params</span><span class="p">)</span>

    <span class="c1"># OPTIMIZER PARAM GROUPS ARE SET USING DEFAULT OR MODEL SPECIFIC INIT</span>
    <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="s1">&#39;initialize_param_groups&#39;</span><span class="p">):</span>
        <span class="c1"># INITIALIZE_PARAM_GROUPS MUST RETURN A LIST OF DICTS WITH &#39;named_params&#39; AND OPTIMIZER&#39;s ATTRIBUTES PER GROUP</span>
        <span class="n">net_named_params</span> <span class="o">=</span> <span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="o">.</span><span class="n">initialize_param_groups</span><span class="p">(</span><span class="n">lr</span><span class="p">,</span> <span class="n">training_params</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">net_named_params</span> <span class="o">=</span> <span class="p">[{</span><span class="s1">&#39;named_params&#39;</span><span class="p">:</span> <span class="n">net</span><span class="o">.</span><span class="n">named_parameters</span><span class="p">()}]</span>

    <span class="k">if</span> <span class="n">training_params</span><span class="o">.</span><span class="n">zero_weight_decay_on_bias_and_bn</span><span class="p">:</span>
        <span class="n">optimizer_training_params</span> <span class="o">=</span> <span class="n">separate_zero_wd_params_groups_for_optimizer</span><span class="p">(</span>
            <span class="n">net</span><span class="o">.</span><span class="n">module</span><span class="p">,</span> <span class="n">net_named_params</span><span class="p">,</span> <span class="n">training_params</span><span class="o">.</span><span class="n">optimizer_params</span><span class="p">[</span><span class="s1">&#39;weight_decay&#39;</span><span class="p">]</span>
        <span class="p">)</span>

    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># Overwrite groups to include params instead of named params</span>
        <span class="k">for</span> <span class="n">ind_group</span><span class="p">,</span> <span class="n">param_group</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">net_named_params</span><span class="p">):</span>
            <span class="n">param_group</span><span class="p">[</span><span class="s1">&#39;params&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="p">[</span><span class="n">param</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="k">for</span> <span class="n">param</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="n">param_group</span><span class="p">[</span><span class="s1">&#39;named_params&#39;</span><span class="p">])]</span>
            <span class="k">del</span> <span class="n">param_group</span><span class="p">[</span><span class="s1">&#39;named_params&#39;</span><span class="p">]</span>
            <span class="n">net_named_params</span><span class="p">[</span><span class="n">ind_group</span><span class="p">]</span> <span class="o">=</span> <span class="n">param_group</span>
        <span class="n">optimizer_training_params</span> <span class="o">=</span> <span class="n">net_named_params</span>

    <span class="c1"># CREATE AN OPTIMIZER OBJECT AND INITIALIZE IT</span>
    <span class="n">optimizer</span> <span class="o">=</span> <span class="n">optimizer_cls</span><span class="p">(</span><span class="n">optimizer_training_params</span><span class="p">,</span> <span class="n">lr</span><span class="o">=</span><span class="n">lr</span><span class="p">,</span> <span class="o">**</span><span class="n">training_params</span><span class="o">.</span><span class="n">optimizer_params</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">optimizer</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.optimizers.rmsprop_tf &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../../" id="documentation_options" src="../../../../../_static/documentation_options.js"></script>
        <script src="../../../../../_static/jquery.js"></script>
        <script src="../../../../../_static/underscore.js"></script>
        <script src="../../../../../_static/doctools.js"></script>
    <script src="../../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.optimizers.rmsprop_tf</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.optimizers.rmsprop_tf</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.optim</span> <span class="kn">import</span> <span class="n">Optimizer</span>

<span class="sd">&quot;&quot;&quot;</span>
<span class="sd">This implementation is taken from timm&#39;s github:</span>
<span class="sd">https://github.com/rwightman/pytorch-image-models/blob/master/timm/optim/rmsprop_tf.py</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="sd">&quot;&quot;&quot; RMSProp modified to behave like Tensorflow impl</span>
<span class="sd">Originally cut &amp; paste from PyTorch RMSProp</span>
<span class="sd">https://github.com/pytorch/pytorch/blob/063946d2b3f3f1e953a2a3b54e0b34f1393de295/torch/optim/rmsprop.py</span>
<span class="sd">Licensed under BSD-Clause 3 (ish), https://github.com/pytorch/pytorch/blob/master/LICENSE</span>
<span class="sd">Modifications Copyright 2020 Ross Wightman</span>
<span class="sd">&quot;&quot;&quot;</span>


<div class="viewcode-block" id="RMSpropTF"><a class="viewcode-back" href="../../../../../super_gradients.training.utils.optimizers.html#super_gradients.training.utils.optimizers.rmsprop_tf.RMSpropTF">[docs]</a><span class="k">class</span> <span class="nc">RMSpropTF</span><span class="p">(</span><span class="n">Optimizer</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Implements RMSprop algorithm (TensorFlow style epsilon)</span>
<span class="sd">    NOTE: This is a direct cut-and-paste of PyTorch RMSprop with eps applied before sqrt</span>
<span class="sd">    and a few other modifications to closer match Tensorflow for matching hyper-params.</span>
<span class="sd">    Noteworthy changes include:</span>
<span class="sd">    1. Epsilon applied inside square-root</span>
<span class="sd">    2. square_avg initialized to ones</span>
<span class="sd">    3. LR scaling of update accumulated in momentum buffer</span>
<span class="sd">    Proposed by G. Hinton in his</span>
<span class="sd">    `course &lt;http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf&gt;`_.</span>
<span class="sd">    The centered version first appears in `Generating Sequences</span>
<span class="sd">    With Recurrent Neural Networks &lt;https://arxiv.org/pdf/1308.0850v5.pdf&gt;`_.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</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">lr</span><span class="o">=</span><span class="mf">1e-2</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.9</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="mf">1e-10</span><span class="p">,</span> <span class="n">weight_decay</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">momentum</span><span class="o">=</span><span class="mf">0.</span><span class="p">,</span> <span class="n">centered</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                 <span class="n">decoupled_decay</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">lr_in_momentum</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;RMSprop optimizer that follows the tf&#39;s RMSprop characteristics</span>
<span class="sd">        :param params (iterable): iterable of parameters to optimize or dicts defining parameter groups</span>
<span class="sd">        :param lr (float, optional): learning rate</span>
<span class="sd">        :param momentum (float, optional): momentum factor</span>
<span class="sd">        :param alpha (float, optional): smoothing (decay) constant</span>
<span class="sd">        :param eps (float, optional): term added to the denominator to improve numerical stability</span>
<span class="sd">        :param centered (bool, optional) : if ``True``, compute the centered RMSProp, the gradient is normalized by an</span>
<span class="sd">         estimation of its variance</span>
<span class="sd">        :param weight_decay (float, optional): weight decay (L2 penalty)</span>
<span class="sd">        :param decoupled_decay (bool, optional): decoupled weight decay as per https://arxiv.org/abs/1711.05101</span>
<span class="sd">        :param lr_in_momentum (bool, optional): learning rate scaling is included in the momentum buffer update as per</span>
<span class="sd">         defaults in Tensorflow</span>
<span class="sd">    &quot;&quot;&quot;</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">lr</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid learning rate: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">lr</span><span class="p">))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">eps</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid epsilon value: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">eps</span><span class="p">))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">momentum</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid momentum value: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">momentum</span><span class="p">))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">weight_decay</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid weight_decay value: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">weight_decay</span><span class="p">))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="mf">0.0</span> <span class="o">&lt;=</span> <span class="n">alpha</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid alpha value: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">alpha</span><span class="p">))</span>

        <span class="n">defaults</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">lr</span><span class="o">=</span><span class="n">lr</span><span class="p">,</span> <span class="n">momentum</span><span class="o">=</span><span class="n">momentum</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span> <span class="n">centered</span><span class="o">=</span><span class="n">centered</span><span class="p">,</span> <span class="n">weight_decay</span><span class="o">=</span><span class="n">weight_decay</span><span class="p">,</span>
                        <span class="n">decoupled_decay</span><span class="o">=</span><span class="n">decoupled_decay</span><span class="p">,</span> <span class="n">lr_in_momentum</span><span class="o">=</span><span class="n">lr_in_momentum</span><span class="p">)</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">RMSpropTF</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><span class="n">params</span><span class="p">,</span> <span class="n">defaults</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">__setstate__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">state</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">RMSpropTF</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">__setstate__</span><span class="p">(</span><span class="n">state</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">group</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">param_groups</span><span class="p">:</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;momentum&#39;</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
            <span class="n">group</span><span class="o">.</span><span class="n">setdefault</span><span class="p">(</span><span class="s1">&#39;centered&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>

<div class="viewcode-block" id="RMSpropTF.step"><a class="viewcode-back" href="../../../../../super_gradients.training.utils.optimizers.html#super_gradients.training.utils.optimizers.rmsprop_tf.RMSpropTF.step">[docs]</a>    <span class="k">def</span> <span class="nf">step</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">closure</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: C901</span>
        <span class="sd">&quot;&quot;&quot;Performs a single optimization step.</span>
<span class="sd">        Arguments:</span>
<span class="sd">            closure (callable, optional): A closure that reevaluates the model</span>
<span class="sd">                and returns the loss.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="n">closure</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">loss</span> <span class="o">=</span> <span class="n">closure</span><span class="p">()</span>

        <span class="k">for</span> <span class="n">group</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">param_groups</span><span class="p">:</span>
            <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;params&#39;</span><span class="p">]:</span>
                <span class="k">if</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="k">continue</span>
                <span class="n">grad</span> <span class="o">=</span> <span class="n">p</span><span class="o">.</span><span class="n">grad</span><span class="o">.</span><span class="n">data</span>
                <span class="k">if</span> <span class="n">grad</span><span class="o">.</span><span class="n">is_sparse</span><span class="p">:</span>
                    <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s1">&#39;RMSprop does not support sparse gradients&#39;</span><span class="p">)</span>
                <span class="n">state</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">state</span><span class="p">[</span><span class="n">p</span><span class="p">]</span>

                <span class="c1"># State initialization</span>
                <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">state</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">state</span><span class="p">[</span><span class="s1">&#39;step&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
                    <span class="n">state</span><span class="p">[</span><span class="s1">&#39;square_avg&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones_like</span><span class="p">(</span><span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="p">)</span>  <span class="c1"># PyTorch inits to zero</span>
                    <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;momentum&#39;</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                        <span class="n">state</span><span class="p">[</span><span class="s1">&#39;momentum_buffer&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="p">)</span>
                    <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;centered&#39;</span><span class="p">]:</span>
                        <span class="n">state</span><span class="p">[</span><span class="s1">&#39;grad_avg&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="p">)</span>

                <span class="n">square_avg</span> <span class="o">=</span> <span class="n">state</span><span class="p">[</span><span class="s1">&#39;square_avg&#39;</span><span class="p">]</span>
                <span class="n">one_minus_alpha</span> <span class="o">=</span> <span class="mf">1.</span> <span class="o">-</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;alpha&#39;</span><span class="p">]</span>

                <span class="n">state</span><span class="p">[</span><span class="s1">&#39;step&#39;</span><span class="p">]</span> <span class="o">+=</span> <span class="mi">1</span>

                <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;weight_decay&#39;</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="k">if</span> <span class="s1">&#39;decoupled_decay&#39;</span> <span class="ow">in</span> <span class="n">group</span> <span class="ow">and</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;decoupled_decay&#39;</span><span class="p">]:</span>
                        <span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="o">-</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;weight_decay&#39;</span><span class="p">],</span> <span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="p">)</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="n">grad</span> <span class="o">=</span> <span class="n">grad</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;weight_decay&#39;</span><span class="p">],</span> <span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="p">)</span>

                <span class="c1"># Tensorflow order of ops for updating squared avg</span>
                <span class="n">square_avg</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">one_minus_alpha</span><span class="p">,</span> <span class="n">grad</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span> <span class="o">-</span> <span class="n">square_avg</span><span class="p">)</span>
                <span class="c1"># square_avg.mul_(alpha).addcmul_(1 - alpha, grad, grad)  # PyTorch original</span>

                <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;centered&#39;</span><span class="p">]:</span>
                    <span class="n">grad_avg</span> <span class="o">=</span> <span class="n">state</span><span class="p">[</span><span class="s1">&#39;grad_avg&#39;</span><span class="p">]</span>
                    <span class="n">grad_avg</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">one_minus_alpha</span><span class="p">,</span> <span class="n">grad</span> <span class="o">-</span> <span class="n">grad_avg</span><span class="p">)</span>
                    <span class="c1"># grad_avg.mul_(alpha).add_(1 - alpha, grad)  # PyTorch original</span>
                    <span class="n">avg</span> <span class="o">=</span> <span class="n">square_avg</span><span class="o">.</span><span class="n">addcmul</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">grad_avg</span><span class="p">,</span> <span class="n">grad_avg</span><span class="p">)</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;eps&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">sqrt_</span><span class="p">()</span>  <span class="c1"># eps moved in sqrt</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">avg</span> <span class="o">=</span> <span class="n">square_avg</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;eps&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">sqrt_</span><span class="p">()</span>  <span class="c1"># eps moved in sqrt</span>

                <span class="k">if</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;momentum&#39;</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">buf</span> <span class="o">=</span> <span class="n">state</span><span class="p">[</span><span class="s1">&#39;momentum_buffer&#39;</span><span class="p">]</span>
                    <span class="c1"># Tensorflow accumulates the LR scaling in the momentum buffer</span>
                    <span class="k">if</span> <span class="s1">&#39;lr_in_momentum&#39;</span> <span class="ow">in</span> <span class="n">group</span> <span class="ow">and</span> <span class="n">group</span><span class="p">[</span><span class="s1">&#39;lr_in_momentum&#39;</span><span class="p">]:</span>
                        <span class="n">buf</span><span class="o">.</span><span class="n">mul_</span><span class="p">(</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;momentum&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">addcdiv_</span><span class="p">(</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;lr&#39;</span><span class="p">],</span> <span class="n">grad</span><span class="p">,</span> <span class="n">avg</span><span class="p">)</span>
                        <span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="o">-</span><span class="n">buf</span><span class="p">)</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="c1"># PyTorch scales the param update by LR</span>
                        <span class="n">buf</span><span class="o">.</span><span class="n">mul_</span><span class="p">(</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;momentum&#39;</span><span class="p">])</span><span class="o">.</span><span class="n">addcdiv_</span><span class="p">(</span><span class="n">grad</span><span class="p">,</span> <span class="n">avg</span><span class="p">)</span>
                        <span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="o">-</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;lr&#39;</span><span class="p">],</span> <span class="n">buf</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">p</span><span class="o">.</span><span class="n">data</span><span class="o">.</span><span class="n">addcdiv_</span><span class="p">(</span><span class="o">-</span><span class="n">group</span><span class="p">[</span><span class="s1">&#39;lr&#39;</span><span class="p">],</span> <span class="n">grad</span><span class="p">,</span> <span class="n">avg</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">loss</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.regularization_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.regularization_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.regularization_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">nn</span>


<div class="viewcode-block" id="DropPath"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.regularization_utils.DropPath">[docs]</a><span class="k">class</span> <span class="nc">DropPath</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).</span>

<span class="sd">    Code taken from TIMM (https://github.com/rwightman/pytorch-image-models)</span>
<span class="sd">    Apache License 2.0</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">drop_prob</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DropPath</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">drop_prob</span> <span class="o">=</span> <span class="n">drop_prob</span>

<div class="viewcode-block" id="DropPath.forward"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.regularization_utils.DropPath.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">drop_prob</span> <span class="o">==</span> <span class="mf">0.</span> <span class="ow">or</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">training</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">x</span>

        <span class="n">keep_prob</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">drop_prob</span>
        <span class="n">shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],)</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span> <span class="o">*</span> <span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">ndim</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>  <span class="c1"># work with diff dim tensors, not just 2D ConvNets</span>
        <span class="n">random_tensor</span> <span class="o">=</span> <span class="n">keep_prob</span> <span class="o">+</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">x</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">x</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
        <span class="n">random_tensor</span><span class="o">.</span><span class="n">floor_</span><span class="p">()</span>  <span class="c1"># binarize</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">div</span><span class="p">(</span><span class="n">keep_prob</span><span class="p">)</span> <span class="o">*</span> <span class="n">random_tensor</span>
        <span class="k">return</span> <span class="n">output</span></div></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.segmentation_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.segmentation_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.segmentation_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">cv2</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Callable</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">torch.nn.functional</span> <span class="k">as</span> <span class="nn">F</span>
<span class="kn">from</span> <span class="nn">torchvision.utils</span> <span class="kn">import</span> <span class="n">draw_segmentation_masks</span>

<span class="c1"># FIXME: REFACTOR AUGMENTATIONS, CONSIDER USING A MORE EFFICIENT LIBRARIES SUCH AS, IMGAUG, DALI ETC.</span>
<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span>


<div class="viewcode-block" id="coco_sub_classes_inclusion_tuples_list"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.coco_sub_classes_inclusion_tuples_list">[docs]</a><span class="k">def</span> <span class="nf">coco_sub_classes_inclusion_tuples_list</span><span class="p">():</span>
    <span class="k">return</span> <span class="p">[(</span><span class="mi">0</span><span class="p">,</span> <span class="s1">&#39;background&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">5</span><span class="p">,</span> <span class="s1">&#39;airplane&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="s1">&#39;bicycle&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">16</span><span class="p">,</span> <span class="s1">&#39;bird&#39;</span><span class="p">),</span>
            <span class="p">(</span><span class="mi">9</span><span class="p">,</span> <span class="s1">&#39;boat&#39;</span><span class="p">),</span>
            <span class="p">(</span><span class="mi">44</span><span class="p">,</span> <span class="s1">&#39;bottle&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">6</span><span class="p">,</span> <span class="s1">&#39;bus&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="s1">&#39;car&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">17</span><span class="p">,</span> <span class="s1">&#39;cat&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">62</span><span class="p">,</span> <span class="s1">&#39;chair&#39;</span><span class="p">),</span>
            <span class="p">(</span><span class="mi">21</span><span class="p">,</span> <span class="s1">&#39;cow&#39;</span><span class="p">),</span>
            <span class="p">(</span><span class="mi">67</span><span class="p">,</span> <span class="s1">&#39;dining table&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">18</span><span class="p">,</span> <span class="s1">&#39;dog&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">19</span><span class="p">,</span> <span class="s1">&#39;horse&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="s1">&#39;motorcycle&#39;</span><span class="p">),</span>
            <span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="s1">&#39;person&#39;</span><span class="p">),</span>
            <span class="p">(</span><span class="mi">64</span><span class="p">,</span> <span class="s1">&#39;potted plant&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">20</span><span class="p">,</span> <span class="s1">&#39;sheep&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">63</span><span class="p">,</span> <span class="s1">&#39;couch&#39;</span><span class="p">),</span> <span class="p">(</span><span class="mi">7</span><span class="p">,</span> <span class="s1">&#39;train&#39;</span><span class="p">),</span>
            <span class="p">(</span><span class="mi">72</span><span class="p">,</span> <span class="s1">&#39;tv&#39;</span><span class="p">)]</span></div>


<div class="viewcode-block" id="to_one_hot"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.to_one_hot">[docs]</a><span class="k">def</span> <span class="nf">to_one_hot</span><span class="p">(</span><span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">ignore_index</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Target label to one_hot tensor. labels and ignore_index must be consecutive numbers.</span>
<span class="sd">    :param target: Class labels long tensor, with shape [N, H, W]</span>
<span class="sd">    :param num_classes: num of classes in datasets excluding ignore label, this is the output channels of the one hot</span>
<span class="sd">        result.</span>
<span class="sd">    :return: one hot tensor with shape [N, num_classes, H, W]</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">num_classes</span> <span class="o">=</span> <span class="n">num_classes</span> <span class="k">if</span> <span class="n">ignore_index</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">num_classes</span> <span class="o">+</span> <span class="mi">1</span>

    <span class="n">one_hot</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">one_hot</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">)</span><span class="o">.</span><span class="n">permute</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="mi">3</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="k">if</span> <span class="n">ignore_index</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># remove ignore_index channel</span>
        <span class="n">one_hot</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="n">one_hot</span><span class="p">[:,</span> <span class="p">:</span><span class="n">ignore_index</span><span class="p">],</span> <span class="n">one_hot</span><span class="p">[:,</span> <span class="n">ignore_index</span> <span class="o">+</span> <span class="mi">1</span><span class="p">:]],</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">one_hot</span></div>


<div class="viewcode-block" id="reverse_imagenet_preprocessing"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.reverse_imagenet_preprocessing">[docs]</a><span class="k">def</span> <span class="nf">reverse_imagenet_preprocessing</span><span class="p">(</span><span class="n">im_tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    :param im_tensor: images in a batch after preprocessing for inference, RGB, (B, C, H, W)</span>
<span class="sd">    :return:          images in a batch in cv2 format, BGR, (B, H, W, C)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">im_np</span> <span class="o">=</span> <span class="n">im_tensor</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>
    <span class="n">im_np</span> <span class="o">=</span> <span class="n">im_np</span><span class="p">[:,</span> <span class="p">::</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:]</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">im_np</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">.229</span><span class="p">,</span> <span class="mf">.224</span><span class="p">,</span> <span class="mf">.225</span><span class="p">][::</span><span class="o">-</span><span class="mi">1</span><span class="p">]]])</span>
    <span class="n">im_np</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">.485</span><span class="p">,</span> <span class="mf">.456</span><span class="p">,</span> <span class="mf">.406</span><span class="p">][::</span><span class="o">-</span><span class="mi">1</span><span class="p">]]])</span>
    <span class="n">im_np</span> <span class="o">*=</span> <span class="mf">255.</span>
    <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">ascontiguousarray</span><span class="p">(</span><span class="n">im_np</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span></div>


<div class="viewcode-block" id="BinarySegmentationVisualization"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.BinarySegmentationVisualization">[docs]</a><span class="k">class</span> <span class="nc">BinarySegmentationVisualization</span><span class="p">:</span>

    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">_visualize_image</span><span class="p">(</span><span class="n">image_np</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">pred_mask</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">target_mask</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                         <span class="n">image_scale</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">checkpoint_dir</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">image_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
        <span class="n">pred_mask</span> <span class="o">=</span> <span class="n">pred_mask</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="n">image_np</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">from_numpy</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">moveaxis</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">))</span>

        <span class="n">pred_mask</span> <span class="o">=</span> <span class="n">pred_mask</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">newaxis</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:]</span> <span class="o">&gt;</span> <span class="mf">0.5</span>
        <span class="n">target_mask</span> <span class="o">=</span> <span class="n">target_mask</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">newaxis</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:]</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">bool</span><span class="p">)</span>
        <span class="n">tp_mask</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="n">pred_mask</span><span class="p">,</span> <span class="n">target_mask</span><span class="p">)</span>
        <span class="n">fp_mask</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="n">pred_mask</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">logical_not</span><span class="p">(</span><span class="n">target_mask</span><span class="p">))</span>
        <span class="n">fn_mask</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">logical_and</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">logical_not</span><span class="p">(</span><span class="n">pred_mask</span><span class="p">),</span> <span class="n">target_mask</span><span class="p">)</span>
        <span class="n">overlay</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">from_numpy</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">tp_mask</span><span class="p">,</span> <span class="n">fp_mask</span><span class="p">,</span> <span class="n">fn_mask</span><span class="p">]))</span>

        <span class="c1"># SWITCH BETWEEN BLUE AND RED IF WE SAVE THE IMAGE ON THE DISC AS OTHERWISE WE CHANGE CHANNEL ORDERING</span>
        <span class="n">colors</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;green&#39;</span><span class="p">,</span> <span class="s1">&#39;red&#39;</span><span class="p">,</span> <span class="s1">&#39;blue&#39;</span><span class="p">]</span>
        <span class="n">res_image</span> <span class="o">=</span> <span class="n">draw_segmentation_masks</span><span class="p">(</span><span class="n">image_np</span><span class="p">,</span> <span class="n">overlay</span><span class="p">,</span> <span class="n">colors</span><span class="o">=</span><span class="n">colors</span><span class="p">)</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>
        <span class="n">res_image</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">([</span><span class="n">res_image</span><span class="p">[</span><span class="n">ch</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:,</span> <span class="n">np</span><span class="o">.</span><span class="n">newaxis</span><span class="p">]</span> <span class="k">for</span> <span class="n">ch</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">3</span><span class="p">)],</span> <span class="mi">2</span><span class="p">)</span>
        <span class="n">res_image</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">res_image</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">),</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">fx</span><span class="o">=</span><span class="n">image_scale</span><span class="p">,</span> <span class="n">fy</span><span class="o">=</span><span class="n">image_scale</span><span class="p">,</span>
                               <span class="n">interpolation</span><span class="o">=</span><span class="n">cv2</span><span class="o">.</span><span class="n">INTER_NEAREST</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">checkpoint_dir</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">res_image</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">cv2</span><span class="o">.</span><span class="n">imwrite</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">checkpoint_dir</span><span class="p">,</span> <span class="nb">str</span><span class="p">(</span><span class="n">image_name</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;.jpg&#39;</span><span class="p">),</span> <span class="n">res_image</span><span class="p">)</span>

<div class="viewcode-block" id="BinarySegmentationVisualization.visualize_batch"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.BinarySegmentationVisualization.visualize_batch">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">visualize_batch</span><span class="p">(</span><span class="n">image_tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">pred_mask</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">target_mask</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                        <span class="n">batch_name</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">str</span><span class="p">],</span> <span class="n">checkpoint_dir</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                        <span class="n">undo_preprocessing_func</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="n">reverse_imagenet_preprocessing</span><span class="p">,</span>
                        <span class="n">image_scale</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        A helper function to visualize detections predicted by a network:</span>
<span class="sd">        saves images into a given path with a name that is {batch_name}_{imade_idx_in_the_batch}.jpg, one batch per call.</span>
<span class="sd">        Colors are generated on the fly: uniformly sampled from color wheel to support all given classes.</span>

<span class="sd">        :param image_tensor:            rgb images, (B, H, W, 3)</span>
<span class="sd">        :param pred_boxes:              boxes after NMS for each image in a batch, each (Num_boxes, 6),</span>
<span class="sd">                                        values on dim 1 are: x1, y1, x2, y2, confidence, class</span>
<span class="sd">        :param target_boxes:            (Num_targets, 6), values on dim 1 are: image id in a batch, class, x y w h</span>
<span class="sd">                                        (coordinates scaled to [0, 1])</span>
<span class="sd">        :param batch_name:              id of the current batch to use for image naming</span>

<span class="sd">        :param checkpoint_dir:          a path where images with boxes will be saved. if None, the result images will</span>
<span class="sd">                                        be returns as a list of numpy image arrays</span>

<span class="sd">        :param undo_preprocessing_func: a function to convert preprocessed images tensor into a batch of cv2-like images</span>
<span class="sd">        :param image_scale:             scale factor for output image</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">image_np</span> <span class="o">=</span> <span class="n">undo_preprocessing_func</span><span class="p">(</span><span class="n">image_tensor</span><span class="o">.</span><span class="n">detach</span><span class="p">())</span>
        <span class="n">pred_mask</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">(</span><span class="n">pred_mask</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="p">:])</span>  <span class="c1"># comment out</span>

        <span class="n">out_images</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">image_np</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]):</span>
            <span class="n">preds</span> <span class="o">=</span> <span class="n">pred_mask</span><span class="p">[</span><span class="n">i</span><span class="p">]</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>
            <span class="n">targets</span> <span class="o">=</span> <span class="n">target_mask</span><span class="p">[</span><span class="n">i</span><span class="p">]</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>

            <span class="n">image_name</span> <span class="o">=</span> <span class="s1">&#39;_&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="nb">str</span><span class="p">(</span><span class="n">batch_name</span><span class="p">),</span> <span class="nb">str</span><span class="p">(</span><span class="n">i</span><span class="p">)])</span>
            <span class="n">res_image</span> <span class="o">=</span> <span class="n">BinarySegmentationVisualization</span><span class="o">.</span><span class="n">_visualize_image</span><span class="p">(</span><span class="n">image_np</span><span class="p">[</span><span class="n">i</span><span class="p">],</span> <span class="n">preds</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">image_scale</span><span class="p">,</span>
                                                                         <span class="n">checkpoint_dir</span><span class="p">,</span> <span class="n">image_name</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">res_image</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">out_images</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">res_image</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">out_images</span></div></div>


<div class="viewcode-block" id="visualize_batches"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.visualize_batches">[docs]</a><span class="k">def</span> <span class="nf">visualize_batches</span><span class="p">(</span><span class="n">dataloader</span><span class="p">,</span> <span class="n">module</span><span class="p">,</span> <span class="n">visualization_path</span><span class="p">,</span> <span class="n">num_batches</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">undo_preprocessing_func</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="n">os</span><span class="o">.</span><span class="n">makedirs</span><span class="p">(</span><span class="n">visualization_path</span><span class="p">,</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">batch_i</span><span class="p">,</span> <span class="p">(</span><span class="n">imgs</span><span class="p">,</span> <span class="n">targets</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">dataloader</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">batch_i</span> <span class="o">==</span> <span class="n">num_batches</span><span class="p">:</span>
            <span class="k">return</span>
        <span class="n">imgs</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">tensor_container_to_device</span><span class="p">(</span><span class="n">imgs</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">(</span><span class="s1">&#39;cuda:0&#39;</span><span class="p">))</span>
        <span class="n">targets</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">tensor_container_to_device</span><span class="p">(</span><span class="n">targets</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">(</span><span class="s1">&#39;cuda:0&#39;</span><span class="p">))</span>
        <span class="n">pred_mask</span> <span class="o">=</span> <span class="n">module</span><span class="p">(</span><span class="n">imgs</span><span class="p">)</span>

        <span class="c1"># Visualize the batch</span>
        <span class="k">if</span> <span class="n">undo_preprocessing_func</span><span class="p">:</span>
            <span class="n">BinarySegmentationVisualization</span><span class="o">.</span><span class="n">visualize_batch</span><span class="p">(</span><span class="n">imgs</span><span class="p">,</span> <span class="n">pred_mask</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">batch_i</span><span class="p">,</span> <span class="n">visualization_path</span><span class="p">,</span>
                                                            <span class="n">undo_preprocessing_func</span><span class="o">=</span><span class="n">undo_preprocessing_func</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">BinarySegmentationVisualization</span><span class="o">.</span><span class="n">visualize_batch</span><span class="p">(</span><span class="n">imgs</span><span class="p">,</span> <span class="n">pred_mask</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">batch_i</span><span class="p">,</span> <span class="n">visualization_path</span><span class="p">)</span></div>


<div class="viewcode-block" id="one_hot_to_binary_edge"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.one_hot_to_binary_edge">[docs]</a><span class="k">def</span> <span class="nf">one_hot_to_binary_edge</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                           <span class="n">kernel_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                           <span class="n">flatten_channels</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Utils function to create edge feature maps.</span>
<span class="sd">    :param x: input tensor, must be one_hot tensor with shape [B, C, H, W]</span>
<span class="sd">    :param kernel_size: kernel size of dilation erosion convolutions. The result edge widths depends on this argument as</span>
<span class="sd">        follows: `edge_width = kernel - 1`</span>
<span class="sd">    :param flatten_channels: Whether to apply logical_or across channels dimension, if at least one pixel class is</span>
<span class="sd">        considered as edge pixel flatten value is 1. If set as `False` the output tensor shape is [B, C, H, W], else</span>
<span class="sd">        [B, 1, H, W]. Default is `True`.</span>
<span class="sd">    :return: one_hot edge torch.Tensor.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">kernel_size</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">kernel_size</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;kernel size must be an odd positive values, such as [1, 3, 5, ..], found: </span><span class="si">{</span><span class="n">kernel_size</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="n">_kernel</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="mi">1</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">kernel_size</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">float32</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">x</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
    <span class="n">padding</span> <span class="o">=</span> <span class="p">(</span><span class="n">kernel_size</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="o">//</span> <span class="mi">2</span>
    <span class="c1"># Use replicate padding to prevent class shifting and edge formation at the image boundaries.</span>
    <span class="n">padded_x</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">pad</span><span class="p">(</span><span class="n">x</span><span class="o">.</span><span class="n">float</span><span class="p">(),</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;replicate&quot;</span><span class="p">,</span> <span class="n">pad</span><span class="o">=</span><span class="p">[</span><span class="n">padding</span><span class="p">]</span> <span class="o">*</span> <span class="mi">4</span><span class="p">)</span>
    <span class="c1"># The binary edges feature map is created by subtracting dilated features from erosed features.</span>
    <span class="c1"># First the positive one value masks are expanded (dilation) by applying a sliding window filter of one values.</span>
    <span class="c1"># The resulted output is then clamped to binary format to [0, 1], this way the one-hot boundaries are expanded by</span>
    <span class="c1"># (kernel_size - 1) / 2.</span>
    <span class="n">dilation</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span>
        <span class="n">F</span><span class="o">.</span><span class="n">conv2d</span><span class="p">(</span><span class="n">padded_x</span><span class="p">,</span> <span class="n">_kernel</span><span class="p">,</span> <span class="n">groups</span><span class="o">=</span><span class="n">x</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">)),</span>
        <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span>
    <span class="p">)</span>
    <span class="c1"># Similar to dilation, erosion (can be seen as inverse of dilation) is applied to contract the one-hot features by</span>
    <span class="c1"># applying a dilation operation on the inverse of the one-hot features.</span>
    <span class="n">erosion</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">torch</span><span class="o">.</span><span class="n">clamp</span><span class="p">(</span>
        <span class="n">F</span><span class="o">.</span><span class="n">conv2d</span><span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">padded_x</span><span class="p">,</span> <span class="n">_kernel</span><span class="p">,</span> <span class="n">groups</span><span class="o">=</span><span class="n">x</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">)),</span>
        <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span>
    <span class="p">)</span>
    <span class="c1"># Finally the edge features are the result of subtracting dilation by erosion.</span>
    <span class="c1"># i.e for a simple 1D one-hot input:    [0, 0, 0, 1, 1, 1, 0, 0, 0], using sliding kernel with size 3: [1, 1, 1]</span>
    <span class="c1"># Dilated features:                     [0, 0, 1, 1, 1, 1, 1, 0, 0]</span>
    <span class="c1"># Erosed inverse features:              [0, 0, 0, 0, 1, 0, 0, 0, 0]</span>
    <span class="c1"># Edge features: dilation - erosion:    [0, 0, 1, 1, 0, 1, 1, 0, 0]</span>
    <span class="n">edge</span> <span class="o">=</span> <span class="n">dilation</span> <span class="o">-</span> <span class="n">erosion</span>
    <span class="k">if</span> <span class="n">flatten_channels</span><span class="p">:</span>
        <span class="c1"># use max operator across channels. Equivalent to logical or for input with binary values [0, 1].</span>
        <span class="n">edge</span> <span class="o">=</span> <span class="n">edge</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">edge</span></div>


<div class="viewcode-block" id="target_to_binary_edge"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.segmentation_utils.target_to_binary_edge">[docs]</a><span class="k">def</span> <span class="nf">target_to_binary_edge</span><span class="p">(</span><span class="n">target</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
                          <span class="n">num_classes</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                          <span class="n">kernel_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                          <span class="n">ignore_index</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                          <span class="n">flatten_channels</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Utils function to create edge feature maps from target.</span>
<span class="sd">    :param target: Class labels long tensor, with shape [N, H, W]</span>
<span class="sd">    :param num_classes: num of classes in datasets excluding ignore label, this is the output channels of the one hot</span>
<span class="sd">        result.</span>
<span class="sd">    :param kernel_size: kernel size of dilation erosion convolutions. The result edge widths depends on this argument as</span>
<span class="sd">        follows: `edge_width = kernel - 1`</span>
<span class="sd">    :param flatten_channels: Whether to apply logical or across channels dimension, if at least one pixel class is</span>
<span class="sd">        considered as edge pixel flatten value is 1. If set as `False` the output tensor shape is [B, C, H, W], else</span>
<span class="sd">        [B, 1, H, W]. Default is `True`.</span>
<span class="sd">    :return: one_hot edge torch.Tensor.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">one_hot</span> <span class="o">=</span> <span class="n">to_one_hot</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="n">num_classes</span><span class="o">=</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">one_hot_to_binary_edge</span><span class="p">(</span><span class="n">one_hot</span><span class="p">,</span> <span class="n">kernel_size</span><span class="o">=</span><span class="n">kernel_size</span><span class="p">,</span> <span class="n">flatten_channels</span><span class="o">=</span><span class="n">flatten_channels</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.sg_model_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.sg_model_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.sg_model_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">sys</span>
<span class="kn">import</span> <span class="nn">socket</span>
<span class="kn">import</span> <span class="nn">time</span>
<span class="kn">from</span> <span class="nn">dataclasses</span> <span class="kn">import</span> <span class="n">dataclass</span>
<span class="kn">from</span> <span class="nn">multiprocessing</span> <span class="kn">import</span> <span class="n">Process</span>
<span class="kn">from</span> <span class="nn">pathlib</span> <span class="kn">import</span> <span class="n">Path</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Dict</span>
<span class="kn">import</span> <span class="nn">random</span>

<span class="kn">from</span> <span class="nn">treelib</span> <span class="kn">import</span> <span class="n">Tree</span>
<span class="kn">from</span> <span class="nn">termcolor</span> <span class="kn">import</span> <span class="n">colored</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.utils.tensorboard</span> <span class="kn">import</span> <span class="n">SummaryWriter</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.exceptions.dataset_exceptions</span> <span class="kn">import</span> <span class="n">UnsupportedBatchItemsFormat</span>


<span class="c1"># TODO: These utils should move to sg_model package as internal (private) helper functions</span>

<span class="n">IS_BETTER_COLOR</span> <span class="o">=</span> <span class="p">{</span><span class="kc">True</span><span class="p">:</span> <span class="s2">&quot;green&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">:</span> <span class="s2">&quot;red&quot;</span><span class="p">}</span>
<span class="n">IS_GREATER_SYMBOLS</span> <span class="o">=</span> <span class="p">{</span><span class="kc">True</span><span class="p">:</span> <span class="s2">&quot;↗&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">:</span> <span class="s2">&quot;↘&quot;</span><span class="p">}</span>


<div class="viewcode-block" id="MonitoredValue"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.MonitoredValue">[docs]</a><span class="nd">@dataclass</span>
<span class="k">class</span> <span class="nc">MonitoredValue</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Store a value and some indicators relative to its past iterations.</span>

<span class="sd">    The value can be a metric/loss, and the iteration can be epochs/batch.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">name</span><span class="p">:</span> <span class="nb">str</span>
    <span class="n">greater_is_better</span><span class="p">:</span> <span class="nb">bool</span>
    <span class="n">current</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">previous</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">best</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">change_from_previous</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">change_from_best</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">is_better_than_previous</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">change_from_best</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">None</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">change_from_previous</span> <span class="o">&gt;=</span> <span class="mi">0</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">change_from_previous</span> <span class="o">&lt;</span> <span class="mi">0</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">is_best_value</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">change_from_best</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">None</span>
        <span class="k">elif</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">change_from_best</span> <span class="o">&gt;=</span> <span class="mi">0</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">change_from_best</span> <span class="o">&lt;</span> <span class="mi">0</span></div>


<div class="viewcode-block" id="update_monitored_value"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.update_monitored_value">[docs]</a><span class="k">def</span> <span class="nf">update_monitored_value</span><span class="p">(</span><span class="n">previous_monitored_value</span><span class="p">:</span> <span class="n">MonitoredValue</span><span class="p">,</span> <span class="n">new_value</span><span class="p">:</span> <span class="nb">float</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">MonitoredValue</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Update the given ValueToMonitor object (could be a loss or a metric) with the new value</span>

<span class="sd">    :param previous_monitored_value: The stats about the value that is monitored throughout epochs.</span>
<span class="sd">    :param new_value: The value of the current epoch that will be used to update previous_monitored_value</span>
<span class="sd">    :return:</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">previous_value</span><span class="p">,</span> <span class="n">previous_best_value</span> <span class="o">=</span> <span class="n">previous_monitored_value</span><span class="o">.</span><span class="n">current</span><span class="p">,</span> <span class="n">previous_monitored_value</span><span class="o">.</span><span class="n">best</span>
    <span class="n">name</span><span class="p">,</span> <span class="n">greater_is_better</span> <span class="o">=</span> <span class="n">previous_monitored_value</span><span class="o">.</span><span class="n">name</span><span class="p">,</span> <span class="n">previous_monitored_value</span><span class="o">.</span><span class="n">greater_is_better</span>

    <span class="k">if</span> <span class="n">previous_best_value</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">previous_best_value</span> <span class="o">=</span> <span class="n">previous_value</span>
    <span class="k">elif</span> <span class="n">greater_is_better</span><span class="p">:</span>
        <span class="n">previous_best_value</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">previous_value</span><span class="p">,</span> <span class="n">previous_best_value</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">previous_best_value</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">previous_value</span><span class="p">,</span> <span class="n">previous_best_value</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">previous_value</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">change_from_previous</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">change_from_best</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">change_from_previous</span> <span class="o">=</span> <span class="n">new_value</span> <span class="o">-</span> <span class="n">previous_value</span>
        <span class="n">change_from_best</span> <span class="o">=</span> <span class="n">new_value</span> <span class="o">-</span> <span class="n">previous_best_value</span>

    <span class="k">return</span> <span class="n">MonitoredValue</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">,</span> <span class="n">current</span><span class="o">=</span><span class="n">new_value</span><span class="p">,</span> <span class="n">previous</span><span class="o">=</span><span class="n">previous_value</span><span class="p">,</span> <span class="n">best</span><span class="o">=</span><span class="n">previous_best_value</span><span class="p">,</span>
                          <span class="n">change_from_previous</span><span class="o">=</span><span class="n">change_from_previous</span><span class="p">,</span> <span class="n">change_from_best</span><span class="o">=</span><span class="n">change_from_best</span><span class="p">,</span>
                          <span class="n">greater_is_better</span><span class="o">=</span><span class="n">greater_is_better</span><span class="p">)</span></div>


<div class="viewcode-block" id="update_monitored_values_dict"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.update_monitored_values_dict">[docs]</a><span class="k">def</span> <span class="nf">update_monitored_values_dict</span><span class="p">(</span><span class="n">monitored_values_dict</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">MonitoredValue</span><span class="p">],</span>
                                 <span class="n">new_values_dict</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="nb">float</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">MonitoredValue</span><span class="p">]:</span>
    <span class="sd">&quot;&quot;&quot;Update the given ValueToMonitor object (could be a loss or a metric) with the new value</span>

<span class="sd">    :param monitored_values_dict: Dict mapping value names to their stats throughout epochs.</span>
<span class="sd">    :param new_values_dict: Dict mapping value names to their new (i.e. current epoch) value.</span>
<span class="sd">    :return: Updated monitored_values_dict</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">for</span> <span class="n">monitored_value_name</span> <span class="ow">in</span> <span class="n">monitored_values_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
        <span class="n">monitored_values_dict</span><span class="p">[</span><span class="n">monitored_value_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">update_monitored_value</span><span class="p">(</span>
            <span class="n">new_value</span><span class="o">=</span><span class="n">new_values_dict</span><span class="p">[</span><span class="n">monitored_value_name</span><span class="p">],</span>
            <span class="n">previous_monitored_value</span><span class="o">=</span><span class="n">monitored_values_dict</span><span class="p">[</span><span class="n">monitored_value_name</span><span class="p">],</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">monitored_values_dict</span></div>


<div class="viewcode-block" id="display_epoch_summary"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.display_epoch_summary">[docs]</a><span class="k">def</span> <span class="nf">display_epoch_summary</span><span class="p">(</span><span class="n">epoch</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">n_digits</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
                          <span class="n">train_monitored_values</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">MonitoredValue</span><span class="p">],</span>
                          <span class="n">valid_monitored_values</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">MonitoredValue</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;Display a summary of loss/metric of interest, for a given epoch.</span>

<span class="sd">        :param epoch: the number of epoch.</span>
<span class="sd">        :param n_digits: number of digits to display on screen for float values</span>
<span class="sd">        :param train_monitored_values: mapping of loss/metric with their stats that will be displayed</span>
<span class="sd">        :param valid_monitored_values: mapping of loss/metric with their stats that will be displayed</span>
<span class="sd">        :return:</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">_format_to_str</span><span class="p">(</span><span class="n">val</span><span class="p">:</span> <span class="nb">float</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
        <span class="k">return</span> <span class="nb">str</span><span class="p">(</span><span class="nb">round</span><span class="p">(</span><span class="n">val</span><span class="p">,</span> <span class="n">n_digits</span><span class="p">))</span>

    <span class="k">def</span> <span class="nf">_generate_tree</span><span class="p">(</span><span class="n">value_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">monitored_value</span><span class="p">:</span> <span class="n">MonitoredValue</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tree</span><span class="p">:</span>
        <span class="sd">&quot;&quot;&quot;Generate a tree that represents the stats of a given loss/metric.&quot;&quot;&quot;</span>

        <span class="n">current</span> <span class="o">=</span> <span class="n">_format_to_str</span><span class="p">(</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">current</span><span class="p">)</span>
        <span class="n">root_id</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="nb">hash</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">value_name</span><span class="si">}</span><span class="s2"> = </span><span class="si">{</span><span class="n">current</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">))</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">random</span><span class="o">.</span><span class="n">random</span><span class="p">())</span>

        <span class="n">tree</span> <span class="o">=</span> <span class="n">Tree</span><span class="p">()</span>
        <span class="n">tree</span><span class="o">.</span><span class="n">create_node</span><span class="p">(</span><span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">value_name</span><span class="o">.</span><span class="n">capitalize</span><span class="p">()</span><span class="si">}</span><span class="s2"> = </span><span class="si">{</span><span class="n">current</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">identifier</span><span class="o">=</span><span class="n">root_id</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">monitored_value</span><span class="o">.</span><span class="n">previous</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">previous</span> <span class="o">=</span> <span class="n">_format_to_str</span><span class="p">(</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">previous</span><span class="p">)</span>
            <span class="n">best</span> <span class="o">=</span> <span class="n">_format_to_str</span><span class="p">(</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">best</span><span class="p">)</span>
            <span class="n">change_from_previous</span> <span class="o">=</span> <span class="n">_format_to_str</span><span class="p">(</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">change_from_previous</span><span class="p">)</span>
            <span class="n">change_from_best</span> <span class="o">=</span> <span class="n">_format_to_str</span><span class="p">(</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">change_from_best</span><span class="p">)</span>

            <span class="n">diff_with_prev_colored</span> <span class="o">=</span> <span class="n">colored</span><span class="p">(</span>
                <span class="n">text</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">IS_GREATER_SYMBOLS</span><span class="p">[</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">change_from_previous</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2"> </span><span class="si">{</span><span class="n">change_from_previous</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">IS_BETTER_COLOR</span><span class="p">[</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">is_better_than_previous</span><span class="p">]</span>
            <span class="p">)</span>
            <span class="n">diff_with_best_colored</span> <span class="o">=</span> <span class="n">colored</span><span class="p">(</span>
                <span class="n">text</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">IS_GREATER_SYMBOLS</span><span class="p">[</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">change_from_best</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2"> </span><span class="si">{</span><span class="n">change_from_best</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">IS_BETTER_COLOR</span><span class="p">[</span><span class="n">monitored_value</span><span class="o">.</span><span class="n">is_best_value</span><span class="p">]</span>
            <span class="p">)</span>

            <span class="n">tree</span><span class="o">.</span><span class="n">create_node</span><span class="p">(</span>
                <span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;Epoch N-1      = </span><span class="si">{</span><span class="n">previous</span><span class="si">:</span><span class="s2">6</span><span class="si">}</span><span class="s2"> (</span><span class="si">{</span><span class="n">diff_with_prev_colored</span><span class="si">:</span><span class="s2">8</span><span class="si">}</span><span class="s2">)&quot;</span><span class="p">,</span>
                <span class="n">identifier</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;0_previous_</span><span class="si">{</span><span class="n">root_id</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
                <span class="n">parent</span><span class="o">=</span><span class="n">root_id</span>
            <span class="p">)</span>
            <span class="n">tree</span><span class="o">.</span><span class="n">create_node</span><span class="p">(</span>
                <span class="n">tag</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;Best until now = </span><span class="si">{</span><span class="n">best</span><span class="si">:</span><span class="s2">6</span><span class="si">}</span><span class="s2"> (</span><span class="si">{</span><span class="n">diff_with_best_colored</span><span class="si">:</span><span class="s2">8</span><span class="si">}</span><span class="s2">)&quot;</span><span class="p">,</span>
                <span class="n">identifier</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;1_best_</span><span class="si">{</span><span class="n">root_id</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
                <span class="n">parent</span><span class="o">=</span><span class="n">root_id</span>
            <span class="p">)</span>
        <span class="k">return</span> <span class="n">tree</span>

    <span class="n">train_tree</span> <span class="o">=</span> <span class="n">Tree</span><span class="p">()</span>
    <span class="n">train_tree</span><span class="o">.</span><span class="n">create_node</span><span class="p">(</span><span class="s2">&quot;Training&quot;</span><span class="p">,</span> <span class="s2">&quot;Training&quot;</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">train_monitored_values</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="n">train_tree</span><span class="o">.</span><span class="n">paste</span><span class="p">(</span><span class="s1">&#39;Training&#39;</span><span class="p">,</span> <span class="n">new_tree</span><span class="o">=</span><span class="n">_generate_tree</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">monitored_value</span><span class="o">=</span><span class="n">value</span><span class="p">))</span>

    <span class="n">valid_tree</span> <span class="o">=</span> <span class="n">Tree</span><span class="p">()</span>
    <span class="n">valid_tree</span><span class="o">.</span><span class="n">create_node</span><span class="p">(</span><span class="s2">&quot;Validation&quot;</span><span class="p">,</span> <span class="s2">&quot;Validation&quot;</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">valid_monitored_values</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="n">valid_tree</span><span class="o">.</span><span class="n">paste</span><span class="p">(</span><span class="s1">&#39;Validation&#39;</span><span class="p">,</span> <span class="n">new_tree</span><span class="o">=</span><span class="n">_generate_tree</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">monitored_value</span><span class="o">=</span><span class="n">value</span><span class="p">))</span>

    <span class="n">summary_tree</span> <span class="o">=</span> <span class="n">Tree</span><span class="p">()</span>
    <span class="n">summary_tree</span><span class="o">.</span><span class="n">create_node</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;SUMMARY OF EPOCH </span><span class="si">{</span><span class="n">epoch</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="s2">&quot;Summary&quot;</span><span class="p">)</span>
    <span class="n">summary_tree</span><span class="o">.</span><span class="n">paste</span><span class="p">(</span><span class="s2">&quot;Summary&quot;</span><span class="p">,</span> <span class="n">train_tree</span><span class="p">)</span>
    <span class="n">summary_tree</span><span class="o">.</span><span class="n">paste</span><span class="p">(</span><span class="s2">&quot;Summary&quot;</span><span class="p">,</span> <span class="n">valid_tree</span><span class="p">)</span>
    <span class="n">summary_tree</span><span class="o">.</span><span class="n">show</span><span class="p">()</span></div>


<div class="viewcode-block" id="try_port"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.try_port">[docs]</a><span class="k">def</span> <span class="nf">try_port</span><span class="p">(</span><span class="n">port</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    try_port - Helper method for tensorboard port binding</span>
<span class="sd">    :param port:</span>
<span class="sd">    :return:</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">sock</span> <span class="o">=</span> <span class="n">socket</span><span class="o">.</span><span class="n">socket</span><span class="p">(</span><span class="n">socket</span><span class="o">.</span><span class="n">AF_INET</span><span class="p">,</span> <span class="n">socket</span><span class="o">.</span><span class="n">SOCK_STREAM</span><span class="p">)</span>
    <span class="n">is_port_available</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">sock</span><span class="o">.</span><span class="n">bind</span><span class="p">((</span><span class="s2">&quot;localhost&quot;</span><span class="p">,</span> <span class="n">port</span><span class="p">))</span>
        <span class="n">is_port_available</span> <span class="o">=</span> <span class="kc">True</span>

    <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">ex</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Port &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">port</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39; is in use&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">ex</span><span class="p">))</span>

    <span class="n">sock</span><span class="o">.</span><span class="n">close</span><span class="p">()</span>
    <span class="k">return</span> <span class="n">is_port_available</span></div>


<div class="viewcode-block" id="launch_tensorboard_process"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.launch_tensorboard_process">[docs]</a><span class="k">def</span> <span class="nf">launch_tensorboard_process</span><span class="p">(</span><span class="n">checkpoints_dir_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">sleep_postpone</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> <span class="n">port</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Process</span><span class="p">,</span> <span class="nb">int</span><span class="p">]:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    launch_tensorboard_process - Default behavior is to scan all free ports from 6006-6016 and try using them</span>
<span class="sd">                                 unless port is defined by the user</span>
<span class="sd">        :param checkpoints_dir_path:</span>
<span class="sd">        :param sleep_postpone:</span>
<span class="sd">        :param port:</span>
<span class="sd">        :return: tuple of tb process, port</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">logdir_path</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">Path</span><span class="p">(</span><span class="n">checkpoints_dir_path</span><span class="p">)</span><span class="o">.</span><span class="n">parent</span><span class="o">.</span><span class="n">absolute</span><span class="p">())</span>
    <span class="n">tb_cmd</span> <span class="o">=</span> <span class="s1">&#39;tensorboard --logdir=&#39;</span> <span class="o">+</span> <span class="n">logdir_path</span> <span class="o">+</span> <span class="s1">&#39; --bind_all&#39;</span>
    <span class="k">if</span> <span class="n">port</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">tb_ports</span> <span class="o">=</span> <span class="p">[</span><span class="n">port</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">tb_ports</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="mi">6006</span><span class="p">,</span> <span class="mi">6016</span><span class="p">)</span>

    <span class="k">for</span> <span class="n">tb_port</span> <span class="ow">in</span> <span class="n">tb_ports</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">try_port</span><span class="p">(</span><span class="n">tb_port</span><span class="p">):</span>
            <span class="k">continue</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Starting Tensor-Board process on port: &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">tb_port</span><span class="p">))</span>
            <span class="n">tensor_board_process</span> <span class="o">=</span> <span class="n">Process</span><span class="p">(</span><span class="n">target</span><span class="o">=</span><span class="n">os</span><span class="o">.</span><span class="n">system</span><span class="p">,</span> <span class="n">args</span><span class="o">=</span><span class="p">([</span><span class="n">tb_cmd</span> <span class="o">+</span> <span class="s1">&#39; --port=&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">tb_port</span><span class="p">)]))</span>
            <span class="n">tensor_board_process</span><span class="o">.</span><span class="n">daemon</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="n">tensor_board_process</span><span class="o">.</span><span class="n">start</span><span class="p">()</span>

            <span class="c1"># LET THE TENSORBOARD PROCESS START</span>
            <span class="k">if</span> <span class="n">sleep_postpone</span><span class="p">:</span>
                <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
            <span class="k">return</span> <span class="n">tensor_board_process</span><span class="p">,</span> <span class="n">tb_port</span>

    <span class="c1"># RETURNING IRRELEVANT VALUES</span>
    <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Failed to initialize Tensor-Board process on port: &#39;</span> <span class="o">+</span> <span class="s1">&#39;, &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">tb_ports</span><span class="p">)))</span>
    <span class="k">return</span> <span class="kc">None</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span></div>


<div class="viewcode-block" id="init_summary_writer"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.init_summary_writer">[docs]</a><span class="k">def</span> <span class="nf">init_summary_writer</span><span class="p">(</span><span class="n">tb_dir</span><span class="p">,</span> <span class="n">checkpoint_loaded</span><span class="p">,</span> <span class="n">user_prompt</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Remove previous tensorboard files from directory and launch a tensor board process&quot;&quot;&quot;</span>
    <span class="c1"># If the training is from scratch, Walk through destination folder and delete existing tensorboard logs</span>
    <span class="n">user</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">checkpoint_loaded</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">filename</span> <span class="ow">in</span> <span class="n">os</span><span class="o">.</span><span class="n">listdir</span><span class="p">(</span><span class="n">tb_dir</span><span class="p">):</span>
            <span class="k">if</span> <span class="s1">&#39;events&#39;</span> <span class="ow">in</span> <span class="n">filename</span><span class="p">:</span>
                <span class="k">if</span> <span class="ow">not</span> <span class="n">user_prompt</span><span class="p">:</span>
                    <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;&quot;</span><span class="si">{}</span><span class="s1">&quot; will not be deleted&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">filename</span><span class="p">))</span>
                    <span class="k">continue</span>

                <span class="k">while</span> <span class="kc">True</span><span class="p">:</span>
                    <span class="c1"># Verify with user before deleting old tensorboard files</span>
                    <span class="n">user</span> <span class="o">=</span> <span class="nb">input</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\n</span><span class="s1">OLDER TENSORBOARD FILES EXISTS IN EXPERIMENT FOLDER:</span><span class="se">\n</span><span class="s1">&quot;</span><span class="si">{}</span><span class="s1">&quot;</span><span class="se">\n</span><span class="s1">&#39;</span>
                                 <span class="s1">&#39;DO YOU WANT TO DELETE THEM? [y/n]&#39;</span>
                                 <span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">filename</span><span class="p">))</span> <span class="k">if</span> <span class="p">(</span><span class="n">user</span> <span class="o">!=</span> <span class="s1">&#39;n&#39;</span> <span class="ow">or</span> <span class="n">user</span> <span class="o">!=</span> <span class="s1">&#39;y&#39;</span><span class="p">)</span> <span class="k">else</span> <span class="n">user</span>
                    <span class="k">if</span> <span class="n">user</span> <span class="o">==</span> <span class="s1">&#39;y&#39;</span><span class="p">:</span>
                        <span class="n">os</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="s1">&#39;</span><span class="si">{}</span><span class="s1">/</span><span class="si">{}</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">tb_dir</span><span class="p">,</span> <span class="n">filename</span><span class="p">))</span>
                        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;DELETED: </span><span class="si">{}</span><span class="s1">!&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">filename</span><span class="p">))</span>
                        <span class="k">break</span>
                    <span class="k">elif</span> <span class="n">user</span> <span class="o">==</span> <span class="s1">&#39;n&#39;</span><span class="p">:</span>
                        <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;&quot;</span><span class="si">{}</span><span class="s1">&quot; will not be deleted&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">filename</span><span class="p">))</span>
                        <span class="k">break</span>
                    <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Unknown answer...&#39;</span><span class="p">)</span>

    <span class="c1"># Launch a tensorboard process</span>
    <span class="k">return</span> <span class="n">SummaryWriter</span><span class="p">(</span><span class="n">tb_dir</span><span class="p">)</span></div>


<div class="viewcode-block" id="add_log_to_file"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.add_log_to_file">[docs]</a><span class="k">def</span> <span class="nf">add_log_to_file</span><span class="p">(</span><span class="n">filename</span><span class="p">,</span> <span class="n">results_titles_list</span><span class="p">,</span> <span class="n">results_values_list</span><span class="p">,</span> <span class="n">epoch</span><span class="p">,</span> <span class="n">max_epochs</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Add a message to the log file&quot;&quot;&quot;</span>
    <span class="c1"># -Note: opening and closing the file every time is in-efficient. It is done for experimental purposes</span>
    <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">filename</span><span class="p">,</span> <span class="s1">&#39;a&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
        <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\n</span><span class="s1">Epoch (</span><span class="si">%d</span><span class="s1">/</span><span class="si">%d</span><span class="s1">)  - &#39;</span> <span class="o">%</span> <span class="p">(</span><span class="n">epoch</span><span class="p">,</span> <span class="n">max_epochs</span><span class="p">))</span>
        <span class="k">for</span> <span class="n">result_title</span><span class="p">,</span> <span class="n">result_value</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">results_titles_list</span><span class="p">,</span> <span class="n">results_values_list</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">result_value</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
                <span class="n">result_value</span> <span class="o">=</span> <span class="n">result_value</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
            <span class="n">f</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="n">result_title</span> <span class="o">+</span> <span class="s1">&#39;: &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">result_value</span><span class="p">)</span> <span class="o">+</span> <span class="s1">&#39;</span><span class="se">\t</span><span class="s1">&#39;</span><span class="p">)</span></div>


<div class="viewcode-block" id="write_training_results"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.write_training_results">[docs]</a><span class="k">def</span> <span class="nf">write_training_results</span><span class="p">(</span><span class="n">writer</span><span class="p">,</span> <span class="n">results_titles_list</span><span class="p">,</span> <span class="n">results_values_list</span><span class="p">,</span> <span class="n">epoch</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Stores the training and validation loss and accuracy for current epoch in a tensorboard file&quot;&quot;&quot;</span>
    <span class="k">for</span> <span class="n">res_key</span><span class="p">,</span> <span class="n">res_val</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">results_titles_list</span><span class="p">,</span> <span class="n">results_values_list</span><span class="p">):</span>
        <span class="c1"># USE ONLY LOWER-CASE LETTERS AND REPLACE SPACES WITH &#39;_&#39; TO AVOID MANY TITLES FOR THE SAME KEY</span>
        <span class="n">corrected_res_key</span> <span class="o">=</span> <span class="n">res_key</span><span class="o">.</span><span class="n">lower</span><span class="p">()</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s1">&#39; &#39;</span><span class="p">,</span> <span class="s1">&#39;_&#39;</span><span class="p">)</span>
        <span class="n">writer</span><span class="o">.</span><span class="n">add_scalar</span><span class="p">(</span><span class="n">corrected_res_key</span><span class="p">,</span> <span class="n">res_val</span><span class="p">,</span> <span class="n">epoch</span><span class="p">)</span>
    <span class="n">writer</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span></div>


<div class="viewcode-block" id="write_hpms"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.write_hpms">[docs]</a><span class="k">def</span> <span class="nf">write_hpms</span><span class="p">(</span><span class="n">writer</span><span class="p">,</span> <span class="n">hpmstructs</span><span class="o">=</span><span class="p">[],</span> <span class="n">special_conf</span><span class="o">=</span><span class="p">{}):</span>
    <span class="sd">&quot;&quot;&quot;Stores the training and dataset hyper params in the tensorboard file&quot;&quot;&quot;</span>
    <span class="n">hpm_string</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
    <span class="k">for</span> <span class="n">hpm</span> <span class="ow">in</span> <span class="n">hpmstructs</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">hpm</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="n">hpm_string</span> <span class="o">+=</span> <span class="s1">&#39;</span><span class="si">{}</span><span class="s1">: </span><span class="si">{}</span><span class="s1">  </span><span class="se">\n</span><span class="s1">  &#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">val</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">special_conf</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="n">hpm_string</span> <span class="o">+=</span> <span class="s1">&#39;</span><span class="si">{}</span><span class="s1">: </span><span class="si">{}</span><span class="s1">  </span><span class="se">\n</span><span class="s1">  &#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">val</span><span class="p">)</span>
    <span class="n">writer</span><span class="o">.</span><span class="n">add_text</span><span class="p">(</span><span class="s2">&quot;Hyper_parameters&quot;</span><span class="p">,</span> <span class="n">hpm_string</span><span class="p">)</span>
    <span class="n">writer</span><span class="o">.</span><span class="n">flush</span><span class="p">()</span></div>


<span class="c1"># TODO: This should probably move into datasets/datasets_utils.py?</span>
<div class="viewcode-block" id="unpack_batch_items"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.unpack_batch_items">[docs]</a><span class="k">def</span> <span class="nf">unpack_batch_items</span><span class="p">(</span><span class="n">batch_items</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">tuple</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Adds support for unpacking batch items in train/validation loop.</span>

<span class="sd">    @param batch_items: (Union[tuple, torch.Tensor]) returned by the data loader, which is expected to be in one of</span>
<span class="sd">         the following formats:</span>
<span class="sd">            1. torch.Tensor or tuple, s.t inputs = batch_items[0], targets = batch_items[1] and len(batch_items) = 2</span>
<span class="sd">            2. tuple: (inputs, targets, additional_batch_items)</span>

<span class="sd">         where inputs are fed to the network, targets are their corresponding labels and additional_batch_items is a</span>
<span class="sd">         dictionary (format {additional_batch_item_i_name: additional_batch_item_i ...}) which can be accessed through</span>
<span class="sd">         the phase context under the attribute additional_batch_item_i_name, using a phase callback.</span>


<span class="sd">    @return: inputs, target, additional_batch_items</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">additional_batch_items</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">batch_items</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">inputs</span><span class="p">,</span> <span class="n">target</span> <span class="o">=</span> <span class="n">batch_items</span>

    <span class="k">elif</span> <span class="nb">len</span><span class="p">(</span><span class="n">batch_items</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
        <span class="n">inputs</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">additional_batch_items</span> <span class="o">=</span> <span class="n">batch_items</span>

    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="n">UnsupportedBatchItemsFormat</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">inputs</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">additional_batch_items</span></div>


<div class="viewcode-block" id="log_uncaught_exceptions"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.sg_model_utils.log_uncaught_exceptions">[docs]</a><span class="k">def</span> <span class="nf">log_uncaught_exceptions</span><span class="p">(</span><span class="n">logger</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Makes logger log uncaught exceptions</span>
<span class="sd">    @param logger: logging.Logger</span>

<span class="sd">    @return: None</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">handle_exception</span><span class="p">(</span><span class="n">exc_type</span><span class="p">,</span> <span class="n">exc_value</span><span class="p">,</span> <span class="n">exc_traceback</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">exc_type</span><span class="p">,</span> <span class="ne">KeyboardInterrupt</span><span class="p">):</span>
            <span class="n">sys</span><span class="o">.</span><span class="n">__excepthook__</span><span class="p">(</span><span class="n">exc_type</span><span class="p">,</span> <span class="n">exc_value</span><span class="p">,</span> <span class="n">exc_traceback</span><span class="p">)</span>
            <span class="k">return</span>

        <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="s2">&quot;Uncaught exception&quot;</span><span class="p">,</span> <span class="n">exc_info</span><span class="o">=</span><span class="p">(</span><span class="n">exc_type</span><span class="p">,</span> <span class="n">exc_value</span><span class="p">,</span> <span class="n">exc_traceback</span><span class="p">))</span>

    <span class="n">sys</span><span class="o">.</span><span class="n">excepthook</span> <span class="o">=</span> <span class="n">handle_exception</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.ssd_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.ssd_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.ssd_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">itertools</span>
<span class="kn">from</span> <span class="nn">math</span> <span class="kn">import</span> <span class="n">sqrt</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">torch</span>

<span class="kn">from</span> <span class="nn">super_gradients.training.utils.detection_utils</span> <span class="kn">import</span> <span class="n">non_max_suppression</span><span class="p">,</span> <span class="n">NMS_Type</span><span class="p">,</span> \
    <span class="n">matrix_non_max_suppression</span><span class="p">,</span> <span class="n">DetectionPostPredictionCallback</span>


<div class="viewcode-block" id="DefaultBoxes"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ssd_utils.DefaultBoxes">[docs]</a><span class="k">class</span> <span class="nc">DefaultBoxes</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Default Boxes, (aka: anchor boxes or priors boxes) used by SSD model</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fig_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">feat_size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">scales</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">aspect_ratios</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]],</span>
                 <span class="n">scale_xy</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">scale_wh</span><span class="o">=</span><span class="mf">0.2</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        For each feature map i (each predicting level, grids) the anchors (a.k.a. default boxes) will be:</span>
<span class="sd">        [</span>
<span class="sd">            [s, s], [sqrt(s * s_next), sqrt(s * s_next)],</span>
<span class="sd">            [s * sqrt(alpha1), s / sqrt(alpha1)], [s / sqrt(alpha1), s * sqrt(alpha1)],</span>
<span class="sd">            ...</span>
<span class="sd">            [s * sqrt(alphaN), s / sqrt(alphaN)], [s / sqrt(alphaN), s * sqrt(alphaN)]</span>
<span class="sd">        ] / fig_size</span>
<span class="sd">        where:</span>
<span class="sd">            * s = scale[i] - this level&#39;s scale</span>
<span class="sd">            * s_next = scale[i + 1] - next level&#39;s scale</span>
<span class="sd">            * alpha1, ... alphaN - this level&#39;s alphas, e.g. [2, 3]</span>
<span class="sd">            * fig_size - input image resolution</span>

<span class="sd">        Because of division by image resolution, the anchors will be in image coordinates normalized to [0, 1]</span>

<span class="sd">        :param fig_size:        input image resolution</span>
<span class="sd">        :param feat_size:       resolution of all feature maps with predictions (grids)</span>
<span class="sd">        :param scales:          anchor sizes in pixels for each feature level;</span>
<span class="sd">                                one value per level will be used to generate anchors based on the formula above</span>
<span class="sd">        :param aspect_ratios:   lists of alpha values for each feature map</span>
<span class="sd">        :param scale_xy:        predicted boxes will be with a factor scale_xy</span>
<span class="sd">                                so will be multiplied by scale_xy during post-prediction processing;</span>
<span class="sd">                                e.g. scale 0.1 means that prediction will be 10 times bigger</span>
<span class="sd">                                (improves predictions quality)</span>
<span class="sd">        :param scale_wh:        same logic as in scale_xy, but for width and height.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">feat_size</span> <span class="o">=</span> <span class="n">feat_size</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fig_size</span> <span class="o">=</span> <span class="n">fig_size</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">scale_xy_</span> <span class="o">=</span> <span class="n">scale_xy</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">scale_wh_</span> <span class="o">=</span> <span class="n">scale_wh</span>
        <span class="c1"># According to https://github.com/weiliu89/caffe</span>
        <span class="c1"># Calculation method slightly different from paper</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">scales</span> <span class="o">=</span> <span class="n">scales</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">aspect_ratios</span> <span class="o">=</span> <span class="n">aspect_ratios</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">default_boxes</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">num_anchors</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="c1"># size of feature and number of feature</span>
        <span class="k">for</span> <span class="n">idx</span><span class="p">,</span> <span class="n">sfeat</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">feat_size</span><span class="p">):</span>

            <span class="n">sk1</span> <span class="o">=</span> <span class="n">scales</span><span class="p">[</span><span class="n">idx</span><span class="p">]</span>
            <span class="n">sk2</span> <span class="o">=</span> <span class="n">scales</span><span class="p">[</span><span class="n">idx</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]</span>
            <span class="n">sk3</span> <span class="o">=</span> <span class="n">sqrt</span><span class="p">(</span><span class="n">sk1</span> <span class="o">*</span> <span class="n">sk2</span><span class="p">)</span>
            <span class="n">all_sizes</span> <span class="o">=</span> <span class="p">[(</span><span class="n">sk1</span><span class="p">,</span> <span class="n">sk1</span><span class="p">),</span> <span class="p">(</span><span class="n">sk3</span><span class="p">,</span> <span class="n">sk3</span><span class="p">)]</span>

            <span class="k">for</span> <span class="n">alpha</span> <span class="ow">in</span> <span class="n">aspect_ratios</span><span class="p">[</span><span class="n">idx</span><span class="p">]:</span>
                <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">sk1</span> <span class="o">*</span> <span class="n">sqrt</span><span class="p">(</span><span class="n">alpha</span><span class="p">),</span> <span class="n">sk1</span> <span class="o">/</span> <span class="n">sqrt</span><span class="p">(</span><span class="n">alpha</span><span class="p">)</span>
                <span class="n">all_sizes</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">))</span>
                <span class="n">all_sizes</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">))</span>

            <span class="n">all_sizes</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="n">all_sizes</span><span class="p">)</span> <span class="o">/</span> <span class="n">fig_size</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">num_anchors</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">all_sizes</span><span class="p">))</span>
            <span class="k">for</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="ow">in</span> <span class="n">all_sizes</span><span class="p">:</span>
                <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span> <span class="ow">in</span> <span class="n">itertools</span><span class="o">.</span><span class="n">product</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">sfeat</span><span class="p">),</span> <span class="n">repeat</span><span class="o">=</span><span class="mi">2</span><span class="p">):</span>
                    <span class="n">cx</span><span class="p">,</span> <span class="n">cy</span> <span class="o">=</span> <span class="p">(</span><span class="n">j</span> <span class="o">+</span> <span class="mf">0.5</span><span class="p">)</span> <span class="o">/</span> <span class="n">sfeat</span><span class="p">,</span> <span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="mf">0.5</span><span class="p">)</span> <span class="o">/</span> <span class="n">sfeat</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">default_boxes</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">cx</span><span class="p">,</span> <span class="n">cy</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">))</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</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="bp">self</span><span class="o">.</span><span class="n">default_boxes</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">float</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="nb">min</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="nb">max</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>

        <span class="c1"># For IoU calculation</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dboxes_xyxy</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dboxes_xyxy</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dboxes_xyxy</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dboxes_xyxy</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dboxes_xyxy</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span><span class="p">[:,</span> <span class="mi">3</span><span class="p">]</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">scale_xy</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale_xy_</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">scale_wh</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">scale_wh_</span>

    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">order</span><span class="o">=</span><span class="s2">&quot;xyxy&quot;</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">order</span> <span class="o">==</span> <span class="s2">&quot;xyxy&quot;</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes_xyxy</span>
        <span class="k">if</span> <span class="n">order</span> <span class="o">==</span> <span class="s2">&quot;xywh&quot;</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">dboxes</span></div>


<div class="viewcode-block" id="SSDPostPredictCallback"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ssd_utils.SSDPostPredictCallback">[docs]</a><span class="k">class</span> <span class="nc">SSDPostPredictCallback</span><span class="p">(</span><span class="n">DetectionPostPredictionCallback</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    post prediction callback module to convert and filter predictions coming from the SSD net to a format</span>
<span class="sd">    used by all other detection models</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">conf</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.001</span><span class="p">,</span> <span class="n">iou</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.6</span><span class="p">,</span> <span class="n">classes</span><span class="p">:</span> <span class="nb">list</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
                 <span class="n">max_predictions</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">300</span><span class="p">,</span>
                 <span class="n">nms_type</span><span class="p">:</span> <span class="n">NMS_Type</span> <span class="o">=</span> <span class="n">NMS_Type</span><span class="o">.</span><span class="n">ITERATIVE</span><span class="p">,</span>
                 <span class="n">multi_label_per_box</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Predictions of SSD contain unnormalized probabilities for a background class,</span>
<span class="sd">        together with confidences for all the dataset classes. Background will be utilized and discarded,</span>
<span class="sd">        so this callback will return 0-based classes without background</span>
<span class="sd">        :param conf: confidence threshold</span>
<span class="sd">        :param iou: IoU threshold</span>
<span class="sd">        :param classes: (optional list) filter by class</span>
<span class="sd">        :param nms_type: the type of nms to use (iterative or matrix)</span>
<span class="sd">        :param multi_label_per_box: whether to use re-use each box with all possible labels</span>
<span class="sd">                                    (instead of the maximum confidence all confidences above threshold</span>
<span class="sd">                                    will be sent to NMS)</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">SSDPostPredictCallback</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>
        <span class="bp">self</span><span class="o">.</span><span class="n">conf</span> <span class="o">=</span> <span class="n">conf</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">iou</span> <span class="o">=</span> <span class="n">iou</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">nms_type</span> <span class="o">=</span> <span class="n">nms_type</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">classes</span> <span class="o">=</span> <span class="n">classes</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">max_predictions</span> <span class="o">=</span> <span class="n">max_predictions</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">multi_label_per_box</span> <span class="o">=</span> <span class="n">multi_label_per_box</span>

<div class="viewcode-block" id="SSDPostPredictCallback.forward"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.ssd_utils.SSDPostPredictCallback.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">predictions</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="n">nms_input</span> <span class="o">=</span> <span class="n">predictions</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">nms_type</span> <span class="o">==</span> <span class="n">NMS_Type</span><span class="o">.</span><span class="n">ITERATIVE</span><span class="p">:</span>
            <span class="n">nms_res</span> <span class="o">=</span> <span class="n">non_max_suppression</span><span class="p">(</span><span class="n">nms_input</span><span class="p">,</span> <span class="n">conf_thres</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">conf</span><span class="p">,</span> <span class="n">iou_thres</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">iou</span><span class="p">,</span>
                                          <span class="n">multi_label_per_box</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">multi_label_per_box</span><span class="p">,</span> <span class="n">with_confidence</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">nms_res</span> <span class="o">=</span> <span class="n">matrix_non_max_suppression</span><span class="p">(</span><span class="n">nms_input</span><span class="p">,</span> <span class="n">conf_thres</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">conf</span><span class="p">,</span>
                                                 <span class="n">max_num_of_detections</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">max_predictions</span><span class="p">)</span>

        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_filter_max_predictions</span><span class="p">(</span><span class="n">nms_res</span><span class="p">)</span></div>

    <span class="k">def</span> <span class="nf">_filter_max_predictions</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">res</span><span class="p">:</span> <span class="n">List</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">:</span>
        <span class="n">res</span><span class="p">[:]</span> <span class="o">=</span> <span class="p">[</span><span class="n">im</span><span class="p">[:</span><span class="bp">self</span><span class="o">.</span><span class="n">max_predictions</span><span class="p">]</span> <span class="k">if</span> <span class="p">(</span><span class="n">im</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">im</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">max_predictions</span><span class="p">)</span> <span class="k">else</span> <span class="n">im</span> <span class="k">for</span> <span class="n">im</span> <span class="ow">in</span> <span class="n">res</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">res</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
@@ -3,10 +3,11 @@
 
                             <head>
                
 
                               <meta charset="utf-8" />
                
 
                               <meta name="viewport" content="width=device-width, initial-scale=1.0" />
                
 
                            -  <title>super_gradients.training.utils.utils &mdash; SuperGradients 1.0 documentation</title>
                
 
                            +  <title>super_gradients.training.utils.utils &mdash; SuperGradients 3.0.3 documentation</title>
                
 
                                   <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
                
 
                                   <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
                
 
                            +      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
                
 
                               <!--[if lt IE 9]>
                
 
                                 <script src="../../../../_static/js/html5shiv.min.js"></script>
                
 
                               <![endif]-->
                
@@ -14,7 +15,9 @@
 
                                     <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
                
 
                                     <script src="../../../../_static/jquery.js"></script>
                
 
                                     <script src="../../../../_static/underscore.js"></script>
                
 
                            +        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
                
 
                                     <script src="../../../../_static/doctools.js"></script>
                
 
                            +        <script src="../../../../_static/sphinx_highlight.js"></script>
                
 
                                 <script src="../../../../_static/js/theme.js"></script>
                
 
                                 <link rel="index" title="Index" href="../../../../genindex.html" />
                
 
                                 <link rel="search" title="Search" href="../../../../search.html" /> 
                
@@ -35,30 +38,29 @@
 
                               </form>
                
 
                             </div>
                
 
                                     </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
                
 
                            -              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                            +              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
                
 
                             <ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
                
 
                            +<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
                
 
                             </ul>
                
 
                            -<p class="caption"><span class="caption-text">Technical Documentation</span></p>
                
 
                            +<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
                
 
                             <ul>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
                
 
                             <li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
                
 
                            -</ul>
                
 
                            -<p class="caption"><span class="caption-text">User Guide</span></p>
                
 
                            -<ul>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
                
 
                            -<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
                
 
                             </ul>
                
 
                                     </div>
                
@@ -90,7 +92,7 @@
 
                             <span class="kn">import</span> <span class="nn">time</span>
                
 
                             <span class="kn">from</span> <span class="nn">functools</span> <span class="kn">import</span> <span class="n">lru_cache</span>
                
 
                             <span class="kn">from</span> <span class="nn">pathlib</span> <span class="kn">import</span> <span class="n">Path</span>
                
 
                            -<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Mapping</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">List</span>
                
 
                            +<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Mapping</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Dict</span>
                
 
                             <span class="kn">from</span> <span class="nn">zipfile</span> <span class="kn">import</span> <span class="n">ZipFile</span>
                
 
                             <span class="kn">import</span> <span class="nn">os</span>
                
 
                             <span class="kn">from</span> <span class="nn">jsonschema</span> <span class="kn">import</span> <span class="n">validate</span>
                
@@ -112,28 +114,40 @@
 
                             <span class="n">logger</span> <span class="o">=</span> <span class="n">get_logger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="convert_to_tensor"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.convert_to_tensor">[docs]</a><span class="k">def</span> <span class="nf">convert_to_tensor</span><span class="p">(</span><span class="n">array</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">empty_list</span><span class="p">():</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;Instantiate an empty list. This is a workaround to generate a list with a function call in hydra, instead of the &quot;[]&quot;.&quot;&quot;&quot;</span>
                
 
                            +    <span class="k">return</span> <span class="nb">list</span><span class="p">()</span>
                
 
                            +
                
 
                            +
                
 
                            +<div class="viewcode-block" id="convert_to_tensor"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.convert_to_tensor">[docs]</a><span class="k">def</span> <span class="nf">convert_to_tensor</span><span class="p">(</span><span class="n">array</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Converts numpy arrays and lists to Torch tensors before calculation losses</span>
                
 
                             <span class="sd">    :param array: torch.tensor / Numpy array / List</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                                 <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">FloatTensor</span><span class="p">(</span><span class="n">array</span><span class="p">)</span> <span class="k">if</span> <span class="nb">type</span><span class="p">(</span><span class="n">array</span><span class="p">)</span> <span class="o">!=</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span> <span class="k">else</span> <span class="n">array</span></div>
                
 
                            -<div class="viewcode-block" id="HpmStruct"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.HpmStruct">[docs]</a><span class="k">class</span> <span class="nc">HpmStruct</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="HpmStruct"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.HpmStruct">[docs]</a><span class="k">class</span> <span class="nc">HpmStruct</span><span class="p">:</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">entries</span><span class="p">):</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">entries</span><span class="p">)</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">schema</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                            -<div class="viewcode-block" id="HpmStruct.set_schema"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.HpmStruct.set_schema">[docs]</a>    <span class="k">def</span> <span class="nf">set_schema</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">schema</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="HpmStruct.set_schema"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.HpmStruct.set_schema">[docs]</a>    <span class="k">def</span> <span class="nf">set_schema</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">schema</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">schema</span> <span class="o">=</span> <span class="n">schema</span></div>
                
 
                            -<div class="viewcode-block" id="HpmStruct.override"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.HpmStruct.override">[docs]</a>    <span class="k">def</span> <span class="nf">override</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">entries</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="HpmStruct.override"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.HpmStruct.override">[docs]</a>    <span class="k">def</span> <span class="nf">override</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">entries</span><span class="p">):</span>
                
 
                                     <span class="n">recursive_override</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">,</span> <span class="n">entries</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="HpmStruct.to_dict"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.HpmStruct.to_dict">[docs]</a>    <span class="k">def</span> <span class="nf">to_dict</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            -        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span></div>
                
 
                            +<div class="viewcode-block" id="HpmStruct.to_dict"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.HpmStruct.to_dict">[docs]</a>    <span class="k">def</span> <span class="nf">to_dict</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">include_schema</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span
                
 
                            +        <span class="sd">&quot;&quot;&quot;Convert this HpmStruct instance into a dict.</span>
                
 
                            +<span class="sd">        :param include_schema: If True, also return the field &quot;schema&quot;</span>
                
 
                            +<span class="sd">        :return: Dict representation of this HpmStruct instance.</span>
                
 
                            +<span class="sd">        &quot;&quot;&quot;</span>
                
 
                            +        <span class="n">out_dict</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
                
 
                            +        <span class="k">if</span> <span class="ow">not</span> <span class="n">include_schema</span><span class="p">:</span>
                
 
                            +            <span class="n">out_dict</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;schema&quot;</span><span class="p">)</span>
                
 
                            +        <span class="k">return</span> <span class="n">out_dict</span></div>
                
 
                            -<div class="viewcode-block" id="HpmStruct.validate"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.HpmStruct.validate">[docs]</a>    <span class="k">def</span> <span class="nf">validate</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="HpmStruct.validate"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.HpmStruct.validate">[docs]</a>    <span class="k">def</span> <span class="nf">validate</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">        Validate the current dict values according to the provided schema</span>
                
 
                             <span class="sd">        :raises</span>
                
@@ -147,16 +161,16 @@
 
                                         <span class="n">validate</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">schema</span><span class="p">)</span></div></div>
                
 
                            -<div class="viewcode-block" id="WrappedModel"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.WrappedModel">[docs]</a><span class="k">class</span> <span class="nc">WrappedModel</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="WrappedModel"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.WrappedModel">[docs]</a><span class="k">class</span> <span class="nc">WrappedModel</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
                
 
                                 <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">module</span><span class="p">):</span>
                
 
                                     <span class="nb">super</span><span class="p">(</span><span class="n">WrappedModel</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>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">module</span> <span class="o">=</span> <span class="n">module</span>  <span class="c1"># that I actually define.</span>
                
 
                            -<div class="viewcode-block" id="WrappedModel.forward"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.WrappedModel.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="WrappedModel.forward"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.WrappedModel.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
                
 
                                     <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">module</span><span class="p">(</span><span class="n">x</span><span class="p">)</span></div></div>
                
 
                            -<div class="viewcode-block" id="Timer"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.Timer">[docs]</a><span class="k">class</span> <span class="nc">Timer</span><span class="p">:</span>
                
 
                            +<div class="viewcode-block" id="Timer"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.Timer">[docs]</a><span class="k">class</span> <span class="nc">Timer</span><span class="p">:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;A class to measure time handling both GPU &amp; CPU processes</span>
                
 
                             <span class="sd">    Returns time in milliseconds&quot;&quot;&quot;</span>
                
@@ -174,13 +188,13 @@
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">starter</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">ender</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span>
                
 
                            -<div class="viewcode-block" id="Timer.start"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.Timer.start">[docs]</a>    <span class="k">def</span> <span class="nf">start</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="Timer.start"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.Timer.start">[docs]</a>    <span class="k">def</span> <span class="nf">start</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">on_gpu</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">starter</span><span class="o">.</span><span class="n">record</span><span class="p">()</span>
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">starter</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span></div>
                
 
                            -<div class="viewcode-block" id="Timer.stop"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.Timer.stop">[docs]</a>    <span class="k">def</span> <span class="nf">stop</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="Timer.stop"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.Timer.stop">[docs]</a>    <span class="k">def</span> <span class="nf">stop</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">on_gpu</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">ender</span><span class="o">.</span><span class="n">record</span><span class="p">()</span>
                
 
                                         <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">synchronize</span><span class="p">()</span>
                
@@ -193,7 +207,7 @@
 
                                     <span class="k">return</span> <span class="n">timer</span></div></div>
                
 
                            -<div class="viewcode-block" id="AverageMeter"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.AverageMeter">[docs]</a><span class="k">class</span> <span class="nc">AverageMeter</span><span class="p">:</span>
                
 
                            +<span class="k">class</span> <span class="nc">AverageMeter</span><span class="p">:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;A class to calculate the average of a metric, for each batch</span>
                
 
                             <span class="sd">    during training/testing&quot;&quot;&quot;</span>
                
@@ -201,7 +215,7 @@
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">_sum</span> <span class="o">=</span> <span class="kc">None</span>
                
 
                                     <span class="bp">self</span><span class="o">.</span><span class="n">_count</span> <span class="o">=</span> <span class="mi">0</span>
                
 
                            -<div class="viewcode-block" id="AverageMeter.update"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.AverageMeter.update">[docs]</a>    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span cl
                
 
                            +    <span class="k">def</span> <span class="nf">update</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p
                
 
                                     <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
                
 
                                         <span class="n">value</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">value</span><span class="p">)</span>
                
@@ -211,20 +225,20 @@
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                                         <span class="bp">self</span><span class="o">.</span><span class="n">_sum</span> <span class="o">+=</span> <span class="n">value</span> <span class="o">*</span> <span class="n">batch_size</span>
                
 
                            -        <span class="bp">self</span><span class="o">.</span><span class="n">_count</span> <span class="o">+=</span> <span class="n">batch_size</span></div>
                
 
                            +        <span class="bp">self</span><span class="o">.</span><span class="n">_count</span> <span class="o">+=</span> <span class="n">batch_size</span>
                
 
                                 <span class="nd">@property</span>
                
 
                                 <span class="k">def</span> <span class="nf">average</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
                
 
                                     <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_sum</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                                         <span class="k">return</span> <span class="mi">0</span>
                
 
                                     <span class="k">return</span> <span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">_sum</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">_count</span><span class="p">)</span><span class="o">.</span><span class="fm">__float__</span><span class="p">())</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_sum</span><span class="o">.</span><span class="
                
 
                            -            <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_sum</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">_count</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></div>
                
 
                            +            <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_sum</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">_count</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>
                
 
                                     <span class="c1"># return (self._sum / self._count).__float__() if self._sum.dim() &lt; 1 or len(self._sum) == 1 \</span>
                
 
                                     <span class="c1">#     else tuple((self._sum / self._count).cpu().numpy())</span>
                
 
                            -<div class="viewcode-block" id="tensor_container_to_device"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.tensor_container_to_device">[docs]</a><span class="k">def</span> <span class="nf">tensor_container_to_device</span><span class="p">(</span><span class="n">obj</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span 
                
 
                            +<div class="viewcode-block" id="tensor_container_to_device"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.tensor_container_to_device">[docs]</a><span class="k">def</span> <span class="nf">tensor_container_to_device</span><span class="p">(</span><span class="n">obj</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class=
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    recursively send compounded objects to device (sending all tensors to device and maintaining structure)</span>
                
 
                             <span class="sd">        :param obj           the object to send to device (list / tuple / tensor / dict)</span>
                
@@ -245,7 +259,7 @@
 
                                     <span class="k">return</span> <span class="n">obj</span></div>
                
 
                            -<div class="viewcode-block" id="get_param"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.get_param">[docs]</a><span class="k">def</span> <span class="nf">get_param</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="get_param"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.get_param">[docs]</a><span class="k">def</span> <span class="nf">get_param</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">default_val</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Retrieves a param from a parameter object/dict. If the parameter does not exist, will return default_val.</span>
                
 
                             <span class="sd">    In case the default_val is of type dictionary, and a value is found in the params - the function</span>
                
@@ -279,23 +293,23 @@
 
                                     <span class="k">return</span> <span class="n">default_val</span></div>
                
 
                            -<div class="viewcode-block" id="static_vars"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.static_vars">[docs]</a><span class="k">def</span> <span class="nf">static_vars</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">static_vars</span><span class="p">(</span><span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
                
 
                                 <span class="k">def</span> <span class="nf">decorate</span><span class="p">(</span><span class="n">func</span><span class="p">):</span>
                
 
                                     <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">kwargs</span><span class="p">:</span>
                
 
                                         <span class="nb">setattr</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">[</span><span class="n">k</span><span class="p">])</span>
                
 
                                     <span class="k">return</span> <span class="n">func</span>
                
 
                            -    <span class="k">return</span> <span class="n">decorate</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">decorate</span>
                
 
                            -<div class="viewcode-block" id="print_once"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.print_once">[docs]</a><span class="nd">@static_vars</span><span class="p">(</span><span class="n">printed</span><span class="o">=</span><span class="nb">set</span><span class="p">())</span>
                
 
                            +<span class="nd">@static_vars</span><span class="p">(</span><span class="n">printed</span><span class="o">=</span><span class="nb">set</span><span class="p">())</span>
                
 
                             <span class="k">def</span> <span class="nf">print_once</span><span class="p">(</span><span class="n">s</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
                
 
                                 <span class="k">if</span> <span class="n">s</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">print_once</span><span class="o">.</span><span class="n">printed</span><span class="p">:</span>
                
 
                                     <span class="n">print_once</span><span class="o">.</span><span class="n">printed</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">s</span><span class="p">)</span>
                
 
                            -        <span class="nb">print</span><span class="p">(</span><span class="n">s</span><span class="p">)</span></div>
                
 
                            +        <span class="nb">print</span><span class="p">(</span><span class="n">s</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="move_state_dict_to_device"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.move_state_dict_to_device">[docs]</a><span class="k">def</span> <span class="nf">move_state_dict_to_device</span><span class="p">(</span><span class="n">model_sd</span><span class="p">,</span> <span class="n">device</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">move_state_dict_to_device</span><span class="p">(</span><span class="n">model_sd</span><span class="p">,</span> <span class="n">device</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Moving model state dict tensors to target device (cuda or cpu)</span>
                
 
                             <span class="sd">    :param model_sd: model state dict</span>
                
@@ -303,10 +317,10 @@
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                                 <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">model_sd</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                
 
                                     <span class="n">model_sd</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="o">=</span> <span class="n">v</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">)</span>
                
 
                            -    <span class="k">return</span> <span class="n">model_sd</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">model_sd</span>
                
 
                            -<div class="viewcode-block" id="random_seed"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.random_seed">[docs]</a><span class="k">def</span> <span class="nf">random_seed</span><span class="p">(</span><span class="n">is_ddp</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">seed</span><span class="p">):</span>
                
 
                            +<div class="viewcode-block" id="random_seed"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.random_seed">[docs]</a><span class="k">def</span> <span class="nf">random_seed</span><span class="p">(</span><span class="n">is_ddp</span><span class="p">,</span> <span class="n">device</span><span class="p">,</span> <span class="n">seed</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Sets random seed of numpy, torch and random.</span>
                
@@ -321,7 +335,7 @@
 
                                 <span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="n">seed</span> <span class="o">+</span> <span class="n">rank</span><span class="p">)</span></div>
                
 
                            -<div class="viewcode-block" id="load_func"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.load_func">[docs]</a><span class="k">def</span> <span class="nf">load_func</span><span class="p">(</span><span class="n">dotpath</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">load_func</span><span class="p">(</span><span class="n">dotpath</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    load function in module.  function is right-most segment.</span>
                
@@ -332,10 +346,10 @@
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                                 <span class="n">module_</span><span class="p">,</span> <span class="n">func</span> <span class="o">=</span> <span class="n">dotpath</span><span class="o">.</span><span class="n">rsplit</span><span class="p">(</span><span class="s2">&quot;.&quot;</span><span class="p">,</span> <span class="n">maxsplit</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
                
 
                                 <span class="n">m</span> <span class="o">=</span> <span class="n">import_module</span><span class="p">(</span><span class="n">module_</span><span class="p">)</span>
                
 
                            -    <span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">func</span><span class="p">)</span></div>
                
 
                            +    <span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="n">func</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="get_filename_suffix_by_framework"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.get_filename_suffix_by_framework">[docs]</a><span class="k">def</span> <span class="nf">get_filename_suffix_by_framework</span><span class="p">(</span><span class="n">framework</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
                
 
                            +<span class="k">def</span> <span class="nf">get_filename_suffix_by_framework</span><span class="p">(</span><span class="n">framework</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Return the file extension of framework.</span>
                
@@ -359,10 +373,10 @@
 
                                 <span class="k">if</span> <span class="n">framework</span><span class="o">.</span><span class="n">upper</span><span class="p">()</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">frameworks_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
                
 
                                     <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Unsupported framework: </span><span class="si">{</span><span class="n">framework</span><span class="si">}</span><span class="s1">&#39;</span><span class="p">)</span>
                
 
                            -    <span class="k">return</span> <span class="n">frameworks_dict</span><span class="p">[</span><span class="n">framework</span><span class="o">.</span><span class="n">upper</span><span class="p">()]</span></div>
                
 
                            +    <span class="k">return</span> <span class="n">frameworks_dict</span><span class="p">[</span><span class="n">framework</span><span class="o">.</span><span class="n">upper</span><span class="p">()]</span>
                
 
                            -<div class="viewcode-block" id="check_models_have_same_weights"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.check_models_have_same_weights">[docs]</a><span class="k">def</span> <span class="nf">check_models_have_same_weights</span><span class="p">(</span><span class="n">model_1</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module
                
 
                            +<span class="k">def</span> <span class="nf">check_models_have_same_weights</span><span class="p">(</span><span class="n">model_1</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span> <span class="n">model_2</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Checks whether two networks have the same weights</span>
                
@@ -382,10 +396,10 @@
 
                                 <span class="k">if</span> <span class="n">models_differ</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                
 
                                     <span class="k">return</span> <span class="kc">True</span>
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
 
                            -        <span class="k">return</span> <span class="kc">False</span></div>
                
 
                            +        <span class="k">return</span> <span class="kc">False</span>
                
 
                            -<div class="viewcode-block" id="recursive_override"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.recursive_override">[docs]</a><span class="k">def</span> <span class="nf">recursive_override</span><span class="p">(</span><span class="n">base</span><span class="p">:</span> <span class="nb">dict</span><span class="p">,</span> <span class="n">extension</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span
                
 
                            +<span class="k">def</span> <span class="nf">recursive_override</span><span class="p">(</span><span class="n">base</span><span class="p">:</span> <span class="nb">dict</span><span class="p">,</span> <span class="n">extension</span><span class="p">:</span> <span class="nb">dict</span><span class="p">):</span>
                
 
                                 <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">extension</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                
 
                                     <span class="k">if</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">base</span><span class="p">:</span>
                
 
                                         <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">v</span><span class="p">,</span> <span class="n">Mapping</span><span class="p">):</span>
                
@@ -393,10 +407,10 @@
 
                                         <span class="k">else</span><span class="p">:</span>
                
 
                                             <span class="n">base</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="o">=</span> <span class="n">extension</span><span class="p">[</span><span class="n">k</span><span class="p">]</span>
                
 
                                     <span class="k">else</span><span class="p">:</span>
                
 
                            -            <span class="n">base</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="o">=</span> <span class="n">extension</span><span class="p">[</span><span class="n">k</span><span class="p">]</span></div>
                
 
                            +            <span class="n">base</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="o">=</span> <span class="n">extension</span><span class="p">[</span><span class="n">k</span><span class="p">]</span>
                
 
                            -<div class="viewcode-block" id="download_and_unzip_from_url"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.download_and_unzip_from_url">[docs]</a><span class="k">def</span> <span class="nf">download_and_unzip_from_url</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="nb">dir</span><span class="o">=</span><span class="s1">&#39;.&#39;</span><span class="p">,</span> <span class="n">unzip</sp
                
 
                            +<span class="k">def</span> <span class="nf">download_and_unzip_from_url</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="nb">dir</span><span class="o">=</span><span class="s1">&#39;.&#39;</span><span class="p">,</span> <span class="n">unzip</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">delete</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Downloads a zip file from url to dir, and unzips it.</span>
                
@@ -431,10 +445,10 @@
 
                                 <span class="nb">dir</span> <span class="o">=</span> <span class="n">Path</span><span class="p">(</span><span class="nb">dir</span><span class="p">)</span>
                
 
                                 <span class="nb">dir</span><span class="o">.</span><span class="n">mkdir</span><span class="p">(</span><span class="n">parents</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">exist_ok</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>  <span class="c1"># make directory</span>
                
 
                                 <span class="k">for</span> <span class="n">u</span> <span class="ow">in</span> <span class="p">[</span><span class="n">url</span><span class="p">]</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">url</span><span class="p">,</span> <span class="p">(</span><span class="nb">str</span><span class="p">,</span> <span class="n">Path</span><span class="p">))</span> <span class="k">else</span> <span class="n">url</span><span class="p">:</span
                
 
                            -        <span class="n">download_one</span><span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="nb">dir</span><span class="p">)</span></div>
                
 
                            +        <span class="n">download_one</span><span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="nb">dir</span><span class="p">)</span>
                
 
                            -<div class="viewcode-block" id="download_and_untar_from_url"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.download_and_untar_from_url">[docs]</a><span class="k">def</span> <span class="nf">download_and_untar_from_url</span><span class="p">(</span><span class="n">urls</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="nb">dir</span><spa
                
 
                            +<span class="k">def</span> <span class="nf">download_and_untar_from_url</span><span class="p">(</span><span class="n">urls</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span> <span class="nb">dir</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Path</span><span class="p">]</span> <span class="o">=</span> <span cl
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Download a file from url and untar.</span>
                
@@ -460,10 +474,10 @@
 
                                     <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="sa">f</span><span class="s1">&#39;Extracting to </span><span class="si">{</span><span class="nb">dir</span><span class="si">}</span><span class="s1">...&#39;</span><span class="p">)</span>
                
 
                                     <span class="k">with</span> <span class="n">tarfile</span><span class="o">.</span><span class="n">open</span><span class="p">(</span><span class="n">filepath</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="n">modes</span><span class="p">[</span><span class="n">filepath</span><span class="o">.</span><span class="n">suffix</span><span class="p">])</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
                
 
                                         <span class="n">f</span><span class="o">.</span><span class="n">extractall</span><span class="p">(</span><span class="nb">dir</span><span class="p">)</span>
                
 
                            -        <span class="n">filepath</span><span class="o">.</span><span class="n">unlink</span><span class="p">()</span></div>
                
 
                            +        <span class="n">filepath</span><span class="o">.</span><span class="n">unlink</span><span class="p">()</span>
                
 
                            -<div class="viewcode-block" id="make_divisible"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.make_divisible">[docs]</a><span class="k">def</span> <span class="nf">make_divisible</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">divisor</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">ce
                
 
                            +<span class="k">def</span> <span class="nf">make_divisible</span><span class="p">(</span><span class="n">x</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">divisor</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">ceil</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    Returns x evenly divisible by divisor.</span>
                
 
                             <span class="sd">    If ceil=True it will return the closest larger number to the original x, and ceil=False the closest smaller number.</span>
                
@@ -471,10 +485,10 @@
 
                                 <span class="k">if</span> <span class="n">ceil</span><span class="p">:</span>
                
 
                                     <span class="k">return</span> <span class="n">math</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="n">x</span> <span class="o">/</span> <span class="n">divisor</span><span class="p">)</span> <span class="o">*</span> <span class="n">divisor</span>
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
 
                            -        <span class="k">return</span> <span class="n">math</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">x</span> <span class="o">/</span> <span class="n">divisor</span><span class="p">)</span> <span class="o">*</span> <span class="n">divisor</span></div>
                
 
                            +        <span class="k">return</span> <span class="n">math</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">x</span> <span class="o">/</span> <span class="n">divisor</span><span class="p">)</span> <span class="o">*</span> <span class="n">divisor</span>
                
 
                            -<div class="viewcode-block" id="check_img_size_divisibility"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.check_img_size_divisibility">[docs]</a><span class="k">def</span> <span class="nf">check_img_size_divisibility</span><span class="p">(</span><span class="n">img_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">stride</span><span class="p">:</span> <span class="nb">int</spa
                
 
                            +<span class="k">def</span> <span class="nf">check_img_size_divisibility</span><span class="p">(</span><span class="n">img_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">stride</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">32</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">bool</span><span class="p">,</spa
                
 
                                 <span class="sd">&quot;&quot;&quot;</span>
                
 
                             <span class="sd">    :param img_size: Int, the size of the image (H or W).</span>
                
 
                             <span class="sd">    :param stride: Int, the number to check if img_size is divisible by.</span>
                
@@ -486,22 +500,22 @@
 
                                 <span class="k">if</span> <span class="n">new_size</span> <span class="o">!=</span> <span class="n">img_size</span><span class="p">:</span>
                
 
                                     <span class="k">return</span> <span class="kc">False</span><span class="p">,</span> <span class="p">(</span><span class="n">new_size</span><span class="p">,</span> <span class="n">make_divisible</span><span class="p">(</span><span class="n">img_size</span><span class="p">,</span> <span class="nb">int</span><span class="p">(</span><span class="n">stride</span><span class="p">),</span> <span class="n">ceil</span><span class="o">=</span><span class="kc">False</span><span class="p">))</span
                
 
                                 <span class="k">else</span><span class="p">:</span>
                
 
                            -        <span class="k">return</span> <span class="kc">True</span><span class="p">,</span> <span class="kc">None</span></div>
                
 
                            +        <span class="k">return</span> <span class="kc">True</span><span class="p">,</span> <span class="kc">None</span>
                
 
                            -<div class="viewcode-block" id="get_orientation_key"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.get_orientation_key">[docs]</a><span class="nd">@lru_cache</span><span class="p">(</span><span class="kc">None</span><span class="p">)</span>
                
 
                            +<span class="nd">@lru_cache</span><span class="p">(</span><span class="kc">None</span><span class="p">)</span>
                
 
                             <span class="k">def</span> <span class="nf">get_orientation_key</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Get the orientation key according to PIL, which is useful to get the image size for instance</span>
                
 
                             <span class="sd">    :return: Orientation key according to PIL&quot;&quot;&quot;</span>
                
 
                                 <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">ExifTags</span><span class="o">.</span><span class="n">TAGS</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                
 
                                     <span class="k">if</span> <span class="n">value</span> <span class="o">==</span> <span class="s1">&#39;Orientation&#39;</span><span class="p">:</span>
                
 
                            -            <span class="k">return</span> <span class="n">key</span></div>
                
 
                            +            <span class="k">return</span> <span class="n">key</span>
                
 
                            -<div class="viewcode-block" id="exif_size"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.exif_size">[docs]</a><span class="k">def</span> <span class="nf">exif_size</span><span class="p">(</span><span class="n">image</span><span class="p">:</span> <span class="n">Image</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <s
                
 
                            +<span class="k">def</span> <span class="nf">exif_size</span><span class="p">(</span><span class="n">image</span><span class="p">:</span> <span class="n">Image</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Get the size of image.</span>
                
 
                             <span class="sd">    :param image:   The image to get size from</span>
                
 
                            -<span class="sd">    :return:        (width, height)</span>
                
 
                            +<span class="sd">    :return:        (height, width)</span>
                
 
                             <span class="sd">    &quot;&quot;&quot;</span>
                
 
                                 <span class="n">orientation_key</span> <span class="o">=</span> <span class="n">get_orientation_key</span><span class="p">()</span>
                
@@ -519,14 +533,27 @@
 
                                             <span class="n">image_size</span> <span class="o">=</span> <span class="p">(</span><span class="n">image_size</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">image_size</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
                
 
                                 <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">ex</span><span class="p">:</span>
                
 
                                     <span class="nb">print</span><span class="p">(</span><span class="s1">&#39;Caught Exception trying to rotate: &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">image</span><span class="p">)</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">ex</span><span class="p">))</span>
                
 
                            -    <span class="n">height</span><span class="p">,</span> <span class="n">width</span> <span class="o">=</span> <span class="n">image_size</span>
                
 
                            -    <span class="k">return</span> <span class="n">width</span><span class="p">,</span> <span class="n">height</span></div>
                
 
                            +    <span class="n">width</span><span class="p">,</span> <span class="n">height</span> <span class="o">=</span> <span class="n">image_size</span>
                
 
                            +    <span class="k">return</span> <span class="n">height</span><span class="p">,</span> <span class="n">width</span>
                
 
                            -<div class="viewcode-block" id="get_image_size_from_path"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.get_image_size_from_path">[docs]</a><span class="k">def</span> <span class="nf">get_image_size_from_path</span><span class="p">(</span><span class="n">img_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span c
                
 
                            +<span class="k">def</span> <span class="nf">get_image_size_from_path</span><span class="p">(</span><span class="n">img_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]:</span>
                
 
                                 <span class="sd">&quot;&quot;&quot;Get the image size of an image at a specific path&quot;&quot;&quot;</span>
                
 
                                 <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">img_path</span><span class="p">,</span> <span class="s1">&#39;rb&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
                
 
                            -        <span class="k">return</span> <span class="n">exif_size</span><span class="p">(</span><span class="n">Image</span><span class="o">.</span><span class="n">open</span><span class="p">(</span><span class="n">f</span><span class="p">))</span></div>
                
 
                            +        <span class="k">return</span> <span class="n">exif_size</span><span class="p">(</span><span class="n">Image</span><span class="o">.</span><span class="n">open</span><span class="p">(</span><span class="n">f</span><span class="p">))</span>
                
 
                            +
                
 
                            +
                
 
                            +<span class="k">def</span> <span class="nf">override_default_params_without_nones</span><span class="p">(</span><span class="n">params</span><span class="p">:</span> <span class="n">Dict</span><span class="p">,</span> <span class="n">default_params</span><span class="p">:</span> <span class="n">Dict</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Dict</span><span class="p">:</span>
                
 
                            +    <span class="sd">&quot;&quot;&quot;</span>
                
 
                            +<span class="sd">    Helper method for overriding default dictionary&#39;s entries excluding entries with None values.</span>
                
 
                            +<span class="sd">    :param params: dict, output dictionary which will take the defaults.</span>
                
 
                            +<span class="sd">    :param default_params: dict, dictionary for the defaults.</span>
                
 
                            +<span class="sd">    :return: dict, params after manipulation,</span>
                
 
                            +<span class="sd">    &quot;&quot;&quot;</span>
                
 
                            +    <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">default_params</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                
 
                            +        <span class="k">if</span> <span class="n">key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">params</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span> <span class="ow">or</span> <span class="n">params</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                
 
                            +            <span class="n">params</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">val</span>
                
 
                            +    <span class="k">return</span> <span class="n">params</span>
                
 
                             </pre></div>
                
 
                                        </div>
                
@@ -556,4 +583,4 @@
 
                               </script> 
                
 
                             </body>
                
 
                            -</html>
                
 
                            +</html>
                
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.version_utils &mdash; SuperGradients 3.0.3 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/sphinx_highlight.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption" role="heading"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Version 3 is out! Notebooks have been updated!</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#build-with-supergradients">Build with SuperGradients</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#quick-installation">Quick Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#what-s-new">What’s New</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#coming-soon">Coming soon</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#table-of-content">Table of Content</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#getting-started">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#advanced-features">Advanced Features</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#installation-methods">Installation Methods</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#implemented-model-architectures">Implemented Model Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#documentation">Documentation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#contributing">Contributing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#citation">Citation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#community">Community</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#license">License</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#deci-platform">Deci Platform</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.version_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.version_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>

<span class="n">_TORCH_VERSION_MAJOR_MINOR_</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">version</span><span class="o">.</span><span class="n">__version__</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;.&quot;</span><span class="p">)[:</span><span class="mi">2</span><span class="p">]))</span>

<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;torch_version_is_greater_or_equal&quot;</span><span class="p">]</span>


<div class="viewcode-block" id="torch_version_is_greater_or_equal"><a class="viewcode-back" href="../../../../super_gradients.training.html#super_gradients.training.utils.torch_version_is_greater_or_equal">[docs]</a><span class="k">def</span> <span class="nf">torch_version_is_greater_or_equal</span><span class="p">(</span><span class="n">major</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">minor</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
    <span class="n">version</span> <span class="o">=</span> <span class="p">(</span><span class="n">major</span><span class="p">,</span> <span class="n">minor</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">_TORCH_VERSION_MAJOR_MINOR_</span> <span class="o">&gt;=</span> <span class="n">version</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            <!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>super_gradients.training.utils.weight_averaging_utils &mdash; SuperGradients 1.0 documentation</title>
      <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
      <link rel="stylesheet" href="../../../../_static/graphviz.css" type="text/css" />
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
        <script data-url_root="../../../../" id="documentation_options" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
    <script src="../../../../_static/js/theme.js"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav"> 
  <div class="wy-grid-for-nav">
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
            <a href="../../../../index.html" class="icon icon-home"> SuperGradients
          </a>
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>
        </div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
              <p class="caption"><span class="caption-text">Welcome To SuperGradients</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html">Fill our 4-question quick survey! We will raffle free SuperGradients swag between those who will participate -&gt; Fill Survey</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../welcome.html#supergradients">SuperGradients</a></li>
</ul>
<p class="caption"><span class="caption-text">Technical Documentation</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.common.html">Common package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../super_gradients.training.html">Training package</a></li>
</ul>
<p class="caption"><span class="caption-text">User Guide</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html">What is SuperGradients?</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#introducing-the-supergradients-library">Introducing the SuperGradients library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#installation">Installation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#integrating-your-training-code-complete-walkthrough">Integrating Your Training Code - Complete Walkthrough</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#training-parameters">Training Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#logs-and-checkpoints">Logs and Checkpoints</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#dataset-parameters">Dataset Parameters</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#network-architectures">Network Architectures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#pretrained-models">Pretrained Models</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#how-to-reproduce-our-training-recipes">How To Reproduce Our Training Recipes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#professional-tools-integration">Professional Tools Integration</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../user_guide.html#supergradients-faq">SuperGradients FAQ</a></li>
</ul>

        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">SuperGradients</a>
      </nav>

      <div class="wy-nav-content">
        <div class="rst-content">
          <div role="navigation" aria-label="Page navigation">
  <ul class="wy-breadcrumbs">
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
      <li>super_gradients.training.utils.weight_averaging_utils</li>
      <li class="wy-breadcrumbs-aside">
      </li>
  </ul>
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
             
  <h1>Source code for super_gradients.training.utils.weight_averaging_utils</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">pkg_resources</span>
<span class="kn">from</span> <span class="nn">super_gradients.training</span> <span class="kn">import</span> <span class="n">utils</span> <span class="k">as</span> <span class="n">core_utils</span>
<span class="kn">from</span> <span class="nn">super_gradients.training.utils.utils</span> <span class="kn">import</span> <span class="n">move_state_dict_to_device</span>


<div class="viewcode-block" id="ModelWeightAveraging"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.weight_averaging_utils.ModelWeightAveraging">[docs]</a><span class="k">class</span> <span class="nc">ModelWeightAveraging</span><span class="p">:</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Utils class for managing the averaging of the best several snapshots into a single model.</span>
<span class="sd">    A snapshot dictionary file and the average model will be saved / updated at every epoch and evaluated only when</span>
<span class="sd">    training is completed. The snapshot file will only be deleted upon completing the training.</span>
<span class="sd">    The snapshot dict will be managed on cpu.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ckpt_dir</span><span class="p">,</span>
                 <span class="n">greater_is_better</span><span class="p">,</span>
                 <span class="n">source_ckpt_folder_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">metric_to_watch</span><span class="o">=</span><span class="s1">&#39;acc&#39;</span><span class="p">,</span>
                 <span class="n">metric_idx</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">load_checkpoint</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                 <span class="n">number_of_models_to_average</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
                 <span class="n">model_checkpoints_location</span><span class="o">=</span><span class="s1">&#39;local&#39;</span>
                 <span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Init the ModelWeightAveraging</span>
<span class="sd">        :param checkpoint_dir: the directory where the checkpoints are saved</span>
<span class="sd">        :param metric_to_watch: monitoring loss or acc, will be identical to that which determines best_model</span>
<span class="sd">        :param metric_idx:</span>
<span class="sd">        :param load_checkpoint: whether to load pre-existing snapshot dict.</span>
<span class="sd">        :param number_of_models_to_average: number of models to average</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">if</span> <span class="n">source_ckpt_folder_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">source_ckpt_file</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">source_ckpt_folder_name</span><span class="p">,</span> <span class="s1">&#39;averaging_snapshots.pkl&#39;</span><span class="p">)</span>
            <span class="n">source_ckpt_file</span> <span class="o">=</span> <span class="n">pkg_resources</span><span class="o">.</span><span class="n">resource_filename</span><span class="p">(</span><span class="s1">&#39;checkpoints&#39;</span><span class="p">,</span> <span class="n">source_ckpt_file</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">averaging_snapshots_file</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">ckpt_dir</span><span class="p">,</span> <span class="s1">&#39;averaging_snapshots.pkl&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">number_of_models_to_average</span> <span class="o">=</span> <span class="n">number_of_models_to_average</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">metric_to_watch</span> <span class="o">=</span> <span class="n">metric_to_watch</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">metric_idx</span> <span class="o">=</span> <span class="n">metric_idx</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="o">=</span> <span class="n">greater_is_better</span>

        <span class="c1"># if continuing training, copy previous snapshot dict if exist</span>
        <span class="k">if</span> <span class="n">load_checkpoint</span> <span class="ow">and</span> <span class="n">source_ckpt_folder_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">isfile</span><span class="p">(</span><span class="n">source_ckpt_file</span><span class="p">):</span>
            <span class="n">averaging_snapshots_dict</span> <span class="o">=</span> <span class="n">core_utils</span><span class="o">.</span><span class="n">load_checkpoint</span><span class="p">(</span><span class="n">ckpt_destination_dir</span><span class="o">=</span><span class="n">ckpt_dir</span><span class="p">,</span>
                                                                  <span class="n">source_ckpt_folder_name</span><span class="o">=</span><span class="n">source_ckpt_folder_name</span><span class="p">,</span>
                                                                  <span class="n">ckpt_filename</span><span class="o">=</span><span class="s2">&quot;averaging_snapshots.pkl&quot;</span><span class="p">,</span>
                                                                  <span class="n">load_weights_only</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                                                                  <span class="n">model_checkpoints_location</span><span class="o">=</span><span class="n">model_checkpoints_location</span><span class="p">,</span>
                                                                  <span class="n">overwrite_local_ckpt</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

        <span class="k">else</span><span class="p">:</span>
            <span class="n">averaging_snapshots_dict</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;snapshot&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">i</span><span class="p">):</span> <span class="kc">None</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">number_of_models_to_average</span><span class="p">)}</span>
            <span class="c1"># if metric to watch is acc, hold a zero array, if loss hold inf array</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span><span class="p">:</span>
                <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshots_metric&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">inf</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">number_of_models_to_average</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshots_metric&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">inf</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">number_of_models_to_average</span><span class="p">)</span>

        <span class="n">torch</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">averaging_snapshots_dict</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">averaging_snapshots_file</span><span class="p">)</span>

<div class="viewcode-block" id="ModelWeightAveraging.update_snapshots_dict"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.weight_averaging_utils.ModelWeightAveraging.update_snapshots_dict">[docs]</a>    <span class="k">def</span> <span class="nf">update_snapshots_dict</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Update the snapshot dict and returns the updated average model for saving</span>
<span class="sd">        :param model: the latest model</span>
<span class="sd">        :param validation_results_tuple: performance of the latest model</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">averaging_snapshots_dict</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_averaging_snapshots_dict</span><span class="p">()</span>

        <span class="c1"># IF CURRENT MODEL IS BETTER, TAKING HIS PLACE IN ACC LIST AND OVERWRITE THE NEW AVERAGE</span>
        <span class="n">require_update</span><span class="p">,</span> <span class="n">update_ind</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_is_better</span><span class="p">(</span><span class="n">averaging_snapshots_dict</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">require_update</span><span class="p">:</span>
            <span class="c1"># moving state dict to cpu</span>
            <span class="n">new_sd</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">state_dict</span><span class="p">()</span>
            <span class="n">new_sd</span> <span class="o">=</span> <span class="n">move_state_dict_to_device</span><span class="p">(</span><span class="n">new_sd</span><span class="p">,</span> <span class="s1">&#39;cpu&#39;</span><span class="p">)</span>

            <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshot&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">update_ind</span><span class="p">)]</span> <span class="o">=</span> <span class="n">new_sd</span>
            <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshots_metric&#39;</span><span class="p">][</span><span class="n">update_ind</span><span class="p">]</span> <span class="o">=</span> <span class="n">validation_results_tuple</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_idx</span><span class="p">]</span>

        <span class="k">return</span> <span class="n">averaging_snapshots_dict</span></div>

<div class="viewcode-block" id="ModelWeightAveraging.get_average_model"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.weight_averaging_utils.ModelWeightAveraging.get_average_model">[docs]</a>    <span class="k">def</span> <span class="nf">get_average_model</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Returns the averaged model</span>
<span class="sd">        :param model: will be used to determine arch</span>
<span class="sd">        :param validation_results_tuple: if provided, will update the average model before returning</span>
<span class="sd">        :param target_device: if provided, return sd on target device</span>

<span class="sd">        &quot;&quot;&quot;</span>
        <span class="c1"># If validation tuple is provided, update the average model</span>
        <span class="k">if</span> <span class="n">validation_results_tuple</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">averaging_snapshots_dict</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">update_snapshots_dict</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">averaging_snapshots_dict</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_averaging_snapshots_dict</span><span class="p">()</span>

        <span class="n">torch</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">averaging_snapshots_dict</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">averaging_snapshots_file</span><span class="p">)</span>
        <span class="n">average_model_sd</span> <span class="o">=</span> <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshot0&#39;</span><span class="p">]</span>
        <span class="k">for</span> <span class="n">n_model</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">number_of_models_to_average</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshot&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">n_model</span><span class="p">)]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">net_sd</span> <span class="o">=</span> <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshot&#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">n_model</span><span class="p">)]</span>
                <span class="c1"># USING MOVING AVERAGE</span>
                <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="n">average_model_sd</span><span class="p">:</span>
                    <span class="n">average_model_sd</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">true_divide</span><span class="p">(</span>
                        <span class="n">average_model_sd</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">*</span> <span class="n">n_model</span> <span class="o">+</span> <span class="n">net_sd</span><span class="p">[</span><span class="n">key</span><span class="p">],</span>
                        <span class="p">(</span><span class="n">n_model</span> <span class="o">+</span> <span class="mi">1</span><span class="p">))</span>

        <span class="k">return</span> <span class="n">average_model_sd</span></div>

<div class="viewcode-block" id="ModelWeightAveraging.cleanup"><a class="viewcode-back" href="../../../../super_gradients.training.utils.html#super_gradients.training.utils.weight_averaging_utils.ModelWeightAveraging.cleanup">[docs]</a>    <span class="k">def</span> <span class="nf">cleanup</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Delete snapshot file when reaching the last epoch</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">os</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">averaging_snapshots_file</span><span class="p">)</span></div>

    <span class="k">def</span> <span class="nf">_is_better</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">averaging_snapshots_dict</span><span class="p">,</span> <span class="n">validation_results_tuple</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Determines if the new model is better according to the specified metrics</span>
<span class="sd">        :param averaging_snapshots_dict: snapshot dict</span>
<span class="sd">        :param validation_results_tuple: latest model performance</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">snapshot_metric_array</span> <span class="o">=</span> <span class="n">averaging_snapshots_dict</span><span class="p">[</span><span class="s1">&#39;snapshots_metric&#39;</span><span class="p">]</span>
        <span class="n">val</span> <span class="o">=</span> <span class="n">validation_results_tuple</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">metric_idx</span><span class="p">]</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span><span class="p">:</span>
            <span class="n">update_ind</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">argmin</span><span class="p">(</span><span class="n">snapshot_metric_array</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">update_ind</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="n">snapshot_metric_array</span><span class="p">)</span>

        <span class="k">if</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="ow">and</span> <span class="n">val</span> <span class="o">&gt;</span> <span class="n">snapshot_metric_array</span><span class="p">[</span><span class="n">update_ind</span><span class="p">])</span> <span class="ow">or</span> <span class="p">(</span>
                <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater_is_better</span> <span class="ow">and</span> <span class="n">val</span> <span class="o">&lt;</span> <span class="n">snapshot_metric_array</span><span class="p">[</span><span class="n">update_ind</span><span class="p">]):</span>
            <span class="k">return</span> <span class="kc">True</span><span class="p">,</span> <span class="n">update_ind</span>

        <span class="k">return</span> <span class="kc">False</span><span class="p">,</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="nf">_get_averaging_snapshots_dict</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">averaging_snapshots_file</span><span class="p">)</span></div>
</pre></div>

           </div>
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>&#169; Copyright 2021, SuperGradients team.</p>
  </div>

  Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    provided by <a href="https://readthedocs.org">Read the Docs</a>.
   

</footer>
        </div>
      </div>
    </section>
  </div>
  <script>
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script> 

</body>
</html>

          
 
            super\_gradients.common.abstractions
====================================

.. automodule:: super_gradients.common.abstractions

   
   
   

   
   
   

   
   
   

   
   
   




          
 
            super\_gradients.common.auto\_logging
=====================================

.. automodule:: super_gradients.common.auto_logging

   
   
   

   
   
   

   
   
   

   
   
   




          
 
            super\_gradients.common.aws\_connection
=======================================

.. automodule:: super_gradients.common.aws_connection

   
   
   

   
   
   

   
   
   

   
   
   




          
 
            super\_gradients.common.data\_connection
========================================

.. automodule:: super_gradients.common.data_connection

   
   
   

   
   
   

   
   
   

   
   
   




          
 
            super\_gradients.common.data\_interface
=======================================

.. automodule:: super_gradients.common.data_interface

   
   
   

   
   
   

   
   
   

   
   
   




          
 
            super\_gradients.common.data\_types
===================================

.. automodule:: super_gradients.common.data_types

   
   
   

   
   
   

   
   
   

   
   
   




          
 
            super\_gradients.common.decorators
==================================

.. automodule:: super_gradients.common.decorators

   
   
   

   
   
   

   
   
   

   
   
   




          
 
            super\_gradients.common.environment
===================================

.. automodule:: super_gradients.common.environment

   
   
   

   
   
   

   
   
   

   
   
   




          
@@ -20,10 +20,8 @@ Welcome to SuperGradients's documentation!
 
                                super_gradients.training
                
 
                             .. toctree::
                
 
                            -   :maxdepth: 4
                
 
                            -   :caption: User Guide
                
 
                            -
                
 
                            -   user_guide
                
 
                            +.. :maxdepth: 4
                
 
                            +.. :caption: User Guide
                
 
                             Indices and tables
                
 
                             ==================
                
 
            super_gradients
===============

.. toctree::
   :maxdepth: 8

   super_gradients

          
 
            super\_gradients.common.abstractions package
============================================

Submodules
----------

super\_gradients.common.abstractions.abstract\_logger module
------------------------------------------------------------

.. automodule:: super_gradients.common.abstractions.abstract_logger
   :members:
   :undoc-members:
   :show-inheritance:

Module contents
---------------

.. automodule:: super_gradients.common.abstractions
   :members:
   :undoc-members:
   :show-inheritance:

          
 
            super\_gradients.common.auto\_logging package
=============================================

Submodules
----------

super\_gradients.common.auto\_logging.auto\_logger module
---------------------------------------------------------

.. automodule:: super_gradients.common.auto_logging.auto_logger
   :members:
   :undoc-members:
   :show-inheritance:

Module contents
---------------

.. automodule:: super_gradients.common.auto_logging
   :members:
   :undoc-members:
   :show-inheritance: