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)