**vars(config.learning_config.eval_dataset_config)
)
contextnet_trainer = TransducerTrainerGA(
config=config.learning_config.running_config,
text_featurizer=text_featurizer, strategy=strategy
)
with contextnet_trainer.strategy.scope():
# build model
contextnet = ContextNet(**config.model_config, vocabulary_size=text_featurizer.num_classes)
contextnet._build(speech_featurizer.shape)
contextnet.summary(line_length=120)
optimizer = tf.keras.optimizers.Adam(
TransformerSchedule(
d_model=contextnet.dmodel,
warmup_steps=config.learning_config.optimizer_config["warmup_steps"],
max_lr=(0.05 / math.sqrt(contextnet.dmodel))
),
beta_1=config.learning_config.optimizer_config["beta1"],
beta_2=config.learning_config.optimizer_config["beta2"],
epsilon=config.learning_config.optimizer_config["epsilon"]
)
contextnet_trainer.compile(model=contextnet, optimizer=optimizer,
max_to_keep=args.max_ckpts)
contextnet_trainer.fit(train_dataset, eval_dataset,
train_bs=args.tbs, eval_bs=args.ebs, train_acs=args.acs)
conformer_trainer = TransducerTrainerGA(
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"])
conformer_trainer.compile(model=conformer,
optimizer=optimizer,
max_to_keep=args.max_ckpts)
conformer_trainer.fit(train_dataset,
eval_dataset,
train_bs=args.tbs,
eval_bs=args.ebs,
train_acs=args.acs)
augmentations=config.learning_config.augmentations,
stage="train", cache=args.cache, shuffle=True
)
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
)
streaming_transducer_trainer = TransducerTrainerGA(
config=config.learning_config.running_config,
text_featurizer=text_featurizer, strategy=strategy
)
with streaming_transducer_trainer.strategy.scope():
# build model
streaming_transducer = StreamingTransducer(
**config.model_config,
vocabulary_size=text_featurizer.num_classes
)
streaming_transducer._build(speech_featurizer.shape)
streaming_transducer.summary(line_length=150)
optimizer = tf.keras.optimizers.get(config.learning_config.optimizer_config)
streaming_transducer_trainer.compile(model=streaming_transducer, optimizer=optimizer,
max_to_keep=args.max_ckpts)
streaming_transducer_trainer.fit(train_dataset, eval_dataset, train_bs=args.tbs, eval_bs=args.ebs)