Make from_pretrained support parameters defined in the forward pass
Author: NielsRogge
Pull request: no

# 🚀 Feature request

Be able to make the `from_pretrained` method support parameters defined in the forward pass rather than the init of a certain `nn.Module` of a model.

## Motivation

So I was implementing PerceiverIO in PyTorch, and it works fine for the text, image and optical flow modalities. However, for the [multimodal autoencoding](https://github.com/deepmind/deepmind-research/blob/master/perceiver/colabs/video_autoencoding.ipynb) one, I do have some limitations to make it fully HuggingFace-API friendly. The problem is that for this model, modality-specific parameters are defined whose shape depends on the combination of a large number of preprocessing parameters. 

One could come up with (rather complex) formulas to define everything from beforehand based on the several combinations of preprocessing configurations, in order to be able to define the parameters at initialization. This is what I tried at first. However, these formulas become very complex, and hence very error prone. To give an example: 

Suppose that one uses the same preprocessing parameters as defined in the [official video autoencoding Colab notebook](https://github.com/deepmind/deepmind-research/blob/master/perceiver/colabs/video_autoencoding.ipynb), then the 3 modalities of a video (images, audio and class label) have the following shape after preprocessing (preprocessing is part of the model!):
* image modality will have shape (1, 50176, 243) 
* audio modality will have shape (1, 1920, 401) 
* class label modality will have shape (1, 1, 700). 

The Perceiver authors then use trainable modality-specific parameters that pad the respective modality to the highest number of channels + a certain minimum padding size (in this case, the class label modality has the highest number of channels, namely 700, and the minimum padding size is 4, hence all modalities should be padded to have 704 channels). Hence, the image padding parameter will be an `nn.Parameter` of shape `(1, 704-243)` = `(1, 461)`, the audio padding parameter will be an `nn.Parameter` of shape `(1, 704-401)` = `(1, 303)` and the label padding parameter will be an `nn.Parameter` of shape `(1, 704-700)` = `(1, 4)`.

However, this is only for _one specific combination of preprocessing_ (for the image modality e.g., Fourier embeddings with a certain `num_bands`, concatenation with original positions, no additional projection, then concatenation with patches, etc.). These settings are all configurable in the `ImagePreprocessor`, which makes the number of preprocessing combinations large, and hence rather difficult to come up with a general formula to calculate the `num_channels` of the image modality and audio modalities that take into account all possible settings. For this given case, the `num_channels` of the image modality is equal to `temporal_downsample * spatial_downsample * spatial_downsample * in_channels + (2 * d * num_bands + d)`, but that's only one very specific setting of preprocessing, there are a lot more.

In frameworks like Haiku and Flax, one can easily define parameters in the forward pass, as one always first needs to forward a dummy input through the model, in order to do shape inference for all parameters. In PyTorch, this is also possible, by defining the parameters using `register_parameter` in the init method, and then instantiating them in the forward pass, based on the size of the inputs:

```
import torch.nn as nn

class Test(nn.Module):
    def __init__(self):
        # you need to register the parameter names earlier
        super().__init__()
        self.register_parameter('weight', None)

    def forward(self, input):
        if self.weight is None:
            self.weight = nn.Parameter(torch.randn(input.size()))
        return self.weight @ input
```

And it's exactly this idea that would be very handy to use when implementing `PerceiverForMultimodalAutoencoding` (I've currently implemented it this way). However, this doesn't play well with the `from_pretrained()` method, as this method **assumes** that all parameters of a model can be defined at initialization. It doesn't support first forwarding a dummy input through the model in order to instantiate all parameters, and then load the weights from a state dict. Hence, loading this model with the `.from_pretrained()` method will not instantiate these parameters.

For this model, the dummy inputs would look like this:

```
images = torch.randn((1, 16, 3, 224, 224))
audio = torch.randn((1, 30720, 1))
inputs = dict(image=images, audio=audio, label=torch.zeros((images.shape[0], 700)))
```
If we can pass in a `dummy_inputs` parameter to the `from_pretrained` method, then it could use it to first instantiate all parameters of the model, and then load it with the weights from the hub.

Curious to hear your opinions, it might be a bad idea or it could be that there are other (better) options. If this is not possible, then I guess there's no option but to come up with the formulas.

cc @patrickvonplaten @sgugger @LysandreJik @patil-suraj 

Comments:
I think the suggested approach is a bit too heavy. There is an easy way to load those dynamic weights directly inside the `from_pretrained` method by adapting a bit the code. At a first glance, changing the load function [here](https://github.com/huggingface/transformers/blob/269c3d1400267d966a5b2a962c69c56ac8aca5c3/src/transformers/modeling_utils.py#L1540) by adding the following code at the end seems to work pretty well and doesn't require adding any new argument:

```diff
        def load(module: nn.Module, prefix=""):
            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
            args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
            if is_deepspeed_zero3_enabled():
                import deepspeed

                # because zero3 puts placeholders in model params, this context
                # manager gathers (unpartitions) the params of the current layer, then loads from
                # the state dict and then re-partitions them again
                with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0):
                    if torch.distributed.get_rank() == 0:
                        module._load_from_state_dict(*args)
            else:
                module._load_from_state_dict(*args)

            for name, child in module._modules.items():
                if child is not None:
                    load(child, prefix + name + ".")

+           for name, param in module._parameters.items():
+               if param is None and prefix + name in state_dict:
+                   setattr(module, name, torch.nn.Parameter(state_dict[prefix+name]))
```

(the added code is in the last three lines). It would need to be battle-tested a bit more but just creates a proper parameters from the weight inside the state dict if a parameter is None somewhere in the model and has a corresponding weight. This is similar to how you initialize them in the forward pass, just from the model state dict.

I'm not really in favor of adapting PyTorch's way of initializing models to Flax's of Tensorflow's "tracing-the-weights" initialization scheme.

PyTorch's way of initialization scheme has a couple of advantages:

- `.from_pretrained(...)` does not run the forward pass. This is nice for debugging purposes. This might seem like a small thing, but it's a huge change in user-debugging-experience if we would start tracing PyTorch's models IMO. 
- PyTorch forces the code to always define input and output dimensions of the Parameters -> this is very nice for readability. One always knows the weight structure from looking into the code (this is a big advantage to TF and Flax IMO)

> as this method assumes that all parameters of a model can be defined at initialization

-> I think it always holds true though that all parameters can be defined at initialization no? Even if it might require a bit of math to compute a dimension, one can always define all the parameter dimensions at initialization no?

I would prefer to not make any change to the `.from_pretrained(...)` method here actually as IMO it would introduce a bit too much "under-the-hood magic"

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.