| 78 | example. See https://github.com/pytorch/examples/blob/main/word_language_model/model.py |
| 79 | """ |
| 80 | def __init__(self, ps, ntoken, ninp, nhid, nlayers, dropout=0.5): |
| 81 | super(RNNModel, self).__init__() |
| 82 | |
| 83 | # setup embedding table remotely |
| 84 | self.emb_table_rref = rpc.remote(ps, EmbeddingTable, args=(ntoken, ninp, dropout)) |
| 85 | # setup LSTM locally |
| 86 | self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout) |
| 87 | # setup decoder remotely |
| 88 | self.decoder_rref = rpc.remote(ps, Decoder, args=(ntoken, nhid, dropout)) |
| 89 | |
| 90 | def forward(self, input, hidden): |
| 91 | # pass input to the remote embedding table and fetch emb tensor back |