update

Author: Daniel Park

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Release notes
 
+
 ## v.0.0.15
 - fix classification/anomaly detection
 - fix from_pretrained

examples/run_prediction_simple.py

@@ -29,6 +29,12 @@ def parse_args():
     parser.add_argument("--epochs", type=int, default=50, help="Number of training epochs")
     parser.add_argument("--batch_size", type=int, default=16, help="Batch size for training")
     parser.add_argument("--learning_rate", type=float, default=1e-3, help="learning rate")
+    parser.add_argument(
+        "--strategy",
+        type=str,
+        default="auto",
+        help="Distribution strategy: auto/default/one_device/mirrored/multi_worker",
+    )
     parser.add_argument("--manual", action="store_true", help="Use the manual API instead of pipeline")
     return parser.parse_args()
 
@@ -90,7 +96,13 @@ def run_manual(args):
     config = tfts.AutoConfig.for_model(args.use_model)
     model = tfts.AutoModel.from_config(config, predict_sequence_length=args.predict_sequence_length)
 
-    trainer = tfts.Trainer(model)
+    trainer = tfts.Trainer(
+        model,
+        args=tfts.TrainingArguments(
+            output_dir=os.path.join(os.getcwd(), "output"),
+            strategy=getattr(args, "strategy", "auto"),
+        ),
+    )
     trainer.train(
         train,
         valid,

tests/test_demo.py

@@ -8,6 +8,12 @@
 
 import tfts
 from tfts import AutoConfig, AutoModel, KerasTrainer as Trainer
+from tfts.training_args import TrainingArguments
+
+# Smoke tests validate end-to-end training, not device distribution, so they pin a
+# single-device strategy. This keeps them deterministic on CI (no GPU) and on local
+# multi-GPU hosts alike.
+_SINGLE_DEVICE_ARGS = TrainingArguments(output_dir="./weights", strategy="default")
 
 
 class DemoTest(unittest.TestCase):
@@ -22,7 +28,7 @@ def test_demo(self):
         config = AutoConfig.for_model("seq2seq")
         model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
 
-        trainer = Trainer(model)
+        trainer = Trainer(model, args=_SINGLE_DEVICE_ARGS)
         trainer.train((x_train, y_train), (x_valid, y_valid), epochs=1)
 
         pred = trainer.predict(x_valid)
@@ -40,7 +46,7 @@ def test_demo2(self):
         model = AutoModel.from_config(config=config, predict_sequence_length=predict_sequence_length)
         print(x_train.shape, y_train.shape, x_valid.shape, y_valid.shape)
 
-        trainer = Trainer(model)
+        trainer = Trainer(model, args=_SINGLE_DEVICE_ARGS)
         trainer.train((x_train, y_train), optimizer=tf.keras.optimizers.Adam(0.001), epochs=2)
 
         pred = trainer.predict(x_valid)

tests/test_examples/test_prediction.py

@@ -23,6 +23,7 @@ class args(object):
             epochs = 1
             batch_size = 32
             learning_rate = 0.003
+            strategy = "default"
 
         set_seed(args.seed)
         run_manual(args)

tests/test_examples/test_tfts_inputs.py

@@ -5,9 +5,14 @@
 import tensorflow as tf
 
 from tfts import AutoConfig, AutoModel, KerasTrainer
+from tfts.training_args import TrainingArguments
 
 logger = logging.getLogger(__name__)
 
+# Smoke test pinning a single-device strategy so it runs identically on CI and
+# any multi-GPU host.
+_SINGLE_DEVICE_ARGS = TrainingArguments(output_dir="./weights", strategy="default")
+
 
 class InputsTest(unittest.TestCase):
     def setUp(self):
@@ -26,7 +31,7 @@ def test_encoder_array(self):
             print(f"==== Test model {m} ====")
             config = AutoConfig.for_model(m)
             model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
-            trainer = KerasTrainer(model)
+            trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
             trainer.train(
                 train_dataset=(x_train, y_train),
                 valid_dataset=(x_valid, y_valid),
@@ -55,7 +60,7 @@ def test_encoder_decoder_array(self):
         for m in self.test_models:
             config = AutoConfig.for_model(m)
             model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
-            trainer = KerasTrainer(model)
+            trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
             trainer.train((x_train, y_train), (x_valid, y_valid), optimizer=tf.keras.optimizers.Adam(0.003), epochs=1)
 
     def test_encoder_decoder_array2(self):
@@ -82,7 +87,7 @@ def test_encoder_decoder_array2(self):
             print(f"==== Test model {m} ====")
             config = AutoConfig.for_model(m)
             model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
-            trainer = KerasTrainer(model)
+            trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
             trainer.train((x_train, y_train), (x_valid, y_valid), optimizer=tf.keras.optimizers.Adam(0.003), epochs=1)
 
     # def test_encoder_tfdata(self):
@@ -121,7 +126,7 @@ def test_encoder_decoder_tfdata(self):
             print(f"==== Test model {m} ====")
             config = AutoConfig.for_model(m)
             model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
-            trainer = KerasTrainer(model)
+            trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
             trainer.train(
                 train_dataset=train_loader,
                 valid_dataset=valid_loader,

tests/test_models/test_dlinear.py

@@ -6,6 +6,11 @@
 import tfts
 from tfts import AutoConfig, AutoModel, KerasTrainer, Trainer
 from tfts.models.dlinear import DLinear, DLinearConfig
+from tfts.training_args import TrainingArguments
+
+# Smoke test pinning a single-device strategy so it runs identically on CI and
+# any multi-GPU host.
+_SINGLE_DEVICE_ARGS = TrainingArguments(output_dir="./weights", strategy="default")
 
 
 class DLinearTest(unittest.TestCase):
@@ -26,9 +31,7 @@ def test_train(self):
         config.channels = 1  # number of features
 
         model = AutoModel.from_config(config, predict_sequence_length=8)
-        trainer = KerasTrainer(
-            model,
-        )
+        trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
 
         trainer.train(train, valid, optimizer=tf.keras.optimizers.Adam(0.003), epochs=1)
 

tests/test_models/test_informer.py

@@ -11,6 +11,11 @@
 from tfts import AutoConfig, AutoModel, KerasTrainer, Trainer
 from tfts.layers.attention_layer import Attention, ProbAttention
 from tfts.models.informer import Decoder, DecoderLayer, DistilConv, Encoder, EncoderLayer, Informer
+from tfts.training_args import TrainingArguments
+
+# Smoke test pinning a single-device strategy so it runs identically on CI and
+# any multi-GPU host.
+_SINGLE_DEVICE_ARGS = TrainingArguments(output_dir="./weights", strategy="default")
 
 tf.config.run_functions_eagerly(True)
 
@@ -144,5 +149,5 @@ def test_train(self):
 
         config = AutoConfig.for_model("informer")
         model = AutoModel.from_config(config, predict_sequence_length)
-        trainer = KerasTrainer(model)
+        trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
         trainer.train((x_train, y_train), (x_valid, y_valid), optimizer=tf.keras.optimizers.Adam(0.003), epochs=1)

tests/test_models/test_nbeats.py

@@ -8,6 +8,11 @@
 from tfts import AutoConfig, AutoModel, KerasTrainer
 from tfts.data import TimeSeriesSequence, get_data
 from tfts.models.nbeats import NBeats
+from tfts.training_args import TrainingArguments
+
+# Smoke test pinning a single-device strategy so it runs identically on CI and
+# any multi-GPU host.
+_SINGLE_DEVICE_ARGS = TrainingArguments(output_dir="./weights", strategy="default")
 
 
 class NBeatsTest(unittest.TestCase):
@@ -40,5 +45,5 @@ def test_train(self):
         config = AutoConfig.for_model("tft")
 
         model = AutoModel.from_config(config, predict_sequence_length=predict_sequence_length)
-        trainer = KerasTrainer(model)
+        trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
         trainer.train(ts_sequence, epochs=1)

tests/test_models/test_rnn.py

@@ -5,6 +5,11 @@
 import tfts
 from tfts import AutoConfig, AutoModel, KerasTrainer, Trainer
 from tfts.models.rnn import RNN, RNN2
+from tfts.training_args import TrainingArguments
+
+# Smoke test pinning a single-device strategy so it runs identically on CI and
+# any multi-GPU host.
+_SINGLE_DEVICE_ARGS = TrainingArguments(output_dir="./weights", strategy="default")
 
 
 class RNNTest(unittest.TestCase):
@@ -20,7 +25,7 @@ def test_train(self):
         train, valid = tfts.get_data("sine", test_size=0.1)
         config = AutoConfig.for_model("rnn")
         model = AutoModel.from_config(config, predict_sequence_length=8)
-        trainer = KerasTrainer(model)
+        trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
         trainer.train(train, valid, optimizer=tf.keras.optimizers.Adam(0.003), epochs=1)
         y_test = trainer.predict(valid[0])
         self.assertEqual(y_test.shape, valid[1].shape)

tests/test_models/test_tcn.py

@@ -5,6 +5,11 @@
 import tfts
 from tfts import AutoConfig, AutoModel, KerasTrainer, Trainer
 from tfts.models.tcn import TCN
+from tfts.training_args import TrainingArguments
+
+# Smoke test pinning a single-device strategy so it runs identically on CI and
+# any multi-GPU host.
+_SINGLE_DEVICE_ARGS = TrainingArguments(output_dir="./weights", strategy="default")
 
 
 class TCNTest(unittest.TestCase):
@@ -20,9 +25,7 @@ def test_train(self):
         train, valid = tfts.get_data("sine", test_size=0.1)
         config = AutoConfig.for_model("tcn")
         model = AutoModel.from_config(config=config, predict_sequence_length=8)
-        trainer = KerasTrainer(
-            model,
-        )
+        trainer = KerasTrainer(model, args=_SINGLE_DEVICE_ARGS)
         trainer.train(train, valid, optimizer=tf.keras.optimizers.Adam(0.003), epochs=1)
         y_test = trainer.predict(valid[0])
         self.assertEqual(y_test.shape, valid[1].shape)