update

Author: LongxingTan

tests/test_examples/test_tfts_inputs.py

@@ -23,7 +23,7 @@ def test_encoder_array(self):
         y_valid = np.random.rand(1, predict_sequence_length, 1)
 
         for m in self.test_models:
-            logger.info(f"Test model {m}")
+            print(f"==== Test model {m} ====")
             config = AutoConfig.for_model(m)
             model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
             trainer = KerasTrainer(model)
@@ -65,6 +65,7 @@ def test_encoder_decoder_array2(self):
         n_decoder_feature = 3
 
         x_train = (
+            # x, encoder, decoder
             np.random.rand(1, train_length, 1),
             np.random.rand(1, train_length, n_encoder_feature),
             np.random.rand(1, predict_sequence_length, n_decoder_feature),
@@ -78,6 +79,7 @@ def test_encoder_decoder_array2(self):
         y_valid = np.random.rand(1, predict_sequence_length, 1)
 
         for m in self.test_models:
+            print(f"==== Test model {m} ====")
             config = AutoConfig.for_model(m)
             model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
             trainer = KerasTrainer(model)
@@ -116,6 +118,7 @@ def test_encoder_decoder_tfdata(self):
         valid_loader = valid_loader.batch(batch_size=1)
 
         for m in self.test_models:
+            print(f"==== Test model {m} ====")
             config = AutoConfig.for_model(m)
             model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
             trainer = KerasTrainer(model)

tfts/models/seq2seq.py

@@ -214,8 +214,9 @@ def __init__(
         self.num_attention_heads = num_attention_heads
         self.attention_probs_dropout_prob = attention_probs_dropout_prob
 
-    def build(self, decoder_features_shape, decoder_init_input_shape, init_state_shape, **kwargs):
-        rnn_input_size = decoder_features_shape[-1] + decoder_init_input_shape[-1]
+    def build(self, input_shape, **kwargs):
+        batch_size = input_shape[0]
+        rnn_input_size = input_shape[-1] + 1
 
         if self.use_attention:
             encoder_output_shape = kwargs.get("encoder_output_shape")
@@ -228,21 +229,18 @@ def build(self, decoder_features_shape, decoder_init_input_shape, init_state_sha
             )
             self.attention.build(encoder_output_shape)
 
-            # Add attention output size to RNN input size
-            rnn_input_size += encoder_output_shape[-1]
-
         if self.rnn_type == "gru":
             self.rnn_cell = GRUCell(self.rnn_size)
         elif self.rnn_type == "lstm":
             self.rnn_cell = LSTMCell(units=self.rnn_size)
         else:
             raise ValueError(f"Unsupported rnn type: {self.rnn_type}")
 
-        self.rnn_cell.build([None, rnn_input_size])
+        self.rnn_cell.build([batch_size, rnn_input_size])
 
         self.dense = Dense(units=1, activation=None)
-        self.dense.build([None, self.rnn_size])
-        super().build(decoder_features_shape)
+        self.dense.build([batch_size, self.rnn_size])
+        super().build(input_shape)
 
     def call(
         self,

tfts/models/wavenet.py

@@ -111,11 +111,11 @@ def __call__(
         encoder_state, encoder_outputs = self.encoder(encoder_feature)
         decoder_outputs = self.decoder(
             decoder_features=decoder_feature,
-            decoder_init_input=x[:, -1],
+            # imagine the first dim is the predict
+            decoder_init_input=x[:, -1, 0:1],
             teacher=teacher,
             encoder_outputs=encoder_outputs,
         )
-
         return decoder_outputs
 
 
@@ -200,9 +200,9 @@ def __init__(
         self.dilation_rates = dilation_rates
         self.dense_hidden_size = dense_hidden_size
 
-    def build(self, decoder_features_shape, decoder_init_input_shape, encoder_outputs_shape=None, **kwargs):
-        batch_size = decoder_features_shape[0]
-        decoder_input_size = decoder_features_shape[-1] + decoder_init_input_shape[-1]
+    def build(self, input_shape, **kwargs):
+        batch_size = input_shape[0]
+        decoder_input_size = input_shape[-1] + 1
 
         self.dense1 = Dense(self.filters, activation="tanh")
         self.dense1.build([batch_size, decoder_input_size])
@@ -223,7 +223,7 @@ def build(self, decoder_features_shape, decoder_init_input_shape, encoder_output
         self.dense6 = Dense(1)
         self.dense6.build([batch_size, self.dense_hidden_size])
 
-        super().build(decoder_features_shape)
+        super().build(input_shape)
 
     def call(
         self,