Add testing gin configs.

PiperOrigin-RevId: 246423640

Author: T2T Team

tensor2tensor/trax/configs/resnet50_imagenet_8gb_testing.gin

@@ -0,0 +1,44 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.learning_rate
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 32
+batch_fun.bucket_length = 32
+batch_fun.buckets = None
+batch_fun.eval_batch_size = 32
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_image_imagenet224'
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+EvalAdjustingSchedule.constant = 1.0
+MultifactorSchedule.factors = 'constant * linear_warmup'
+MultifactorSchedule.warmup_steps = 400
+
+# Parameters for momentum:
+# ==============================================================================
+momentum.mass = 0.9
+
+
+# Parameters for Resnet50:
+# ==============================================================================
+Resnet50.hidden_size = 64
+Resnet50.num_output_classes = 1001
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 2000
+train.eval_steps = 20
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.Resnet50
+train.optimizer = @trax.optimizers.momentum
+train.train_steps = 100000
+train.lr_schedule = @learning_rate.EvalAdjustingSchedule
+

tensor2tensor/trax/configs/transformer_lm1b_8gb_testing.gin

@@ -0,0 +1,52 @@
+import tensor2tensor.trax.inputs
+import tensor2tensor.trax.models
+import tensor2tensor.trax.optimizers
+import tensor2tensor.trax.trax
+
+# Parameters for batch_fun:
+# ==============================================================================
+batch_fun.batch_size_per_device = 128
+batch_fun.eval_batch_size = 128
+batch_fun.max_eval_length = 2048
+
+# Parameters for inputs:
+# ==============================================================================
+inputs.data_dir = None
+inputs.dataset_name = 't2t_languagemodel_lm1b32k'
+inputs.input_name = 'targets'
+
+# Parameters for mask:
+# ==============================================================================
+masked_mean.mask_id = 0
+
+# Parameters for MultifactorSchedule:
+# ==============================================================================
+MultifactorSchedule.constant = 0.1
+MultifactorSchedule.factors = 'constant * linear_warmup * rsqrt_decay'
+MultifactorSchedule.warmup_steps = 8000
+
+# Parameters for preprocess_fun:
+# ==============================================================================
+shuffle_and_batch_data.preprocess_fun=@trax.inputs.lm1b_preprocess
+lm1b_preprocess.max_target_length = 512
+lm1b_preprocess.max_eval_target_length = 2048
+
+# Parameters for train:
+# ==============================================================================
+train.eval_frequency = 1000
+train.eval_steps = 10
+train.inputs = @trax.inputs.inputs
+train.model = @trax.models.TransformerLM
+train.run_debug_step = False
+train.train_steps = 100000
+
+# Parameters for TransformerLM:
+# ==============================================================================
+TransformerLM.dropout = 0.1
+TransformerLM.feature_depth = 512
+TransformerLM.feedforward_depth = 2048
+TransformerLM.max_len = 2048
+TransformerLM.mode = 'train'
+TransformerLM.num_heads = 8
+TransformerLM.num_layers = 6
+TransformerLM.vocab_size = 32000