def load_model(path):
package = torch.load(path, map_location=lambda storage, loc: storage)
if package['name'] == 'deepspeech':
from deepspeech.model import DeepSpeech
model = DeepSpeech(package['model_conf'], package['audio_conf'],
package['labels'])
else:
raise NotImplementedError
model.load_state_dict(package['state_dict'])
return model, package
def __init__(self,
continue_from=None,
sample_rate=None,
window_size=None,
window_stride=None,
window=None,
noise_dir=None,
noise_prob=None,
noise_min=None,
noise_max=None,
labels_path=None,
hidden_size=None,
hidden_layers=None,
bidirectional=None,
rnn_type=None):
if continue_from: # Starting from previous model
print("Loading checkpoint model %s" % continue_from)
package = torch.load(continue_from,
map_location=lambda storage, loc: storage)
model = DeepSpeech.load_model_package(package)
else:
with open(labels_path) as label_file:
labels = str(''.join(json.load(label_file)))
audio_conf = dict(sample_rate=sample_rate,
window_size=window_size,
window_stride=window_stride,
window=window,
noise_dir=noise_dir,
noise_prob=noise_prob,
noise_levels=(noise_min, noise_max))
rnn_type = rnn_type.lower()
assert rnn_type in supported_rnns, "rnn_type should be either lstm, rnn or gru"
model = DeepSpeech(rnn_hidden_size=hidden_size,
nb_layers=hidden_layers,
labels=labels,
rnn_type=supported_rnns[rnn_type],
audio_conf=audio_conf,
bidirectional=bidirectional)
self._model = model