deploy: SimplexLab/TorchJD@f677df5

Author: ValerianRey

latest/_sources/examples/monitoring.rst.txt

@@ -19,8 +19,8 @@ they have a negative inner product).
 
     import torch
     from torch.nn import Linear, MSELoss, ReLU, Sequential
-    from torch.optim import SGD
     from torch.nn.functional import cosine_similarity
+    from torch.optim import SGD
 
     from torchjd import mtl_backward
     from torchjd.aggregation import UPGrad
@@ -29,7 +29,7 @@ they have a negative inner product).
         """Prints the extracted weights."""
         print(f"Weights: {weights}")
 
-    def print_similarity_with_gd(_, inputs: tuple[torch.Tensor], aggregation: torch.Tensor) -> None:
+    def print_gd_similarity(_, inputs: tuple[torch.Tensor, ...], aggregation: torch.Tensor) -> None:
         """Prints the cosine similarity between the aggregation and the average gradient."""
         matrix = inputs[0]
         gd_output = matrix.mean(dim=0)
@@ -50,7 +50,7 @@ they have a negative inner product).
     aggregator = UPGrad()
 
     aggregator.weighting.register_forward_hook(print_weights)
-    aggregator.register_forward_hook(print_similarity_with_gd)
+    aggregator.register_forward_hook(print_gd_similarity)
 
     inputs = torch.randn(8, 16, 10)  # 8 batches of 16 random input vectors of length 10
     task1_targets = torch.randn(8, 16, 1)  # 8 batches of 16 targets for the first task

latest/examples/monitoring/index.html

@@ -253,8 +253,8 @@ <h1>Monitoring aggregations<a class="headerlink" href="#monitoring-aggregations"
 they have a negative inner product).</p>
 <div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">torch.nn</span><span class="w"> </span><span class="kn">import</span> <span class="n">Linear</span><span class="p">,</span> <span class="n">MSELoss</span><span class="p">,</span> <span class="n">ReLU</span><span class="p">,</span> <span class="n">Sequential</span>
-<span class="kn">from</span><span class="w"> </span><span class="nn">torch.optim</span><span class="w"> </span><span class="kn">import</span> <span class="n">SGD</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">torch.nn.functional</span><span class="w"> </span><span class="kn">import</span> <span class="n">cosine_similarity</span>
+<span class="kn">from</span><span class="w"> </span><span class="nn">torch.optim</span><span class="w"> </span><span class="kn">import</span> <span class="n">SGD</span>
 
 <span class="kn">from</span><span class="w"> </span><span class="nn">torchjd</span><span class="w"> </span><span class="kn">import</span> <span class="n">mtl_backward</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">torchjd.aggregation</span><span class="w"> </span><span class="kn">import</span> <span class="n">UPGrad</span>
@@ -263,7 +263,7 @@ <h1>Monitoring aggregations<a class="headerlink" href="#monitoring-aggregations"
 </span><span class="hll"><span class="w">    </span><span class="sd">&quot;&quot;&quot;Prints the extracted weights.&quot;&quot;&quot;</span>
 </span><span class="hll">    <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Weights: </span><span class="si">{</span><span class="n">weights</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
 </span>
-<span class="hll"><span class="k">def</span><span class="w"> </span><span class="nf">print_similarity_with_gd</span><span class="p">(</span><span class="n">_</span><span class="p">,</span> <span class="n">inputs</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">aggregation</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="kc">None</span><span class="p">:</span>
+<span class="hll"><span class="k">def</span><span class="w"> </span><span class="nf">print_gd_similarity</span><span class="p">(</span><span class="n">_</span><span class="p">,</span> <span class="n">inputs</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="o">...</span><span class="p">],</span> <span class="n">aggregation</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="kc">None</span><span class="p">:</span>
 </span><span class="hll"><span class="w">    </span><span class="sd">&quot;&quot;&quot;Prints the cosine similarity between the aggregation and the average gradient.&quot;&quot;&quot;</span>
 </span><span class="hll">    <span class="n">matrix</span> <span class="o">=</span> <span class="n">inputs</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
 </span><span class="hll">    <span class="n">gd_output</span> <span class="o">=</span> <span class="n">matrix</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
@@ -284,7 +284,7 @@ <h1>Monitoring aggregations<a class="headerlink" href="#monitoring-aggregations"
 <span class="n">aggregator</span> <span class="o">=</span> <span class="n">UPGrad</span><span class="p">()</span>
 
 <span class="hll"><span class="n">aggregator</span><span class="o">.</span><span class="n">weighting</span><span class="o">.</span><span class="n">register_forward_hook</span><span class="p">(</span><span class="n">print_weights</span><span class="p">)</span>
-</span><span class="hll"><span class="n">aggregator</span><span class="o">.</span><span class="n">register_forward_hook</span><span class="p">(</span><span class="n">print_similarity_with_gd</span><span class="p">)</span>
+</span><span class="hll"><span class="n">aggregator</span><span class="o">.</span><span class="n">register_forward_hook</span><span class="p">(</span><span class="n">print_gd_similarity</span><span class="p">)</span>
 </span>
 <span class="n">inputs</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="mi">8</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>  <span class="c1"># 8 batches of 16 random input vectors of length 10</span>
 <span class="n">task1_targets</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="mi">8</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>  <span class="c1"># 8 batches of 16 targets for the first task</span>