| 21 | |
| 22 | |
| 23 | class IMDBModel(model.Model): |
| 24 | |
| 25 | def __init__(self, |
| 26 | hidden_size, |
| 27 | mode='lstm', |
| 28 | return_sequences=False, |
| 29 | bidirectional="False", |
| 30 | num_layers=1): |
| 31 | super().__init__() |
| 32 | batch_first = True |
| 33 | self.lstm = layer.CudnnRNN(hidden_size=hidden_size, |
| 34 | batch_first=batch_first, |
| 35 | rnn_mode=mode, |
| 36 | return_sequences=return_sequences, |
| 37 | num_layers=1, |
| 38 | dropout=0.9, |
| 39 | bidirectional=bidirectional) |
| 40 | self.l1 = layer.Linear(64) |
| 41 | self.l2 = layer.Linear(2) |
| 42 | |
| 43 | def forward(self, x): |
| 44 | y = self.lstm(x) |
| 45 | y = autograd.reshape(y, (y.shape[0], -1)) |
| 46 | y = self.l1(y) |
| 47 | y = autograd.relu(y) |
| 48 | y = self.l2(y) |
| 49 | return y |
| 50 | |
| 51 | def train_one_batch(self, x, y): |
| 52 | out = self.forward(x) |
| 53 | loss = autograd.softmax_cross_entropy(out, y) |
| 54 | self.optimizer(loss) |
| 55 | return out, loss |
| 56 | |
| 57 | def set_opt(self, optimizer): |
| 58 | self.optimizer = optimizer |