def set_model(self):
''' Setup ASR model '''
# Model
self.feat_dim = 120
self.vocab_size = 46
init_adadelta = True
''' Setup ASR model and optimizer '''
# Model
# init_adadelta = self.config['hparas']['optimizer'] == 'Adadelta'
self.model = ASR(self.feat_dim, self.vocab_size, init_adadelta, **
self.src_config['model']).to(self.device)
self.verbose(self.model.create_msg())
if self.finetune_first>0:
names = ["encoder.layers.%d"%i for i in range(self.finetune_first)]
model_paras = [{"params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in names)]}]
else:
model_paras = [{'params': self.model.parameters()}]
# Losses
self.seq_loss = torch.nn.CrossEntropyLoss(ignore_index=0)
# Note: zero_infinity=False is unstable?
self.ctc_loss = torch.nn.CTCLoss(blank=0, zero_infinity=False)
# Plug-ins
self.emb_fuse = False
self.emb_reg = ('emb' in self.config) and (
self.config['emb']['enable'])
if self.emb_reg:
from src.plugin import EmbeddingRegularizer
self.emb_decoder = EmbeddingRegularizer(
self.tokenizer, self.model.dec_dim, **self.config['emb']).to(self.device)
model_paras.append({'params': self.emb_decoder.parameters()})
self.emb_fuse = self.emb_decoder.apply_fuse
if self.emb_fuse:
self.seq_loss = torch.nn.NLLLoss(ignore_index=0)
self.verbose(self.emb_decoder.create_msg())
# Optimizer
self.optimizer = Optimizer(model_paras, **self.src_config['hparas'])
self.verbose(self.optimizer.create_msg())
# Enable AMP if needed
self.enable_apex()
# Automatically load pre-trained model if self.paras.load is given
self.load_ckpt()
# Beam decoder
self.decoder = BeamDecoder(
self.model, self.emb_decoder, **self.config['decode'])
self.verbose(self.decoder.create_msg())
# del self.model
# del self.emb_decoder
self.decoder.to(self.device)
def set_model(self):
''' Setup ASR model '''
# Model
init_adadelta = self.config['hparas']['optimizer'] == 'Adadelta'
self.model = ASR(self.feat_dim, self.vocab_size, init_adadelta,
**self.config['model']).to(self.device)
# Plug-ins
if ('emb' in self.config) and (self.config['emb']['enable']) \
and (self.config['emb']['fuse'] > 0):
from src.plugin import EmbeddingRegularizer
self.emb_decoder = EmbeddingRegularizer(self.tokenizer,
self.model.dec_dim,
**self.config['emb'])
# Load target model in eval mode
self.load_ckpt()
self.ctc_only = False
if self.greedy:
# Greedy decoding: attention-based if the ASR has a decoder, else use CTC
self.decoder = copy.deepcopy(self.model).to(self.device)
else:
if (not self.model.enable_att) or self.config['decode'].get(
'ctc_weight', 0.0) == 1.0:
# Pure CTC Beam Decoder
assert self.config['decode']['beam_size'] <= self.config[
'decode']['vocab_candidate']
self.decoder = CTCBeamDecoder(
self.model.to(self.device),
[1] + [r for r in range(3, self.vocab_size)],
self.config['decode']['beam_size'],
self.config['decode']['vocab_candidate'],
lm_path=self.config['decode']['lm_path'],
lm_config=self.config['decode']['lm_config'],
lm_weight=self.config['decode']['lm_weight'],
device=self.device)
self.ctc_only = True
else:
# Joint CTC-Attention Beam Decoder
self.decoder = BeamDecoder(self.model.cpu(), self.emb_decoder,
**self.config['decode'])
self.verbose(self.decoder.create_msg())
del self.model
del self.emb_decoder
def set_model(self):
''' Setup ASR model '''
# Model
self.model = ASR(self.feat_dim, self.vocab_size,
**self.config['model'])
# Plug-ins
if ('emb' in self.config) and (self.config['emb']['enable']) \
and (self.config['emb']['fuse']>0):
from src.plugin import EmbeddingRegularizer
self.emb_decoder = EmbeddingRegularizer(self.tokenizer,
self.model.dec_dim,
**self.config['emb'])
# Load target model in eval mode
self.load_ckpt()
# self.ctc_only = False
if self.greedy:
self.decoder = copy.deepcopy(self.model).to(self.device)
else:
# Beam decoder
# TODO: CTC decoding function Hidden by author
# if not self.model.enable_att or self.config['decode'].get('ctc_weight', 0.0) == 1.0:
# For CTC only decoding (character level)
# self.decoder = CTCBeamDecoder(self.model.to(self.device),
# range(self.model.vocab_size),
# self.config['decode']['beam_size'],
# self.config['decode']['vocab_candidate'])
# self.ctc_only = True
# else:
# self.decoder = BeamDecoder(self.model, self.emb_decoder, **self.config['decode'])
self.decoder = BeamDecoder(self.model, self.emb_decoder,
**self.config['decode'])
self.verbose(self.decoder.create_msg())
del self.model
del self.emb_decoder
self.emb_decoder = None
def set_model(self):
''' Setup ASR model '''
# Model
init_adadelta = self.src_config['hparas']['optimizer'] == 'Adadelta'
self.model = ASR(self.feat_dim, self.vocab_size, init_adadelta, **
self.config['model']).to(self.device)
# Plug-ins
if ('emb' in self.config) and (self.config['emb']['enable']) \
and (self.config['emb']['fuse'] > 0):
from src.plugin import EmbeddingRegularizer
self.emb_decoder = EmbeddingRegularizer(
self.tokenizer, self.model.dec_dim, **self.config['emb'])
# Load target model in eval mode
self.load_ckpt()
# Beam decoder
self.decoder = BeamDecoder(
self.model.cpu(), self.emb_decoder, **self.config['decode'])
self.verbose(self.decoder.create_msg())
del self.model
del self.emb_decoder