mirror of
https://github.com/wassname/multifit.git
synced 2026-08-24 12:15:19 +08:00
Reuse RNNCore in implementation of BiLM, add accuracy
This commit is contained in:
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user