How to flatten and then unflatten _all_ model parameters

milad: <p>I’m trying to compute some metrics  across <em>all</em> parameters of my model. While these metrics are simple (e.g. <code>topk()</code>) and available as methods on <code>Tensor</code> class I can’t use these methods directly because I get a list of tensors with different sizes via <code>model.parameters()</code>.</p>
<p>So I have found that it’s easier to just flatten everything in <code>model.parameter()</code> list and then concatenate them into a single tensor and then compute the metric.</p>
<p>This works fine but I also need to modify my original parameters based on this computed metric. For example imagine I want to keep the largest 10% parameters in my model and set everything else to zero. What’s the most elegant/pythonic way to do this?</p>

Evan_Vogelbaum: <aside class="quote no-group" data-username="milad" data-post="1" data-topic="34730">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/milad/48/8744_2.png" class="avatar"> milad:</div>
<blockquote>
<p>So I have found that it’s easier to just flatten everything in <code>model.parameter()</code> list and then concatenate them into a single tensor and then compute the metric.</p>
</blockquote>
</aside>
<p>I tried to put something together… let me know if this helps you</p>
<pre><code class="lang-auto">def flatten_params(parameters):
    """
    flattens all parameters into a single column vector. Returns the dictionary to recover them
    :param: parameters: a generator or list of all the parameters
    :return: a dictionary: {"params": [#params, 1],
    "indices": [(start index, end index) for each param] **Note end index in uninclusive**

    """
    l = [torch.flatten(p) for p in parameters]
    indices = []
    s = 0
    for p in l:
        size = p.shape[0]
        indices.append((s, s+size))
        s += size
    flat = torch.cat(l).view(-1, 1)
    return {"params": flat, "indices": indices}


def recover_flattened(flat_params, indices, model):
    """
    Gives a list of recovered parameters from their flattened form
    :param flat_params: [#params, 1]
    :param indices: a list detaling the start and end index of each param [(start, end) for param]
    :param model: the model that gives the params with correct shapes
    :return: the params, reshaped to the ones in the model, with the same order as those in the model
    """
    l = [flat_params[s:e] for (s, e) in indices]
    for i, p in enumerate(model.parameters()):
        l[i] = l[i].view(*p.shape)
    return l

</code></pre>

Ajinkya_Bankar: <p><a class="mention" href="/u/milad">@milad</a> Did you get the solution? If yes, please share. I am also having the same question. Thank you.</p>

Ajinkya_Bankar: <p><a class="mention" href="/u/evan_vogelbaum">@Evan_Vogelbaum</a> How can I get all the flattened model parameters as numpy 1-D array or column vector through flatten_params() function? As <a class="mention" href="/u/milad">@milad</a> mentioned, after some computations on the weights, I want to unflatten the modified weights or replace the modified weights in the original model. Kindly help. Thank you.</p>

eacousineau: <p>Howdy <a class="mention" href="/u/ajinkya_bankar">@Ajinkya_Bankar</a>! Perhaps you’re looking for these functions; I’m sure you could find more succinct usages similar to your use case in Torch tutorials and examples if you search for these fns in their awesome docs:<br>
<a href="https://pytorch.org/docs/stable/generated/torch.from_numpy.html#torch.from_numpy" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.from_numpy.html#torch.from_numpy</a><br>
<a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.numpy" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/tensors.html#torch.Tensor.numpy</a></p>
<p>Be sure to review how <code>.detach()</code> and device transfer works.</p>

eacousineau: <p>To the OP’s post and <a class="mention" href="/u/evan_vogelbaum">@Evan_Vogelbaum</a>’s sln - thanks!</p>
<p>If you inspect the <code>grad_fn</code> from the output of <code>torch.cat()</code>, it does show backprop will still connect to original tensors (I believe), but yeah, aliasing won’t happen.</p>
<p>I’m wondering if there’s a way to stack a tensor s.t. that’s still “raggedly” aliases / views, similar in nature to how fn’s like <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.expand" rel="noopener nofollow ugc"><code>Tensor.expand()</code></a> work. (My naive guess is nah, but still curious nonetheless.)</p>
<p>EDIT: Nahhhh, I ain’t seeing anything in <a href="https://pytorch.org/docs/stable/tensors.html" rel="noopener nofollow ugc"><code>torch.Tensor</code></a> docs; will def. use flatten/unflattening.</p>

eacousineau: <p>As a form of masochism, tinkered with this briefly.</p>
<p>It <em>is</em> possible if you go the other way: reconstruct full vector, then replace each <code>Parameter</code> with one constructed from a tensor view. See this notebook for code + example (much longer / more complex than Evan’s sln):<br>
<a href="https://github.com/EricCousineau-TRI/repro/blob/77942acfda/python/torch/torch_view_alias_no_backprop_tho.ipynb" rel="noopener nofollow ugc">EricCousineau-TRI/repro/blob/77942acfda/python/torch/torch_view_alias_no_backprop_tho.ipynb</a></p>
<p>However, the backprop / <code>grad_fn</code> gets lost b/c <code>nn.Parameter</code> must be a leaf tensor. See:<br>
<a href="https://pytorch.org/docs/1.8.1/_modules/torch/nn/modules/module.html#Module.register_parameter" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/1.8.1/_modules/torch/nn/modules/module.html#Module.register_parameter</a></p>
<p>EDIT: Also, dunno how device transfer may affect aliasing views.</p>

sad_robot: <p>It looks like the way to do this is:</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.nn.utils.parameters_to_vector.html" rel="noopener nofollow ugc">TORCH.NN.UTILS.PARAMETERS_TO_VECTOR</a><br>
and<br>
<a href="https://pytorch.org/docs/stable/generated/torch.nn.utils.vector_to_parameters.html" rel="noopener nofollow ugc">TORCH.NN.UTILS.VECTOR_TO_PARAMETERS</a></p>
<p>Specifically something like:</p>
<pre><code class="lang-auto">param_vec = th.nn.utils.parameters_to_vector(model.parameters().values())
param_vec += something # arbitrary operation
th.nn.utils.vector_to_parameters(param_vec, model.parameters().values())
</code></pre>