def run(args):
config = UserConfig(DEFAULT_YAML, args.config, learning=True)
speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
text_featurizer = TextFeaturizer(config["decoder_config"])
tf.random.set_seed(0)
assert args.saved_model
if args.tfrecords:
test_dataset = ASRTFRecordDataset(
config["learning_config"]["dataset_config"]["test_paths"],
config["learning_config"]["dataset_config"]["tfrecords_dir"],
speech_featurizer,
text_featurizer,
"test",
augmentations=config["learning_config"]["augmentations"],
shuffle=False).create(
config["learning_config"]["running_config"]["batch_size"])
else:
test_dataset = ASRSliceDataset(
stage="test",
speech_featurizer=speech_featurizer,
text_featurizer=text_featurizer,
data_paths=config["learning_config"]["dataset_config"]
["eval_paths"],
shuffle=False).create(
config["learning_config"]["running_config"]["batch_size"])
# build model
f, c = speech_featurizer.compute_feature_dim()
conformer = Conformer(vocabulary_size=text_featurizer.num_classes,
**config["model_config"])
conformer._build([1, 50, f, c])
conformer.summary(line_length=100)
conformer_tester = BaseTester(
config=config["learning_config"]["running_config"],
saved_path=args.saved_model,
from_weights=args.from_weights)
conformer_tester.compile(conformer, speech_featurizer, text_featurizer)
conformer_tester.run(test_dataset)
parser.add_argument("output",
type=str,
default=None,
help="TFLite file path to be exported")
args = parser.parse_args()
assert args.saved and args.output
config = UserConfig(DEFAULT_YAML, args.config, learning=True)
speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
text_featurizer = CharFeaturizer(config["decoder_config"])
# build model
conformer = Conformer(**config["model_config"],
vocabulary_size=text_featurizer.num_classes)
conformer._build(speech_featurizer.shape)
conformer.load_weights(args.saved)
conformer.summary(line_length=150)
conformer.add_featurizers(speech_featurizer, text_featurizer)
concrete_func = conformer.make_tflite_function(
greedy=True).get_concrete_function()
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
tflite_model = converter.convert()
if not os.path.exists(os.path.dirname(args.output)):
if args.tfrecords:
test_dataset = ASRTFRecordDataset(
data_paths=config["learning_config"]["dataset_config"]["test_paths"],
tfrecords_dir=config["learning_config"]["dataset_config"]
["tfrecords_dir"],
speech_featurizer=speech_featurizer,
text_featurizer=text_featurizer,
stage="test",
shuffle=False)
else:
test_dataset = ASRSliceDataset(
data_paths=config["learning_config"]["dataset_config"]["test_paths"],
speech_featurizer=speech_featurizer,
text_featurizer=text_featurizer,
stage="test",
shuffle=False)
# build model
conformer = Conformer(vocabulary_size=text_featurizer.num_classes,
**config["model_config"])
conformer._build(speech_featurizer.shape)
conformer.load_weights(args.saved, by_name=True)
conformer.summary(line_length=120)
conformer.add_featurizers(speech_featurizer, text_featurizer)
conformer_tester = BaseTester(
config=config["learning_config"]["running_config"],
output_name=args.output_name)
conformer_tester.compile(conformer)
conformer_tester.run(test_dataset)
eval_dataset = ASRSliceDataset(
data_paths=config["learning_config"]["dataset_config"]["eval_paths"],
speech_featurizer=speech_featurizer,
text_featurizer=text_featurizer,
stage="eval", cache=args.cache, shuffle=True
)
conformer_trainer = TransducerTrainer(
config=config["learning_config"]["running_config"],
text_featurizer=text_featurizer, strategy=strategy
)
with conformer_trainer.strategy.scope():
# build model
conformer = Conformer(
**config["model_config"],
vocabulary_size=text_featurizer.num_classes
)
conformer._build(speech_featurizer.shape)
conformer.summary(line_length=120)
optimizer_config = config["learning_config"]["optimizer_config"]
optimizer = tf.keras.optimizers.Adam(
TransformerSchedule(
d_model=config["model_config"]["dmodel"],
warmup_steps=optimizer_config["warmup_steps"],
max_lr=(0.05 / math.sqrt(config["model_config"]["dmodel"]))
),
beta_1=optimizer_config["beta1"],
beta_2=optimizer_config["beta2"],
epsilon=optimizer_config["epsilon"]
)
def run(args):
config = UserConfig(DEFAULT_YAML, args.config, learning=True)
speech_featurizer = TFSpeechFeaturizer(config["speech_config"])
text_featurizer = TextFeaturizer(config["decoder_config"])
tf.random.set_seed(2020)
if args.mixed_precision:
policy = tf.keras.mixed_precision.experimental.Policy(
"mixed_float16")
tf.keras.mixed_precision.experimental.set_policy(policy)
print("Enabled mixed precision training")
if args.tfrecords:
train_dataset = ASRTFRecordDataset(
config["learning_config"]["dataset_config"]["train_paths"],
config["learning_config"]["dataset_config"]["tfrecords_dir"],
speech_featurizer,
text_featurizer,
"train",
augmentations=config["learning_config"]["augmentations"],
shuffle=True,
)
eval_dataset = ASRTFRecordDataset(
config["learning_config"]["dataset_config"]["eval_paths"],
config["learning_config"]["dataset_config"]["tfrecords_dir"],
speech_featurizer,
text_featurizer,
"eval",
shuffle=False)
else:
train_dataset = ASRSliceDataset(
stage="train",
speech_featurizer=speech_featurizer,
text_featurizer=text_featurizer,
data_paths=config["learning_config"]["dataset_config"]
["train_paths"],
augmentations=config["learning_config"]["augmentations"],
shuffle=True,
)
eval_dataset = ASRSliceDataset(stage="eval",
speech_featurizer=speech_featurizer,
text_featurizer=text_featurizer,
data_paths=config["learning_config"]
["dataset_config"]["eval_paths"],
shuffle=False)
conformer_trainer = TransducerTrainer(
config=config["learning_config"]["running_config"],
text_featurizer=text_featurizer,
is_mixed_precision=args.mixed_precision)
with conformer_trainer.strategy.scope():
# build model
f, c = speech_featurizer.compute_feature_dim()
conformer = Conformer(**config["model_config"],
vocabulary_size=text_featurizer.num_classes)
conformer._build([1, 50, f, c])
optimizer_config = config["learning_config"]["optimizer_config"]
optimizer = tf.keras.optimizers.Adam(
TransformerSchedule(
d_model=config["model_config"]["dmodel"],
warmup_steps=optimizer_config["warmup_steps"],
max_lr=(0.05 /
math.sqrt(config["model_config"]["dmodel"]))),
beta_1=float(optimizer_config["beta1"]),
beta_2=float(optimizer_config["beta2"]),
epsilon=float(optimizer_config["epsilon"]))
conformer_trainer.compile(model=conformer,
optimizer=optimizer,
max_to_keep=args.max_ckpts)
conformer_trainer.fit(train_dataset, eval_dataset,
args.eval_train_ratio)
if args.export:
if args.from_weights:
conformer_trainer.model.save_weights(args.export)
else:
conformer_trainer.model.save(args.export)