def test_main(deepspeech: DeepSpeech, config: Configuration):
freeze(deepspeech.model)
gpus = get_available_gpus()
extended_model = create_extended_model(deepspeech.model, config, is_gpu=len(gpus) > 0)
assert len(deepspeech.model.layers) + len(config.data['extension']['layers']) == len(extended_model.layers)
*_, first_lstm, second_lstm, softmax = extended_model.layers
assert not first_lstm.trainable
assert second_lstm.trainable and softmax.trainable
optimizer = DeepSpeech.get_optimizer(**config.optimizer)
loss = DeepSpeech.get_loss()
gpus = get_available_gpus()
deepspeech.compiled_model = DeepSpeech.compile_model(extended_model, optimizer, loss, gpus)
assert len(deepspeech.compiled_model.layers) == len(extended_model.layers)
*_, first_lstm, second_lstm, softmax = deepspeech.compiled_model.layers
assert not first_lstm.trainable
assert second_lstm.trainable and softmax.trainable
def load_extended_model(config_path, alphabet_path, pretrained_weights):
deepspeech = DeepSpeech.construct(config_path=config_path,
alphabet_path=alphabet_path)
freeze(deepspeech.model)
gpus = get_available_gpus()
config = Configuration(config_path)
extended_model = create_extended_model(deepspeech.model,
config,
is_gpu=len(gpus) > 0)
optimizer = DeepSpeech.get_optimizer(**config.optimizer)
loss = DeepSpeech.get_loss()
gpus = get_available_gpus()
deepspeech.model = extended_model
deepspeech.compiled_model = DeepSpeech.compile_model(
extended_model, optimizer, loss, gpus)
deepspeech.load(pretrained_weights)
return deepspeech
def main(args):
deepspeech = DeepSpeech.construct(config_path=CONFIG_PATH, alphabet_path=ALPHABET_PATH)
if args.pretrained_weights:
deepspeech.load(args.pretrained_weights)
freeze(deepspeech.model)
gpus = get_available_gpus()
config = Configuration(CONFIG_PATH)
extended_model = create_extended_model(deepspeech.model, config, is_gpu=len(gpus) > 0)
optimizer = DeepSpeech.get_optimizer(**config.optimizer)
loss = DeepSpeech.get_loss()
gpus = get_available_gpus()
deepspeech.model = extended_model
deepspeech.compiled_model = DeepSpeech.compile_model(extended_model, optimizer, loss, gpus)
train_generator, dev_generator = create_generators(deepspeech, args)
deepspeech.fit(train_generator, dev_generator, epochs=args.epochs, shuffle=False)
deepspeech.save(WEIGHTS_PATH)
def test_compile_model(config: Configuration):
model = DeepSpeech.get_model(**config.model, is_gpu=False)
optimizer = DeepSpeech.get_optimizer(**config.optimizer)
loss = DeepSpeech.get_loss()
DeepSpeech.compile_model(model, optimizer, loss)
assert model._is_compiled