Fix LayoutLMv2 init issue and doctest

[ydshieh]

What does this PR do?

(I think I have done all I can do for doctest - there are 2 other failures where I ping other people)

LayoutLMv2Model.visual_segment_embedding's initialization (when it is not in a checkpoint) won't be reproducible even if seed is set. This is because it is defined

self.visual_segment_embedding = nn.Parameter(nn.Embedding(1, config.hidden_size).weight[0])

and we have _fast_init used in loading +_init_weights is not designed to deal with nn.Parameter.

This PR fixes the reproducibility issue by initializing visual_segment_embedding explicitly.

Remark: such places requiring this fix is very few - so we don't destroy the purpose of having _fast_init.

[ydshieh]

this is the main fix

[ydshieh]

The failure in tests_pr_documentation_tests is just missing detectron2. I can install it for that job.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[amyeroberts]

I'm not sure why we can't have this logic in _init_weights ?

[ydshieh]

The definition of _init_weights looks like def _init_weights(self, module):, where we check things like

if isinstance(module, nn.Embedding)

elif isinstance(module, (nn.Linear, nn.Conv1d)

etc. (using layer type to determine the init method)

We don't check if it is nn.Parameter as it is too generic. For example, we also have

self.q_bias = nn.Parameter(torch.zeros(1, 1, self.all_head_size))
self.v_bias = nn.Parameter(torch.zeros(1, 1, self.all_head_size))

(where they are initialized with zeros , and not treated in _init_weights.

However, here it kind special:

self.visual_segment_embedding = nn.Parameter(nn.Embedding(1, config.hidden_size).weight[0])

it was nn.Embedding but wrappered into nn.Parameter. I guess it is because we don't want to have bias (?)
But since it is nn.Parameter, it is not treated in nn.Parameter.

[ydshieh]

For some models (say, VivitPreTrainedModel), we have isinstance(module, nn.Parameter) in _init_weights

         elif isinstance(module, nn.Parameter):
            module.data.normal_(mean=0.0, std=self.config.initializer_range)

But for LayoutLM, it has multiple nn.Parameter and they don't mean to be initialized in the same way.

[ydshieh]

Gentilly ping @amyeroberts 😃

[ydshieh]

@amyeroberts Let me know if you have further question.

[amyeroberts]

Yes, sorry for being late on this. I see. My understanding of the init weights behaviour is that if a weight has been initialized then it's marked as such and skipped it it matches other patterns later on (correct me if I'm wrong here). Would it be possible then to instead something similar to what we see in e.g. wav2vec2's _init_weights where we check if the module is LayoutLMv2Model and then init this param?

I realise it's not a huge difference having this in this init than _init_weights, but there's an increasing amount of logic e.g. with quantization which assumes this behaviour is controlled within the _init_weights module

[ydshieh]

thanks for the insightful suggestion! I will give it a shot and let you know.

[ydshieh]

@amyeroberts works well! Once you approve the changes, I will update the >>> set_seed(x) to make the doctest passing.

[ydshieh]

all these and below have to be updated

[ydshieh]

turns out they work on the CI runners. No need to update

[amyeroberts]

Thanks for fixing and iterating!