Modify EmbeddingDropout and adapt AWD_LSTM to trigger hooks in Embedding layers

[arnaujc91]

Fix for: #2850

@sutt already fixed the awd_lstm.WeightDropout class to be able to run the ActivationStats callback for the model AWD_LSTM. Still the Embedding layers were not recording any statistics. Furthermore the Hooks class had duplicated layers due to how flatten_model works.

In this PR I propose a modification to solve both issues:

1. Trigger the hooks for the Embedding layer
2. Avoid dupicated layers in the Hooks class.

An example of this can be found here: example EmbeddingDropout.

To solve the issue I just modified slightly the class EmbeddingDropout and adapted AWD_LSTM to work exactly the same as before for these changes. I think now the code is more compact than before.

I by no means have your expertise, nevertheless I tried to do my best refactoring the code. If there is anything that I did wrong just let me know.

Thanks a lot!

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[jph00]

Thanks @arnaujc91 . Can you please fix the test failure shown here in CI? It's saying ModuleAttributeError: 'AWD_LSTM' object has no attribute 'encoder_dp'. Also, could you add to the notebook some basic docs and a test or example for EmbeddingDropout so that I can understand the new code you've added?

[arnaujc91]

Ok, I have created a notebook where I tried to explain as good as possible why I am doing this PR and what are the possible solutions that I see to the problem.

See the following notebook: EmbeddingDropout

P.S. I updated again the notebook. Let me know what you think. Thanks a lot!

[jph00]

Thanks @arnaujc91 - I'm still seeing the test failure however. Would you be able to fix that?

[jph00]

Please fix test failure:

ModuleAttributeError: 'AWD_LSTM' object has no attribute 'encoder_dp'

[arnaujc91]

I modified the PR to pass the previous test, but now is failing in another test:

      1 enc = nn.Embedding(10, 7, padding_idx=1)
----> 2 enc_dp = EmbeddingDropout(enc, 0.5)
      3 tst_inp = torch.randint(0,10,(8,))
      4 tst_out = enc_dp(tst_inp)
      5 for i in range(8):

This is because the class EmbeddingDropout expects that you pass it an instance of nn.Embedding to the constructor but precisely the point of my PR is that the new EmbeddingDropout does not need an instance of nn.Embedding. So either I can change this test or the tests are too rigid for this PR to make sense.

As I explained in the previous Notebook there are two possibilities to fix the ActivationStats problem:

1. Modify flatten_model
2. Modify EmbeddingDropout

I decided to modify EmbeddingDropout. The reason was twofold: it simplified the code and it solved the issue with the hooks. My current modification just affects EmbeddingDropout; AWD_LST stays the same. Do you really need that EmbeddingDropout requires an instance of nn.Embedding in other parts of the library?

If the tests can't be modified then I will need to close this PR and maybe check if i can modify flatten_model instead.

Can you or me modify the tests for the layer EmbeddingDropout?

Thanks a lot Jeremy.

[jph00]

Yes please do go ahead and modify the tests. Many thanks!

[arnaujc91]

Please before accepting the modifications of this PR check the Notebook I created explaining how to make full sense of this PR. Right now everything works, but the code can be further cleaned.

Here is the notebook: encoder and encoder_dp are redundant

Thanks a lot!