Reuse RNNCore in implementation of BiLM, add accuracy

This commit is contained in:
Piotr Czapla
2018-11-14 21:16:16 +01:00
parent cf7d93c378
commit 33f9eb2cc7
3 changed files with 46 additions and 126 deletions
+2 -4
View File
@@ -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))]
+29 -118
View File
@@ -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)))
+15 -4
View File
@@ -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: