diff --git a/fastai_contrib/learner.py b/fastai_contrib/learner.py index af72f95..d787867 100644 --- a/fastai_contrib/learner.py +++ b/fastai_contrib/learner.py @@ -28,7 +28,5 @@ def bilm_learner(data:DataBunch, bptt:int=70, emb_sz:int=400, nh:int=1150, nl:in def bilm_split(model:nn.Module) -> List[nn.Module]: "Split a RNN `model` in groups for differential learning rates." - groups = [[rnn, dp] for rnn, dp in zip(model[0].forward_rnns, model[0].hidden_dps)] - groups += [[rnn, dp] for rnn, dp in zip(model[0].backward_rnns, model[0].hidden_dps)] - groups.append([model[0].encoder, model[0].encoder_dp, model[1]]) - return groups + + return [f+b for f,b in zip(lm_split(model.fwd_lm),lm_split(model.bwd_lm))] diff --git a/fastai_contrib/models.py b/fastai_contrib/models.py index dbb8be4..df323af 100644 --- a/fastai_contrib/models.py +++ b/fastai_contrib/models.py @@ -2,133 +2,44 @@ from fastai.torch_core import * from fastai.layers import * from fastai.text.models import * +class BiLMModel(nn.Module): -class BiLMCore(nn.Module): - """ - AWD-LSTM/QRNN inspired by https://arxiv.org/abs/1708.02182. - Inspired by https://github.com/allenai/allennlp/blob/master/allennlp/models/bidirectional_lm.py#L65 - """ - initrange=0.1 - - def __init__(self, vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, pad_token:int, bidir:bool=False, - hidden_p:float=0.2, input_p:float=0.6, embed_p:float=0.1, weight_p:float=0.5, qrnn:bool=False): - + def __init__(self, fwd_lm:nn.Module, bwd_lm:nn.Module): super().__init__() - self.bs,self.qrnn,self.ndir = 1, qrnn,(2 if bidir else 1) - self.emb_sz,self.n_hid,self.n_layers = emb_sz,n_hid,n_layers - # embeddings are shared between forward and backward LMs - self.encoder = nn.Embedding(vocab_sz, emb_sz, padding_idx=pad_token) - self.encoder_dp = EmbeddingDropout(self.encoder, embed_p) - if self.qrnn: - #Using QRNN requires cupy: https://github.com/cupy/cupy - from fastai.text.qrnn.qrnn import QRNNLayer + self.fwd_lm = fwd_lm + self.bwd_lm = bwd_lm - def create_qrnn_layers(): - return [QRNNLayer(emb_sz if l == 0 else n_hid, (n_hid if l != n_layers - 1 else emb_sz)//self.ndir, - save_prev_x=True, zoneout=0, window=2 if l == 0 else 1, output_gate=True, - use_cuda=torch.cuda.is_available()) for l in range(n_layers)] - self.forward_rnns = create_qrnn_layers() - self.backward_rnns = create_qrnn_layers() - for rnn in self.forward_rnns + self.backward_rnns: - rnn.linear = WeightDropout(rnn.linear, weight_p, layer_names=['weight']) - else: - def create_lstm_layers(): - return [nn.LSTM(emb_sz if l == 0 else n_hid, (n_hid if l != n_layers - 1 else emb_sz)//self.ndir, - 1, bidirectional=False) for l in range(n_layers)] - self.forward_rnns = [WeightDropout(rnn, weight_p) for rnn in create_lstm_layers()] - self.backward_rnns = [WeightDropout(rnn, weight_p) for rnn in create_lstm_layers()] - self.forward_rnns = torch.nn.ModuleList(self.forward_rnns) - self.backward_rnns = torch.nn.ModuleList(self.backward_rnns) - self.encoder.weight.data.uniform_(-self.initrange, self.initrange) - self.input_dp = RNNDropout(input_p) - self.hidden_dps = nn.ModuleList([RNNDropout(hidden_p) for l in range(n_layers)]) + def forward(self, input): + sl, bs, tracks = input.size() - def forward(self, input:LongTensor)->Tuple[Tensor,Tensor]: - sl,bs,tracks = input.size() - assert tracks == 2, "It should have two tracks for forward and backward pass" - if bs != self.bs: - self.bs = bs - self.reset() + decoded = [] + raw_outputs = [] + outputs = [] - return [self.fwdlm_forwad(input[..., 0]), self.bwdlm_forwad(input[..., 1])] + fwd_o = self.fwd_lm(input[..., 0]) + bwd_o = self.bwd_lm(input[..., 1]) - def bwdlm_forwad(self, input): - raw_output = self.input_dp(self.encoder_dp(input)) - new_hidden,raw_outputs,outputs = [],[],[] - for l, (rnn,hid_dp) in enumerate(zip(self.backward_rnns, self.hidden_dps)): - raw_output, new_h = rnn(raw_output, self.bwdlm_hidden[l]) - new_hidden.append(new_h) - raw_outputs.append(raw_output) - if l != self.n_layers - 1: raw_output = hid_dp(raw_output) - outputs.append(raw_output) - self.bwdlm_hidden = to_detach(new_hidden) - - return (raw_outputs, outputs) - - def fwdlm_forwad(self, input): - raw_output = self.input_dp(self.encoder_dp(input)) - new_hidden,raw_outputs,outputs = [],[],[] - for l, (rnn,hid_dp) in enumerate(zip(self.forward_rnns, self.hidden_dps)): - raw_output, new_h = rnn(raw_output, self.fwdlm_hidden[l]) - new_hidden.append(new_h) - raw_outputs.append(raw_output) - if l != self.n_layers - 1: raw_output = hid_dp(raw_output) - outputs.append(raw_output) - self.fwdlm_hidden = to_detach(new_hidden) - - return (raw_outputs, outputs) - - def _one_hidden(self, l:int)->Tensor: - "Return one hidden state." - nh = (self.n_hid if l != self.n_layers - 1 else self.emb_sz)//self.ndir - return self.weights.new(self.ndir, self.bs, nh).zero_() + return torch.stack([fwd_o[0], bwd_o[0]], dim=2), (fwd_o[1]+bwd_o[1]), (fwd_o[2] + bwd_o[2]) def reset(self): - "Reset the hidden states." - [r.reset() for r in self.forward_rnns if hasattr(r, 'reset')] - [r.reset() for r in self.backward_rnns if hasattr(r, 'reset')] - self.weights = next(self.parameters()).data - if self.qrnn: self.fwdlm_hidden = [self._one_hidden(l) for l in range(self.n_layers)] - else: self.fwdlm_hidden = [(self._one_hidden(l), self._one_hidden(l)) for l in range(self.n_layers)] - if self.qrnn: self.bwdlm_hidden = [self._one_hidden(l) for l in range(self.n_layers)] - else: self.bwdlm_hidden = [(self._one_hidden(l), self._one_hidden(l)) for l in range(self.n_layers)] - -class BiLinearDecoder(nn.Module): - "To go on top of a RNNCore module and create a Language Model." - - initrange=0.1 - - def __init__(self, n_out:int, n_hid:int, output_p:float, tie_encoder:nn.Module=None, bias:bool=True): - super().__init__() - self.decoder = nn.Linear(n_hid, n_out, bias=bias) - self.decoder.weight.data.uniform_(-self.initrange, self.initrange) - self.output_dp = RNNDropout(output_p) - if bias: self.decoder.bias.data.zero_() - if tie_encoder: self.decoder.weight = tie_encoder.weight - - def forward(self, input:List[Tuple[Tensor,Tensor]])->Tuple[Tensor,Tensor,Tensor]: - decoded=[] - raw_outputs=[] - outputs=[] - for lm_input in input: - d, ro, o = self.one_forward(lm_input) - decoded.append(d) - raw_outputs += ro - outputs += o - return torch.stack(decoded, dim=2), raw_outputs, outputs - - def one_forward(self, input): - raw_outputs, outputs = input - output = self.output_dp(outputs[-1]) - decoded = self.decoder(output.view(output.size(0)*output.size(1), output.size(2))) - return decoded, raw_outputs, outputs - + "Reset the hidden states of underlaying lms." + self.fwd_lm.reset() + self.bwd_lm.reset() def get_bilm(vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, pad_token:int, tie_weights:bool=True, qrnn:bool=False, bias:bool=True, bidir:bool=False, output_p:float=0.4, hidden_p:float=0.2, input_p:float=0.6, embed_p:float=0.1, weight_p:float=0.5)->nn.Module: - "Create a full AWD-LSTM." - rnn_enc = BiLMCore(vocab_sz, emb_sz, n_hid=n_hid, n_layers=n_layers, pad_token=pad_token, qrnn=qrnn, bidir=bidir, - hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p) - enc = rnn_enc.encoder if tie_weights else None - return SequentialRNN(rnn_enc, BiLinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias)) + "Create a two AWD-LSTM one for each direction " + fwd_rnn_enc = RNNCore(vocab_sz, emb_sz, n_hid=n_hid, n_layers=n_layers, pad_token=pad_token, qrnn=qrnn, bidir=bidir, + hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p) + bwd_rnn_enc = RNNCore(vocab_sz, emb_sz, n_hid=n_hid, n_layers=n_layers, pad_token=pad_token, qrnn=qrnn, bidir=bidir, + hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p) + enc = None + if tie_weights: + enc = fwd_rnn_enc.encoder + fwd_rnn_enc.encoder.weight = enc.weight + bwd_rnn_enc.encoder.weight = enc.weight + + return BiLMModel( + fwd_lm=SequentialRNN(fwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias)), + bwd_lm=SequentialRNN(bwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias))) \ No newline at end of file diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 34455a9..55a64db 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -26,10 +26,15 @@ import fastai_contrib.data as contrib_data # conda install -c pytorch -c fastai fastai pytorch-nightly [cuda92] # cupy needs to be installed for QRNN +def accuracy_fwd(input, targs): + return accuracy(input[...,0], targs[...,0]) +def accuracy_bwd(input, targs): + return accuracy(input[...,1], targs[...,1]) + def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, bs=70, bptt=70, name='wt-103', model_dir='models', num_epochs=10, - bidir=False): + bidir=False, ds_pct=1.0): """ :param dir_path: The path to the directory of the file. :param cuda_id: The id of the GPU. Uses GPU 0 by default or no GPU when @@ -58,8 +63,7 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, if qrnn: print('Using QRNNs...') - #trn_path = dir_path / 'wiki.train.tokens' - trn_path = dir_path / 'wiki.valid.tokens' + trn_path = dir_path / 'wiki.train.tokens' val_path = dir_path / 'wiki.valid.tokens' tst_path = dir_path / 'wiki.test.tokens' for path_ in [trn_path, val_path, tst_path]: @@ -69,6 +73,9 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, # read the already whitespace separated data without any preprocessing trn_tok = read_whitespace_file(trn_path) val_tok = read_whitespace_file(val_path) + if ds_pct < 1.0: + trn_tok = trn_tok[:int(len(trn_tok) * ds_pct)] + val_tok = val_tok[:int(len(val_tok) * ds_pct)] # create the vocabulary cnt = Counter(word for sent in trn_tok for word in sent) @@ -123,7 +130,11 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, # compared to standard Adam, we set beta_1 to 0.8 learn.opt_fn = partial(optim.Adam, betas=(0.8, 0.99)) learn.true_wd = False - learn.metrics=[] # accuracy does not work when we have multiple dimensions at the end. + + if bidir: + learn.metrics = [accuracy_fwd, accuracy_bwd] + else: + learn.metrics = [accuracy] # save vocabulary print('Saving vocabulary...') with open(model_dir / f'itos_{name}.pkl', 'wb') as f: