| 203 | |
| 204 | @unittest.skipIf(not singa_api.USE_CUDA, 'CUDA is not enabled') |
| 205 | def test_lstm_model(self, dev=gpu_dev): |
| 206 | hidden_size = 3 |
| 207 | seq_length = 2 |
| 208 | batch_size = 4 |
| 209 | feature_size = 3 |
| 210 | bidirectional = False |
| 211 | directions = 2 if bidirectional else 1 |
| 212 | num_layers = 2 |
| 213 | out_size = hidden_size |
| 214 | return_sequences = False |
| 215 | batch_first = True |
| 216 | rnn_mode = "lstm" |
| 217 | |
| 218 | # manual test case |
| 219 | x_data = np.array([[[0, 0, 1], [0, 1, 0]], [[0, 1, 0], [1, 0, 0]], |
| 220 | [[0, 0, 1], [0, 1, 0]], [[1, 0, 0], [0, 0, 1]]], |
| 221 | dtype=np.float32).reshape(batch_size, seq_length, |
| 222 | hidden_size) # bs, seq, fea |
| 223 | if return_sequences: |
| 224 | y_data = np.array( |
| 225 | [[[0, 1, 0], [1, 0, 0]], [[1, 0, 0], [0, 0, 1]], |
| 226 | [[0, 1, 0], [1, 0, 0]], [[0, 0, 1], [0, 1, 0]]], |
| 227 | dtype=np.float32).reshape(batch_size, seq_length, |
| 228 | hidden_size) # bs, hidden |
| 229 | y_data.reshape(batch_size, -1) |
| 230 | else: |
| 231 | y_data = np.array([[1, 0, 0], [0, 0, 1], [1, 0, 0], [0, 1, 0]], |
| 232 | dtype=np.float32).reshape( |
| 233 | batch_size, hidden_size) # bs, hidden |
| 234 | |
| 235 | x = tensor.Tensor(device=dev, data=x_data) |
| 236 | y_t = tensor.Tensor(device=dev, data=y_data) |
| 237 | |
| 238 | m = LSTMModel(hidden_size, seq_length, batch_size, bidirectional, |
| 239 | num_layers, return_sequences, rnn_mode, batch_first) |
| 240 | m.compile([x], is_train=True, use_graph=False, sequential=False) |
| 241 | |
| 242 | m.train() |
| 243 | for i in range(1000): |
| 244 | y = m.forward(x) |
| 245 | assert y.shape == y_t.shape |
| 246 | loss = autograd.softmax_cross_entropy(y, y_t) |
| 247 | if i % 100 == 0: |
| 248 | print("loss", loss) |
| 249 | m.optimizer(loss) |
| 250 | |
| 251 | m.eval() |
| 252 | y = m.forward(x) |
| 253 | loss = autograd.softmax_cross_entropy(y, y_t) |
| 254 | print("eval loss", loss) |
| 255 | |
| 256 | |
| 257 | class TestModelSaveMethods(unittest.TestCase): |