def init_backend_engine():
"""
Initializes ``engine``, which is either :class:`TFEngine.Engine` or Theano :class:`Engine.Engine`.
"""
BackendEngine.select_engine(config=config)
if BackendEngine.is_theano_selected():
print("Theano:", describe_theano_version(), file=log.v3)
import returnn.theano.util
returnn.theano.util.monkey_patches()
elif BackendEngine.is_tensorflow_selected():
print("TensorFlow:", describe_tensorflow_version(), file=log.v3)
if get_tensorflow_version_tuple()[0] == 0:
print("Warning: TF <1.0 is not supported and likely broken.",
file=log.v2)
if os.environ.get("TF_DEVICE"):
print("Devices: Use %s via TF_DEVICE instead of %s." %
(os.environ.get("TF_DEVICE"),
config.opt_typed_value("device")),
file=log.v4)
config.set("device", os.environ.get("TF_DEVICE"))
if config.is_true("use_horovod"):
import returnn.tf.horovod
hvd = returnn.tf.horovod.get_ctx(config=config)
import socket
if "gpu" in config.value("device", "") or os.environ.get(
"CUDA_VISIBLE_DEVICES", ""):
# We assume that we want to use a GPU.
gpu_opts = config.typed_dict.setdefault("tf_session_opts",
{}).setdefault(
"gpu_options", {})
assert "visible_device_list" not in gpu_opts
gpu_opts["visible_device_list"] = str(hvd.local_rank())
print("Horovod: Hostname %s, pid %i, using GPU %s." %
(socket.gethostname(), os.getpid(),
gpu_opts["visible_device_list"]),
file=log.v3)
else:
if hvd.rank() == 0: # Don't spam in all ranks.
print("Horovod: Not using GPU.", file=log.v3)
if hvd.rank() == 0: # Don't spam in all ranks.
print("Horovod: Reduce type:",
hvd.get_reduce_type(),
file=log.v3)
from returnn.tf.util.basic import debug_register_better_repr, setup_tf_thread_pools, print_available_devices
tf_session_opts = config.typed_value("tf_session_opts", {})
assert isinstance(tf_session_opts, dict)
# This must be done after the Horovod logic, such that we only touch the devices we are supposed to touch.
setup_tf_thread_pools(log_file=log.v3, tf_session_opts=tf_session_opts)
# Print available devices. Also make sure that get_tf_list_local_devices uses the correct TF session opts.
print_available_devices(tf_session_opts=tf_session_opts, file=log.v2)
from returnn.tf.native_op import OpMaker
OpMaker.log_stream = log.v3
debug_register_better_repr()
if config.is_true("distributed_tf"):
import returnn.tf.distributed
returnn.tf.distributed.init_distributed_tf(config)
else:
raise NotImplementedError
def setup():
"""
Calls necessary setups.
"""
# Disable extensive TF debug verbosity. Must come before the first TF import.
import logging
logging.getLogger('tensorflow').disabled = True
# os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
# logging.getLogger("tensorflow").setLevel(logging.INFO)
# Get us some further useful debug messages (in some cases, e.g. CUDA).
# For example: https://github.com/tensorflow/tensorflow/issues/24496
# os.environ["CUDNN_LOGINFO_DBG"] = "1"
# os.environ["CUDNN_LOGDEST_DBG"] = "stdout"
# The following might fix (workaround): Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
# (https://github.com/tensorflow/tensorflow/issues/24496).
# os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
import _setup_returnn_env # noqa
import returnn.util.basic as util
util.init_thread_join_hack()
from returnn.util import better_exchook
better_exchook.install()
better_exchook.replace_traceback_format_tb()
from returnn.log import log
log.initialize(verbosity=[5])
# TF is optional.
# Note that importing TF still has a small side effect:
# BackendEngine._get_default_engine() will return TF by default, if TF is already loaded.
# For most tests, this does not matter.
try:
import tensorflow as tf
except ImportError:
tf = None
if tf:
import returnn.tf.util.basic as tf_util
tf_util.debug_register_better_repr()
import returnn.util.debug as debug
debug.install_lib_sig_segfault()
try:
import faulthandler
# Enable after libSigSegfault, so that we have both,
# because faulthandler will also call the original sig handler.
faulthandler.enable()
except ImportError:
print("no faulthandler")
def main():
parser = argparse.ArgumentParser(description="MB-MelGAN vocoder")
parser.add_argument("--features",
required=True,
help="npy file. via decoder.py --dump_features")
parser.add_argument("--pwg_config",
type=str,
help="ParallelWaveGAN config (.yaml)")
parser.add_argument("--pwg_checkpoint",
type=str,
help="ParallelWaveGAN checkpoint (.pkl)")
args = parser.parse_args()
better_exchook.install()
debug_register_better_repr()
log.initialize(verbosity=[5])
setup_tf_thread_pools()
print_available_devices()
def model_func(wrapped_import, inputs: torch.Tensor):
if typing.TYPE_CHECKING or not wrapped_import:
import torch
from parallel_wavegan import models as pwg_models
from parallel_wavegan import layers as pwg_layers
else:
torch = wrapped_import("torch")
wrapped_import("parallel_wavegan")
pwg_models = wrapped_import("parallel_wavegan.models")
pwg_layers = wrapped_import("parallel_wavegan.layers")
# Initialize PWG
pwg_config = yaml.load(open(args.pwg_config), Loader=yaml.Loader)
pyt_device = torch.device("cpu")
generator = pwg_models.MelGANGenerator(
**pwg_config['generator_params'])
generator.load_state_dict(
torch.load(args.pwg_checkpoint,
map_location="cpu")["model"]["generator"])
generator.remove_weight_norm()
pwg_model = generator.eval().to(pyt_device)
assert pwg_config["generator_params"].get(
"aux_context_window", 0) == 0 # not implemented otherwise
pwg_pqmf = pwg_layers.PQMF(
pwg_config["generator_params"]["out_channels"]).to(pyt_device)
with torch.no_grad():
return pwg_pqmf.synthesis(pwg_model(inputs))
feature_data = numpy.load(args.features)
print("Feature shape:", feature_data.shape)
import pytorch_to_returnn.log
pytorch_to_returnn.log.Verbosity = 6
from pytorch_to_returnn.converter import verify_torch_and_convert_to_returnn
verify_torch_and_convert_to_returnn(model_func,
inputs=feature_data[None, :, :])
# from pytorch_to_returnn.wrapped_import import wrapped_import_demo
# from pytorch_to_returnn import torch as torch_returnn
# model_func(wrapped_import_demo, inputs=torch_returnn.from_numpy(feature_data[None, :, :]))
audio_waveform = model_func(None,
inputs=torch.from_numpy(
feature_data[None, :, :]))
audio_waveform = audio_waveform.view(-1).cpu().numpy()
audio_raw = numpy.asarray(audio_waveform * (2**15 - 1),
dtype="int16").tobytes()
out_fn = "out.wav"
wave_writer = wave.open(out_fn, "wb")
wave_writer.setnchannels(1)
wave_writer.setframerate(16000)
wave_writer.setsampwidth(2)
wave_writer.writeframes(audio_raw)
wave_writer.close()
print("Wrote %s." % out_fn)