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
import sys
import _setup_test_env # noqa
import unittest
from nose.tools import assert_equal, assert_is_instance, assert_in, assert_true, assert_false
from returnn.network_description import LayerNetworkDescription
from returnn.config import Config
from returnn.util.basic import dict_diff_str
from pprint import pprint
from returnn.util import better_exchook
from returnn.util.basic import BackendEngine
try:
# noinspection PyPackageRequirements
import theano
BackendEngine.select_engine(engine=BackendEngine.Theano)
except ImportError:
theano = None
if theano:
import returnn.theano.util
returnn.theano.util.monkey_patches()
from returnn.theano.network import LayerNetwork
from returnn.theano.layers.hidden import ForwardLayer
else:
LayerNetwork = None
ForwardLayer = None
