def get_model(model_name, model_config, cpu_run,
uniform_initialize_bn_weight=False, forward_is_infer=False):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
if forward_is_infer:
class Tacotron2__forward_is_infer(Tacotron2):
def forward(self, inputs, input_lengths):
return self.infer(inputs, input_lengths)
model = Tacotron2__forward_is_infer(**model_config)
else:
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
if forward_is_infer:
class WaveGlow__forward_is_infer(WaveGlow):
def forward(self, spect, sigma=1.0):
return self.infer(spect, sigma)
model = WaveGlow__forward_is_infer(**model_config)
else:
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if uniform_initialize_bn_weight:
init_bn(model)
if not cpu_run:
model = model.cuda()
return model
def load_model(hparams):
model = Tacotron2(hparams).to(device)
if hparams.fp16_run:
model.decoder.attention_layer.score_mask_value = finfo('float16').min
if hparams.distributed_run:
model = apply_gradient_allreduce(model)
return model
def get_model(model_name, model_config, device,
uniform_initialize_bn_weight=False, forward_is_infer=False,
jitable=False):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
if forward_is_infer:
class Tacotron2__forward_is_infer(Tacotron2):
def forward(self, inputs, input_lengths):
return self.infer(inputs, input_lengths)
model = Tacotron2__forward_is_infer(**model_config)
else:
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
if forward_is_infer:
class WaveGlow__forward_is_infer(WaveGlow):
def forward(self, spect, sigma=1.0):
return self.infer(spect, sigma)
model = WaveGlow__forward_is_infer(**model_config)
else:
model = WaveGlow(**model_config)
elif model_name == 'FastPitch':
if forward_is_infer:
if jitable:
class FastPitch__forward_is_infer(_FastPitchJIT):
def forward(self, inputs, input_lengths, pace: float = 1.0,
dur_tgt: Optional[torch.Tensor] = None,
pitch_tgt: Optional[torch.Tensor] = None,
pitch_transform: Optional[bool] = None):
return self.infer(inputs, input_lengths, pace=pace,
dur_tgt=dur_tgt, pitch_tgt=pitch_tgt)
else:
class FastPitch__forward_is_infer(_FastPitch):
def forward(self, inputs, input_lengths, pace: float = 1.0,
dur_tgt: Optional[torch.Tensor] = None,
pitch_tgt: Optional[torch.Tensor] = None,
pitch_transform=None):
return self.infer(inputs, input_lengths, pace=pace,
dur_tgt=dur_tgt, pitch_tgt=pitch_tgt,
pitch_transform=pitch_transform)
model = FastPitch__forward_is_infer(**model_config)
else:
model = _FastPitch(**model_config)
else:
raise NotImplementedError(model_name)
if uniform_initialize_bn_weight:
init_bn(model)
return model.to(device)
def get_model(model_name, model_config, to_cuda):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if to_cuda:
model = model.cuda()
return model
def get_model(model_name: str,
model_config: dict,
cpu_run: bool,
uniform_initialize_bn_weight: bool = False,
forward_is_infer: bool = False):
"""Return a model based on `model_name` defined by `model_config`.
Args:
model_name (str): One of the 'Tacotron2' or 'WaveGlow'.
model_config (dict): [description]
cpu_run (bool): Run on CPU (True) or GPU (False).
uniform_initialize_bn_weight (bool, optional): [description]. Defaults to False.
forward_is_infer (bool, optional): [description]. Defaults to False.
Raises:
NotImplementedError: [description]
Returns:
[type]: [description]
"""
model = None
if model_name == 'Tacotron2':
if forward_is_infer:
class Tacotron2__forward_is_infer(Tacotron2):
def forward(self, inputs, input_lengths):
print(input_lengths) #FIXME
return self.infer(inputs, input_lengths)
model = Tacotron2__forward_is_infer(**model_config)
else:
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
if forward_is_infer:
class WaveGlow__forward_is_infer(WaveGlow):
def forward(self, spect, sigma=1.0):
return self.infer(spect, sigma)
model = WaveGlow__forward_is_infer(**model_config)
else:
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if uniform_initialize_bn_weight:
init_bn(model)
if cpu_run == False:
model = model.cuda()
return model
def get_model(model_name, model_config, to_cuda,
uniform_initialize_bn_weight=False):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if uniform_initialize_bn_weight:
init_bn(model)
if to_cuda:
model = model.cuda()
return model
def get_model(model_name, model_config, to_fp16, to_cuda, training=True):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if to_fp16:
model = batchnorm_to_float(model.half())
model = lstmcell_to_float(model)
if model_name == "WaveGlow":
for k in model.convinv:
k.float()
if to_cuda:
model = model.cuda()
return model
def get_tacotron2_model(args, is_training=True):
config = dict(
# optimization
mask_padding=args.mask_padding,
# audio
n_mel_channels=args.n_mel_channels,
# symbols
n_symbols=args.n_symbols,
symbols_embedding_dim=args.symbols_embedding_dim,
# encoder
encoder_kernel_size=args.encoder_kernel_size,
encoder_n_convolutions=args.encoder_n_convolutions,
encoder_embedding_dim=args.encoder_embedding_dim,
# attention
attention_dim=args.attention_dim,
# attention location
attention_location_n_filters=args.attention_location_n_filters,
attention_location_kernel_size=args.attention_location_kernel_size,
# decoder
n_frames_per_step=args.n_frames_per_step,
decoder_rnn_dim=args.decoder_rnn_dim,
prenet_dim=args.prenet_dim,
max_decoder_steps=args.max_decoder_steps,
gate_threshold=args.gate_threshold,
decoder_n_lstms=args.decoder_n_lstms,
p_decoder_dropout=args.p_decoder_dropout,
# postnet
postnet_embedding_dim=args.postnet_embedding_dim,
postnet_kernel_size=args.postnet_kernel_size,
postnet_n_convolutions=args.postnet_n_convolutions,
)
model = Tacotron2(**config)
if is_training:
_init_bn(model)
return model.cuda()
def get_model(model_name, model_config, to_cuda,
uniform_initialize_bn_weight=False, rename=False):
""" Code chooses a model based on name"""
model = None
if model_name == 'Tacotron2':
if rename:
class Tacotron2_extra(Tacotron2):
def forward(self, inputs, input_lengths):
return self.infer(inputs, input_lengths)
model = Tacotron2_extra(**model_config)
else:
model = Tacotron2(**model_config)
elif model_name == 'WaveGlow':
model = WaveGlow(**model_config)
else:
raise NotImplementedError(model_name)
if uniform_initialize_bn_weight:
init_bn(model)
if to_cuda:
model = model.cuda()
return model