1# Lint as: python3
2# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#         http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16
17"""Test for train.py."""
18
19from __future__ import absolute_import
20from __future__ import division
21from __future__ import print_function
22
23import unittest
24
25import numpy as np
26import tensorflow as tf
27from train import build_cnn
28from train import build_lstm
29from train import load_data
30from train import reshape_function
31
32
33class TestTrain(unittest.TestCase):
34
35    def setUp(self):    # pylint: disable=g-missing-super-call
36        self.seq_length = 128
37        self.train_len, self.train_data, self.valid_len, self.valid_data, \
38                self.test_len, self.test_data = \
39                load_data("./data/train", "./data/valid", "./data/test",
40                                    self.seq_length)
41
42    def test_load_data(self):
43        self.assertIsInstance(self.train_data, tf.data.Dataset)
44        self.assertIsInstance(self.valid_data, tf.data.Dataset)
45        self.assertIsInstance(self.test_data, tf.data.Dataset)
46
47    def test_build_net(self):
48        cnn, cnn_path = build_cnn(self.seq_length)
49        lstm, lstm_path = build_lstm(self.seq_length)
50        cnn_data = np.random.rand(60, 128, 3, 1)
51        lstm_data = np.random.rand(60, 128, 3)
52        cnn_prob = cnn(tf.constant(cnn_data, dtype="float32")).numpy()
53        lstm_prob = lstm(tf.constant(lstm_data, dtype="float32")).numpy()
54        self.assertIsInstance(cnn, tf.keras.Sequential)
55        self.assertIsInstance(lstm, tf.keras.Sequential)
56        self.assertEqual(cnn_path, "./netmodels/CNN")
57        self.assertEqual(lstm_path, "./netmodels/LSTM")
58        self.assertEqual(cnn_prob.shape, (60, 4))
59        self.assertEqual(lstm_prob.shape, (60, 4))
60
61    def test_reshape_function(self):
62        for data, label in self.train_data:
63            original_data_shape = data.numpy().shape
64            original_label_shape = label.numpy().shape
65            break
66        self.train_data = self.train_data.map(reshape_function)
67        for data, label in self.train_data:
68            reshaped_data_shape = data.numpy().shape
69            reshaped_label_shape = label.numpy().shape
70            break
71        self.assertEqual(
72                reshaped_data_shape,
73                (int(original_data_shape[0] * original_data_shape[1] / 3), 3, 1))
74        self.assertEqual(reshaped_label_shape, original_label_shape)
75
76
77if __name__ == "__main__":
78    unittest.main()
79