deploy: SPFlow/SPFlow@fb796f4

Author: braun-steven

.doctrees/environment.pickle

@@ -286,6 +286,21 @@ <h1>Source code for spflow.learn.expectation_maximization</h1><div class="highli
 <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="vm">__name__</span><span class="p">)</span>
 
 
+<span class="k">def</span><span class="w"> </span><span class="nf">_retain_cached_log_likelihood_grads</span><span class="p">(</span><span class="n">cache</span><span class="p">:</span> <span class="n">Cache</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Retain gradients for cached non-leaf likelihood tensors consumed by EM.&quot;&quot;&quot;</span>
+    <span class="k">for</span> <span class="n">lls</span> <span class="ow">in</span> <span class="n">cache</span><span class="p">[</span><span class="s2">&quot;log_likelihood&quot;</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">torch</span><span class="o">.</span><span class="n">is_tensor</span><span class="p">(</span><span class="n">lls</span><span class="p">)</span> <span class="ow">and</span> <span class="n">lls</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
+            <span class="n">lls</span><span class="o">.</span><span class="n">retain_grad</span><span class="p">()</span>
+
+
+<span class="k">def</span><span class="w"> </span><span class="nf">_backward_accumulated_log_likelihood</span><span class="p">(</span><span class="n">acc_ll</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Backpropagate one EM step without retaining the graph.&quot;&quot;&quot;</span>
+    <span class="k">if</span> <span class="ow">not</span> <span class="n">acc_ll</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
+        <span class="k">return</span>
+
+    <span class="n">acc_ll</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
+
+
 <div class="viewcode-block" id="expectation_maximization">
 <a class="viewcode-back" href="../../../api/learning.html#spflow.learn.expectation_maximization.expectation_maximization">[docs]</a>
 <span class="k">def</span><span class="w"> </span><span class="nf">expectation_maximization</span><span class="p">(</span>
@@ -327,14 +342,8 @@ <h1>Source code for spflow.learn.expectation_maximization</h1><div class="highli
         <span class="k">if</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="sa">f</span><span class="s2">&quot;Step </span><span class="si">{</span><span class="n">step</span><span class="si">}</span><span class="s2">: Average log-likelihood: </span><span class="si">{</span><span class="n">avg_ll</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
 
-        <span class="c1"># retain gradients for all module log-likelihoods</span>
-        <span class="k">for</span> <span class="n">lls</span> <span class="ow">in</span> <span class="n">cache</span><span class="p">[</span><span class="s2">&quot;log_likelihood&quot;</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">torch</span><span class="o">.</span><span class="n">is_tensor</span><span class="p">(</span><span class="n">lls</span><span class="p">)</span> <span class="ow">and</span> <span class="n">lls</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
-                <span class="n">lls</span><span class="o">.</span><span class="n">retain_grad</span><span class="p">()</span>
-
-        <span class="c1"># compute gradients (if there are differentiable parameters to begin with)</span>
-        <span class="k">if</span> <span class="n">acc_ll</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
-            <span class="n">acc_ll</span><span class="o">.</span><span class="n">backward</span><span class="p">(</span><span class="n">retain_graph</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+        <span class="n">_retain_cached_log_likelihood_grads</span><span class="p">(</span><span class="n">cache</span><span class="p">)</span>
+        <span class="n">_backward_accumulated_log_likelihood</span><span class="p">(</span><span class="n">acc_ll</span><span class="p">)</span>
 
         <span class="c1"># recursively perform expectation maximization</span>
         <span class="n">module</span><span class="o">.</span><span class="n">_expectation_maximization_step</span><span class="p">(</span>
@@ -393,12 +402,8 @@ <h1>Source code for spflow.learn.expectation_maximization</h1><div class="highli
             <span class="n">epoch_ll</span> <span class="o">=</span> <span class="n">epoch_ll</span> <span class="o">+</span> <span class="n">acc_ll</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span>
             <span class="n">num_samples</span> <span class="o">+=</span> <span class="n">batch_data</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">lls</span> <span class="ow">in</span> <span class="n">cache</span><span class="p">[</span><span class="s2">&quot;log_likelihood&quot;</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">torch</span><span class="o">.</span><span class="n">is_tensor</span><span class="p">(</span><span class="n">lls</span><span class="p">)</span> <span class="ow">and</span> <span class="n">lls</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
-                    <span class="n">lls</span><span class="o">.</span><span class="n">retain_grad</span><span class="p">()</span>
-
-            <span class="k">if</span> <span class="n">acc_ll</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
-                <span class="n">acc_ll</span><span class="o">.</span><span class="n">backward</span><span class="p">(</span><span class="n">retain_graph</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+            <span class="n">_retain_cached_log_likelihood_grads</span><span class="p">(</span><span class="n">cache</span><span class="p">)</span>
+            <span class="n">_backward_accumulated_log_likelihood</span><span class="p">(</span><span class="n">acc_ll</span><span class="p">)</span>
 
             <span class="n">module</span><span class="o">.</span><span class="n">_expectation_maximization_step</span><span class="p">(</span>
                 <span class="n">data</span><span class="o">=</span><span class="n">batch_data</span><span class="p">,</span>