Efficient way to convert an `nn.module` into a single tensor


What is the best way in order to gather the parameters of a neural network (given as an nn.Module object) into a single tensor (vector) and vice-versa?

The simplest option seems to be looping over the self.parameters(), but I am not sure how efficient it is especially when we want to reconstruct the model from a tensor.

Yes, iterating the .parameters() and .bufers() (and flattening the tensors, if desired) sounds like the right approach.

1 Like