A code-first inspection of the AWD-LSTM
AWD-LSTM is a special kind of recurrent neural network that is useful for language models. How can we set it up and customize it with the fastai library?
from fastai2.text.all import *
In the previous post we explored the Norec Norwegian language corpus. We grabbed the reviews for films and TV-shows, parsed the html-text and created labels based on the ratings. In the next few posts I want to use ULMFiT and other methods to predict the sentiment of the reviews based on the text.
ULMFiT has three main steps:
- Train a language model on a large general purpose corpus such as Wikipedia
- Fine-tune the language model on the text your are working with - the style is most likely different than a Wikipedia article
- Combine the encoder of the fine-tuned language model with a linear classifier to predict the class of your text
The core of the ULMFiT method is a type of Recurrent neural network (RNN) called AWD-LSTM. AWD-LSTM is a special kind of Recurrent neural network (RNN) with tuned dropout parameters among other. We need to look into this architecture before we continue with our modeling. For an explanation of what an LSTM actually is i suggest checking out this blog post by Chris Olah. In general, most of Chris' posts and papers are worth reading!
Let's first start by inspecting fastai's language_model_learner. It's a learner class designed to be used for language models, and holds both dataloaders and the architecture along with various hyperparameters. We can use the doc() method to show us the documentation:
doc(language_model_learner)
The documentation tells us that we can pass arch = AWD-LSTM and modify the awd_lstm_lm_config to customize the architecture. The config dictionary specifies various hyperparameters and settings inspired by the aforementioned AWD-LSTM paper. By changing this dictionary we can customize our AWD-LSTM to fit our specific needs:
awd_lstm_lm_config
Let's check the documentation and source code of the AWD-LSTM class. You can check the source code directly in the notebook by appending a ?? behind the method name:
AWD_LSTM??
The source code shows us a few interesting lines we'll look more into in the next few sections:
self.encoder = nn.Embedding(vocab_sz, emb_sz, padding_idx=pad_token)self.rnns = nn.ModuleList([self._one_rnn(emb_sz if l == 0 else n_hid, (n_hid if l != n_layers - 1 else emb_sz)//self.n_dir, bidir, weight_p, l) for l in range(n_layers)])self.input_dp = RNNDropout(input_p)self.hidden_dps = nn.ModuleList([RNNDropout(hidden_p) for l in range(n_layers)])
Once again I'll be lazy and rather refer to another blog that explains embeddings in detail. The blog is by Jay Alammar and has explanations of many deep learning and NLP concepts. The essence is that we'll turn each token in our vocabulary into a vector of some size that represents various aspects of that token. The weights of this vector will be trainable and gives our neural network a lot of flexibility in assigning various properties to each token.
The embedding is created by: self.encoder = nn.Embedding(vocab_sz, emb_sz, padding_idx=pad_token) Here we see that fastai is built on top of pyTorch and relies on pyTorch's fundamental methods in its own code. The encoder layer is a call to nn.Embedding, see documentation. Let's create an embedding of size 10x3 with padding_idx = 0:
embedding = nn.Embedding(num_embeddings=10, embedding_dim=3, padding_idx=0)
embedding.weight
The embedding now has 10 vectors of length 3 with randomly initialized weights. Note that the first one (index 0) is all 0. This is because 0 is our padding index. Next we'll pass some some sample data inp, and inspect the result. We can think of our input as the index of words in a dictionary. E.g. 1='this', 7='is', 4='not' and 3='easy'. 0 will be our padding token. The padding token is a special token that is used to ensure that some text has a certain length. This is useful when stacking various pieces of text into a batch where sizes needs to match.
inp = torch.LongTensor([1,7,4,3,0,0])
emb = embedding(inp)
emb
We see that the embedding produced by feeding the input corresponds to the weights of our original embedding. That is, index 1 of inp is the item '7'. So emb[1] is basically a lookup for embedding.weight[7].
embedding.weight[7], emb[1]
To summarize:We'll need an embedding with the number of embeddings equal to our vocabulary size, and embedding size of 400 and a padding token-id which corresponds to whichever token has been used as padding in our vocabulary.
Secondly we create a list of RNN-layers with various dimensions:
self.rnns = nn.ModuleList([self._one_rnn(emb_sz if l == 0 else n_hid,
(n_hid if l != n_layers - 1 else emb_sz)//self.n_dir, bidir, weight_p, l) for l in(n_layers)])
The code stacks RNN-layers of embedding size x hidden size for the first layer, and hidden size x embedding size for the final. Let's verify that this works for various number of layers:
AWD_LSTM(vocab_sz=10_000, emb_sz=400, n_hid=1152, n_layers=2)
AWD_LSTM(vocab_sz=10_000, emb_sz=400, n_hid=1152, n_layers=5)
We see the first and final layers have similar dimensions in the two examples.
To summarize:We'll use a 3 layer network with input and output dimensions of (400, 1152), (1152, 1152) and (1152, 400) as in the AWD-LSTM paper. This should be handled automatically by the library.
In the module list above, the layers are actually WeightDropout layers. We can verify this from the hidden constructor method that is called when the RNNs are being created. First, a regular nn.LSTM layer is created before being passed to the WeightDropout module.
def _one_rnn(self, n_in, n_out, bidir, weight_p, l):
"Return one of the inner rnn"
rnn = nn.LSTM(n_in, n_out, 1, batch_first=True, bidirectional=bidir)
return WeightDropout(rnn, weight_p)
Lets have a look at an example from the nn.LSTM documentation, also see source code here. We'll make a 1 layer LSTM with input size of 10 and hidden size of 20. Note that in the AWD-LSTM case the input size is equal to the embedding size (400 by default).
inp_s = 10 # input size
hid_s = 20 # hidden size
lstm = nn.LSTM(input_size = inp_s, hidden_size = hid_s, num_layers=1)
The documentation details that the LSTM expects input in the form of input(seq_len, batch, input_size). Seq_len is the length of the part of the text the model will see in each iteration (seq_len = 72 by default in fastais language_model_learner, that is 72 tokens). The batch size is the number of documents the models sees in each iteration.
h0 and c0 are the inital hidden and cell states (set to 0 if not provided). The documentation specifiy their shapes as: (num_layers * num_directions, batch, hidden_size).
bs = 16
n_l, n_d = 1, 1 # we are testing a 1 layer and 1 direction lstm
seq_len = 5
inp = torch.randn(seq_len, bs, inp_s)
h0 = torch.randn(n_l*n_d, bs, hid_s)
c0 = torch.randn(n_l*n_d, bs, hid_s)
inp.shape, h0.shape, c0.shape
The output from the LSTM should be a tuple of output, (h_n, c_n)where output has shape given by: (seq_len, batch, num_directions * hidden_size):
out, (hn, cn) = lstm(inp)
out.shape
Let's also check the actual shape of our weights by looping through the state_dict():
[(key, lstm.state_dict()[key].shape) for key in lstm.state_dict().keys()]
Here we recognize the input size of 10 and hidden size of 20, but where does the 80 come from? The documentation specifies that the weights will be of dimension (4*hidden_size, input_size). The '4' is called gate_size, and we can find this in the source code for the base RNN module also:
if mode == 'LSTM':
gate_size = 4 * hidden_size
To summarize:we expect two sets of weights and biases per LSTM: weight_ih_l0 with a shape of (4hidden_size, input_size)
- weight_hh_l0 with a shape of (4*hidden_size, hidden_size)
- bias_ih_l0 with a shape of (4*hidden_size)
- bias_hh_l0 with a shape of (4*hidden_size)
We can see from the _one_rnn that the nn.LSTM is transformed to a WeightDropout module. The documentation describes the module as 'A module that warps another layer in which some weights will be replaced by 0 during training'. From the source code we can see that it's the weight_hh_l0 weights that will be modified, and that these weights are duplicated with suffix 'raw' in the WeightDropout module: self.register_parameter(f'{layer}_raw', nn.Parameter(w.data)).
Let's see if we can verify this by first checking the weights from the lstm from the above section. Of the 80*20 = 1600 weights in the hh_l0 layer, none are 0:
orig_wts = getattr(lstm, 'weight_hh_l0')
orig_wts.shape, (orig_wts == 0.).sum()
But if pass the lstm through the WeightDroput module, approximately half of the 1600 weights are set to 0. Note that we have to call the model on the input since the weights are only reset during the forward pass. The weights are also only reset for the WeighDropout's internal LSTM module, while a copy with suffix '_raw' retains the original weights.
wd = WeightDropout(lstm, weight_p=0.5)
_,_ = wd(inp) # we don't need the output in this case
The .module attribute of the wd object is our original LSTM:
wd.module
And about half its weights have been set to 0:
(getattr(wd.module, 'weight_hh_l0')==0.0).sum()
The original weights from the lstm matches the '_raw' weights of the WeightDropout module:
test_eq(orig_wts, getattr(wd, 'weight_hh_l0_raw'))
The new layers are as expected:
wd.state_dict().keys()
Finally several RNNDropout layers are being created - one for the input and one for each LSTM. This dropout is applied to the input embedding and on the output of each LSTM. We can test the functionality with inp from the above section.
dp = RNNDropout(0.5)
dp_out = dp(inp)
inp.shape, dp_out.shape
The documentation also says: 'Dropout with probability p that is consistent on the seq_len dimension.' In our input from the above section, seq_len is the first dimension (index 0), and if we check for items equaling 0 and sum along the second dimension (index 1) we see that the same tokens are dropped out for the entire batch (our sample batch size is 16) consistently in approximately half of the instances.
(dp_out == 0).sum((1))
Let's take a look at a minimal IMDB example from the fastai documentation to verify our understanding of the AWD-LSTM architecture.
imdb_path = untar_data(URLs.IMDB_SAMPLE)
df = pd.read_csv(imdb_path/'texts.csv')
dls = TextDataLoaders.from_df(df, path=imdb_path, text_col='text', is_lm=True, valid_col='is_valid')
learn = language_model_learner(dls, AWD_LSTM)
The vocab is of length 7080 and vocab index 1 is 'xxpad':
dls.vocab[:5], len(dls.vocab)
In our model we recognize Embedding(7080, 400, padding_idx=1) as vocab_size x embedding size with the correct padding token. We also see that the (input, output) dimensions of our LSTM-layers are as expected, and with the expected dropout layers added.
learn.model
The model summary shows us the default batch size of 64 and seq_len of 72.
learn.summary()
And finally, the layer names and shapes also is consistent with a gate size of 4 (1152*4 = 4608). Note the enumeration of the layers: 0. is the encoder part of the architecture (including the embedding called encoder) and 1. is the decoder.
for key in learn.model.state_dict().keys():
print(key, '\t', learn.model.state_dict()[key].shape)