XuehaiPan
commented
1 year ago Hi @patel-zeel, thanks for raising this! This is a very useful feature. And thank for the very detailed comment and the code snippet.
Currently, I have written my own version of
ravel_pytreeusing existing function ofoptree:def ravel_pytree(pytree): leaves, structure = optree.tree_flatten(pytree) shapes = [leaf.shape for leaf in leaves] sizes = [leaf.numel() for leaf in leaves] flat_params = torch.cat([leaf.flatten() for leaf in leaves]) def unravel_function(flat_params): assert flat_params.numel() == sum(sizes), f"Invalid size {flat_params.numel()} != {sum(sizes)}" assert len(flat_params.shape) == 1, f"Invalid shape {flat_params.shape}" flat_leaves = flat_params.split(sizes) leaves = [leaf.reshape(shape) for leaf, shape in zip(flat_leaves, shapes)] return optree.tree_unflatten(structure, leaves) return flat_params, unravel_function
I think the ravel_pytree API should be available for all pytrees that contain NDArray(s). We would to have a generic implementation for torch.Tensor, numpy.ndarray, jax.Array, etc. There are some small differences between these data structures. For example, for numpy.ndarray and jax.Array, there is no method numel. We can do this with arr.ravel().shape[0].
Here is the initial API design:
concatenate_func: Callable[[Sequence[ArrayLike]], ArrayLike]
split_func: Callable[[ArrayLike, Sequence[int]], List[ArrayLike]]
def ravel_pytree(pytree, concatenate_func, split_func):
leaves, treespec = optree.tree_flatten(pytree)
flat_leaves = [leaf.ravel() for leaf in leaves]
flat_params = concatenate_func(flat_leaves)
shapes = [leaf.shape for leaf in leaves]
sizes = [flat_leaf.shape[0] for flat_leaf in flat_leaves]
total_size = flat_params.shape[0]
def unravel_function(flat_params):
assert len(flat_params.shape) == 1, f"Invalid shape {flat_params.shape}"
assert flat_params.shape[0] == total_size, f"Invalid size {flat_params.shape[0]} != {total_size}"
flat_leaves = split_func(flat_params, sizes)
leaves = [flat_leaf.reshape(shape) for flat_leaf, shape in zip(flat_leaves, shapes)]
return optree.tree_unflatten(treespec, leaves)
return flat_params, unravel_functionNow optree is going to be the C++ implementation for the pytree utility in the PyTorch package (maybe GA in PyTorch 2.2). Would you like to have this ravel_pytree to be implemented in the PyTorch library or the optree package here? Any thoughts are much appreciated.
patel-zeel
Required prerequisites
Motivation
Hi,
What does this function do?
This function is motivated from JAX's
ravel_pytreewhich converts a PyTree into a 1D tensor and provides a function to convert that 1D tensor back to the original PyTree format.This function is useful for these applications: i) Ease the process of defining "Hypernetworks"; ii) Hessian computation of neural networks (e.g., Laplace approximation)
Challanges in Hypernetworks
In "Hypernetworks" line of work, we have a hypernetwork
hypernet, which generates the parameters of another networknet. Typically,hypernet's last layer will be aLinear/Denselayer, and thus it will be in a "flattened" format (1D tensor). It is challenging to convert this flat tensor into parameters ofnet.Challenges in computing Hessians in Laplace approximation.
Laplace approximation of neural network requires one to compute the hessian of the negative log joint w.r.t. parameters and then invert it. The inverse of Hessian is the covariance matrix of the posterior normal distribution. One has to flatten the parameters to get the hessian. Also, a sample parameter from this posterior is a flat array that needs to be restructured before inserting into the network.
Solution
Currently, I have written my own version of
ravel_pytreeusing the existing functions inoptree:Here is how it can be used for different applications.
Hypernetworks
We would need to transform the flat output of
hypernetto a structure that matchesnet.named_parameters()and then trigger a functional call withtorch.func.functional_call(net, named_parameters, inputs). A code snipped would look like this:Laplace Approximation
Alternatives
An alternative is to use the manually defined function above or to use some other way.
Additional context
No response