(self, rnn_type, ntoken, ninp, nhid, nlayers, dropout=0.5, tie_weights=False)
| 7 | """Container module with an encoder, a recurrent module, and a decoder.""" |
| 8 | |
| 9 | def __init__(self, rnn_type, ntoken, ninp, nhid, nlayers, dropout=0.5, tie_weights=False): |
| 10 | super(RNNModel, self).__init__() |
| 11 | self.ntoken = ntoken |
| 12 | self.drop = nn.Dropout(dropout) |
| 13 | self.encoder = nn.Embedding(ntoken, ninp) |
| 14 | if rnn_type in ['LSTM', 'GRU']: |
| 15 | self.rnn = getattr(nn, rnn_type)(ninp, nhid, nlayers, dropout=dropout) |
| 16 | else: |
| 17 | try: |
| 18 | nonlinearity = {'RNN_TANH': 'tanh', 'RNN_RELU': 'relu'}[rnn_type] |
| 19 | except KeyError as e: |
| 20 | raise ValueError( """An invalid option for `--model` was supplied, |
| 21 | options are ['LSTM', 'GRU', 'RNN_TANH' or 'RNN_RELU']""") from e |
| 22 | self.rnn = nn.RNN(ninp, nhid, nlayers, nonlinearity=nonlinearity, dropout=dropout) |
| 23 | self.decoder = nn.Linear(nhid, ntoken) |
| 24 | |
| 25 | # Optionally tie weights as in: |
| 26 | # "Using the Output Embedding to Improve Language Models" (Press & Wolf 2016) |
| 27 | # https://arxiv.org/abs/1608.05859 |
| 28 | # and |
| 29 | # "Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling" (Inan et al. 2016) |
| 30 | # https://arxiv.org/abs/1611.01462 |
| 31 | if tie_weights: |
| 32 | if nhid != ninp: |
| 33 | raise ValueError('When using the tied flag, nhid must be equal to emsize') |
| 34 | self.decoder.weight = self.encoder.weight |
| 35 | |
| 36 | self.init_weights() |
| 37 | |
| 38 | self.rnn_type = rnn_type |
| 39 | self.nhid = nhid |
| 40 | self.nlayers = nlayers |
| 41 | |
| 42 | def init_weights(self): |
| 43 | initrange = 0.1 |
no test coverage detected