SprintInterface.py

"""
This is a Sprint interface implementation, i.e. you would specify this module in your Sprint config.
(Sprint = the RWTH ASR toolkit.)
Note that there are multiple Sprint interface implementations provided.
This one would be used explicitly, e.g. for forwarding in recognition
or wherever else Sprint needs posteriors (a FeatureScorer).
Most of the other Sprint interfaces will be used automatically,
e.g. via ExternSprintDataset, when it spawns its Sprint subprocess.
"""

from __future__ import print_function

import os
import sys
import time
from threading import Event, Thread
import typing
import numpy

from Dataset import Dataset, init_dataset
from SprintDataset import SprintDatasetBase
from EngineBase import EngineBase
from Log import log
from EngineUtil import assign_dev_data_single_seq
import Debug
from Util import get_gpu_names, interrupt_main, to_bool, BackendEngine
import TaskSystem
import rnn

if typing.TYPE_CHECKING:
  try:
    import TFEngine
  except ImportError:
    TFEngine = None
  try:
    import Engine as TheanoEngine
  except ImportError:
    TheanoEngine = None

_rnn_file = rnn.__file__
_main_file = getattr(sys.modules["__main__"], "__file__", "")
if _rnn_file.endswith(".pyc"):
  _rnn_file = _rnn_file[:-1]
if _main_file.endswith(".pyc"):
  _main_file = _main_file[:-1]
if os.path.realpath(_rnn_file) == os.path.realpath(_main_file):
  rnn = sys.modules["__main__"]

DefaultSprintCrnnConfig = "config/crnn.config"

startTime = None  # type: typing.Optional[float]
isInitialized = False
isTrainThreadStarted = False
isExited = False
InputDim = None  # type: typing.Optional[int]
OutputDim = None  # type: typing.Optional[int]
MaxSegmentLength = 1
TargetMode = None  # type: typing.Optional[str]
Task = "train"

Engine = None  # type: typing.Optional[typing.Union[typing.Type[TFEngine.Engine],typing.Type[TheanoEngine.Engine]]]
config = None  # type: typing.Optional[rnn.Config]
sprintDataset = None  # type: typing.Optional[SprintDatasetBase]
customDataset = None  # type: typing.Optional[Dataset]
engine = None  # type: typing.Optional[typing.Union[TFEngine.Engine,TheanoEngine.Engine,Engine]]


# <editor-fold desc="generic init">
# Generic interface, should be compatible to any PythonControl-based, and PythonTrainer. {

def init(name=None, sprint_unit=None, **kwargs):
  """
  This will get called by various Sprint interfaces.
  Depending on `name` and `sprint_unit`, we can figure out which interface it is.
  For all PythonControl-based interfaces, we must return an object which will be used for further callbacks.

  :param str|None name:
  :param str|None sprint_unit:
  :return: some object or None
  :rtype: None|object
  """
  print("CRNN Python SprintInterface init: name %r, sprint_unit %r, pid %i, kwargs %r" % (
    name, sprint_unit, os.getpid(), kwargs))
  if name is None:
    return init_python_trainer(**kwargs)
  elif name == "Sprint.PythonControl":
    # Any PythonControl interface.
    if sprint_unit == "PythonFeatureScorer":
      return init_python_feature_scorer(**kwargs)
    else:
      raise Exception(
        "SprintInterface: Did not expect init() PythonControl with sprint_unit=%r, kwargs=%r",
        (sprint_unit, kwargs))
  else:
    raise Exception(
      "SprintInterface: Did not expect init() with name=%r, sprint_unit=%r, kwargs=%r",
      (name, sprint_unit, kwargs))

# }
# </editor-fold>


# <editor-fold desc="PythonFeatureScorer">
# Start Sprint PythonFeatureScorer interface. {

# noinspection PyShadowingNames
def init_python_feature_scorer(config, **kwargs):
  """
  :param str config:
  :rtype: PythonFeatureScorer
  """
  sprint_opts = {key: value for (key, value) in [s.split(":", 1) for s in config.split(",") if s]}

  epoch = sprint_opts.get("epoch", None)
  if epoch is not None:
    epoch = int(epoch)
    assert epoch >= 1

  # see init_python_trainer()
  configfile = sprint_opts.get("configfile", None)
  assert sprint_opts.get("action", None) in (None, "forward"), "invalid action: %r" % sprint_opts["action"]

  _init_base(target_mode="forward", configfile=configfile, epoch=epoch, sprint_opts=sprint_opts)

  cls = PythonFeatureScorer
  if rnn.config.has("SprintInterfacePythonFeatureScorer"):
    cls = rnn.config.typed_value("SprintInterfacePythonFeatureScorer")
  return cls(sprint_opts=sprint_opts, **kwargs)


class PythonFeatureScorer(object):
  """
  Sprint API.
  """

  def __init__(self, callback, version_number, sprint_opts, **kwargs):
    """
    :param (str,)->object callback:
    :param int version_number:
    :param dict[str,str] sprint_opts:
    """
    print("SprintInterface: PythonFeatureScorer(%s): version %i, sprint_opts %r, other %r" % (
      self.__class__.__name__, version_number, sprint_opts, kwargs))
    self.input_dim = None
    self.output_dim = None
    self.callback = callback
    self.sprint_opts = sprint_opts
    self.priors = None  # type: typing.Optional[numpy.ndarray]
    self.segment_count = 0
    self.features = []  # type: typing.List[numpy.ndarray]
    self.scores = None  # type: typing.Optional[numpy.ndarray]

  def init(self, input_dim, output_dim):
    """
    Called by Sprint.

    :param int input_dim:
    :param int output_dim: number of emission classes
    """
    self.input_dim = input_dim
    self.output_dim = output_dim

    # see init_python_trainer()
    global InputDim, OutputDim
    InputDim = input_dim
    OutputDim = output_dim
    sprintDataset.set_dimensions(self.input_dim, self.output_dim)
    sprintDataset.initialize()

    _prepare_forwarding()
    self._load_priors()

    global startTime
    startTime = time.time()

  def _load_priors(self):
    """
    This will optionally initialize self.priors of shape (self.output_dim,), in -log space,
    already multiplied by any prior scale.

    :return: nothing
    """
    scale = float(self.sprint_opts["prior_scale"])
    if not scale:
      return
    filename = self.sprint_opts["prior_file"]
    # We expect a filename to the priors, stored as txt, in +log space.
    assert isinstance(filename, str)
    assert os.path.exists(filename)
    from Util import load_txt_vector
    prior = load_txt_vector(filename)  # +log space
    self.priors = -numpy.array(prior, dtype="float32") * numpy.float32(scale)  # -log space
    assert self.priors.shape == (self.output_dim,), "dim mismatch: %r != %i" % (self.priors.shape, self.output_dim)

  # noinspection PyMethodMayBeStatic
  def exit(self):
    """
    Called by Sprint at exit.
    """
    print("SprintInterface: PythonFeatureScorer: exit()")

  # noinspection PyMethodMayBeStatic
  def get_feature_buffer_size(self):
    """
    Called by Sprint.

    :return: -1 -> no limit
    """
    return -1

  # noinspection PyShadowingNames
  def add_feature(self, feature, time):
    """
    Called by Sprint.

    :param numpy.ndarray feature: shape (input_dim,)
    :param int time:
    """
    assert time == len(self.features)
    assert feature.shape == (self.input_dim,)
    self.features.append(feature)

  def reset(self, num_frames):
    """
    Called by Sprint.
    Called when we shall flush any buffers.

    :param int num_frames:
    """
    if num_frames > 0:
      self.segment_count += 1
    assert num_frames == len(self.features)
    del self.features[:]
    self.scores = None

  def get_segment_name(self):
    """
    :rtype: str
    """
    return "unknown-seq-name-%i" % self.segment_count

  def get_features(self, num_frames=None):
    """
    :param int|None num_frames:
    :return: shape (input_dim, num_frames)
    :rtype: numpy.ndarray
    """
    if num_frames is not None:
      assert 0 < num_frames == len(self.features)
    return numpy.stack(self.features, axis=1)

  def get_posteriors(self, num_frames=None):
    """
    :param int|None num_frames:
    :return: shape (output_dim, num_frames)
    :rtype: numpy.ndarray
    """
    if num_frames is None:
      num_frames = len(self.features)
    assert 0 < num_frames == len(self.features)
    posteriors = _forward(
      segment_name=self.get_segment_name(),
      features=self.get_features(num_frames=num_frames))
    assert posteriors.shape == (self.output_dim, num_frames)
    return posteriors

  def features_to_dataset(self, num_frames=None):
    """
    :param int|None num_frames:
    :return: (dataset, seq_idx)
    :rtype: (Dataset.Dataset, int)
    """
    segment_name = self.get_segment_name()
    features = self.get_features(num_frames=num_frames)
    return features_to_dataset(features=features, segment_name=segment_name)

  @property
  def engine(self):
    """
    :rtype: TFEngine.Engine|Engine.Engine
    """
    return rnn.engine

  @property
  def config(self):
    """
    :rtype: Config.Config
    """
    return rnn.config

  def compute(self, num_frames):
    """
    Called by Sprint.
    All the features which we received so far should be evaluated.

    :param int num_frames:
    """
    assert 0 < num_frames == len(self.features)
    posteriors = self.get_posteriors(num_frames=num_frames)
    assert posteriors.shape == (self.output_dim, num_frames)
    scores = -numpy.log(posteriors)  # transfer to -log space
    if self.priors is not None:
      scores -= numpy.expand_dims(self.priors, axis=1)
    # We must return in -log space.
    self.scores = scores

  # noinspection PyShadowingNames
  def get_scores(self, time):
    """
    Called by Sprint.

    :param int time:
    :return: shape (output_dim,)
    :rtype: numpy.ndarray
    """
    # print("get scores, time", time, "max_frames", self.scores.shape[1])
    return self.scores[:, time]

# }
# </editor-fold>


# <editor-fold desc="PythonSegmentOrder">
# Start Sprint PythonSegmentOrder interface. {

# Names need to stay that way for compatibility.
# noinspection PyPep8Naming
def getSegmentList(corpusName, segmentList, **kwargs):
  """
  Called by Sprint PythonSegmentOrder.
  Set python-segment-order = true in Sprint to use this.

  If this is used, this gets called really early.
  If it is used together with the Sprint PythonTrainer,
  it will get called way earlier before the init() below.
  It might also get called multiple times, e.g. if
  Sprint is in interactive mode to calc the seg count.
  This is optional. You can use the SprintInterface
  only for the PythonTrainer.

  :type corpusName: str
  :type segmentList: list[str]
  :rtype: list[str]
  :returns segment list. Can also be an iterator.
  """
  print("Sprint: getSegmentList(%r)" % corpusName)
  print("Corpus segments #: %i" % len(segmentList))
  print("(This can be further filtered in Sprint by a whitelist or so.)")

  # Init what we need. These can be called multiple times.
  # If we use both the PythonSegmentOrder and the PythonTrainer, this will be called first.
  # The PythonTrainer will be called lazily once it gets the first data.
  _init_base(configfile=kwargs.get('config', None))
  sprintDataset.use_multiple_epochs()

  finalEpoch = get_final_epoch()
  startEpoch, startSegmentIdx = Engine.get_train_start_epoch_batch(config)
  print("Sprint: Starting with epoch %i, segment-idx %s." % (startEpoch, startSegmentIdx))
  print("Final epoch is: %i" % finalEpoch)

  # Loop over multiple epochs. Epochs start at 1.
  for curEpoch in range(startEpoch, finalEpoch + 1):
    if isTrainThreadStarted:
      # So that the CRNN train thread always has the SprintDatasetBase in a sane state before we reset it.
      sprintDataset.wait_for_returnn_epoch(curEpoch)
    sprintDataset.init_sprint_epoch(curEpoch)

    index_list = sprintDataset.get_seq_order_for_epoch(curEpoch, len(segmentList))
    orderedSegmentList = [segmentList[i] for i in index_list]
    assert len(orderedSegmentList) == len(segmentList)

    print("Sprint epoch: %i" % curEpoch)
    startSegmentIdx = 0
    if curEpoch == startEpoch:
      startSegmentIdx = startSegmentIdx
    for curSegmentIdx in range(startSegmentIdx, len(orderedSegmentList)):
      sprintDataset.set_complete_frac(float(curSegmentIdx - startSegmentIdx + 1) /
                                      (len(orderedSegmentList) - startSegmentIdx))
      yield orderedSegmentList[curSegmentIdx]

    print("Sprint finished epoch %i" % curEpoch)
    sprintDataset.finish_sprint_epoch()

    if isTrainThreadStarted:
      assert sprintDataset.get_num_timesteps() > 0, \
        "We did not received any seqs. You are probably using a buffered feature extractor and the buffer is " + \
        "bigger than the total number of time frames in the corpus."

  sprintDataset.finalize_sprint()

# End Sprint PythonSegmentOrder interface. }
# </editor-fold>


# <editor-fold desc="PythonTrainer">
# Start Sprint PythonTrainer interface. {

# noinspection PyPep8Naming,PyShadowingNames
def init_python_trainer(inputDim, outputDim, config, targetMode, **kwargs):
  """
  Called by Sprint when it initializes the PythonTrainer.
  Set trainer = python-trainer in Sprint to enable.
  Note that Sprint will call this, i.e. the trainer init lazily quite late,
  only once it sees the first data.

  :type inputDim: int
  :type outputDim: int
  :param str config: config string, passed by Sprint. assumed to be ","-separated
  :param str targetMode: "target-alignment" or "criterion-by-sprint" or so
  :return: not expected to return anything
  :rtype: None
  """
  print("SprintInterface[pid %i] init()" % (os.getpid(),))
  print("inputDim:", inputDim)
  print("outputDim:", outputDim)
  print("config:", config)
  print("targetMode:", targetMode)
  print("other args:", kwargs)
  global InputDim, OutputDim, MaxSegmentLength
  InputDim = inputDim
  OutputDim = outputDim

  MaxSegmentLength = kwargs.get('maxSegmentLength', MaxSegmentLength)

  config = config.split(",")
  config = {key: value for (key, value) in [s.split(":", 1) for s in config if s]}

  if to_bool(config.get("EnableAutoNumpySharedMemPickling", False)) and not TaskSystem.SharedMemNumpyConfig["enabled"]:
    TaskSystem.SharedMemNumpyConfig["enabled"] = True
    print("SprintInterface[pid %i] EnableAutoNumpySharedMemPickling = True" % (os.getpid(),))

  epoch = config.get("epoch", None)
  if epoch is not None:
    epoch = int(epoch)
    assert epoch >= 1

  configfile = config.get("configfile", None)

  global Task
  action = config["action"]
  Task = action
  if action == "train":
    pass
  elif action == "forward":
    assert targetMode in ["criterion-by-sprint", "forward-only"]
    targetMode = "forward"
  elif action == "nop":
    targetMode = None
  else:
    assert False, "unknown action: %r" % action

  _init_base(target_mode=targetMode, configfile=configfile, epoch=epoch, sprint_opts=config)
  sprintDataset.set_dimensions(inputDim, outputDim)
  sprintDataset.initialize()

  if Task == "train":
    _start_train_thread(epoch)
  elif Task == "forward":
    _prepare_forwarding()

  global startTime
  startTime = time.time()


# noinspection PyShadowingBuiltins
def exit():
  """
  Called by Sprint at exit.
  """
  print("SprintInterface[pid %i] exit()" % (os.getpid(),))
  assert isInitialized
  global isExited
  if isExited:
    print("SprintInterface[pid %i] exit called multiple times" % (os.getpid(),))
    return
  isExited = True
  if isTrainThreadStarted:
    engine.stop_train_after_epoch_request = True
    sprintDataset.finish_sprint_epoch()  # In case this was not called yet. (No PythonSegmentOrdering.)
    sprintDataset.finalize_sprint()  # In case this was not called yet. (No PythonSegmentOrdering.)
    trainThread.join()
  rnn.finalize()
  if startTime:
    print("SprintInterface[pid %i]: elapsed total time: %f" % (os.getpid(), time.time() - startTime), file=log.v3)
  else:
    print("SprintInterface[pid %i]: finished (unknown start time)" % os.getpid(), file=log.v3)


# Keep name for compatibility.
# noinspection PyPep8Naming,PyUnusedLocal
def feedInput(features, weights=None, segmentName=None):
  """
  :param numpy.ndarray features:
  :param numpy.array|None weights:
  :param str|None segmentName:
  :return: posteriors
  :rtype: numpy.ndarray
  """
  assert features.shape[0] == InputDim
  if Task == "train":
    posteriors = _train(segmentName, features)
  elif Task == "forward":
    posteriors = _forward(segmentName, features)
  else:
    assert False, "invalid task: %r" % Task
  assert posteriors.shape == (OutputDim * MaxSegmentLength, features.shape[1])
  return posteriors


# Keep name for compatibility.
# noinspection PyPep8Naming
def finishDiscard():
  """
  Discard some segment.
  """
  print("finishDiscard()")
  raise NotImplementedError  # TODO ...


# Keep name for compatibility.
# noinspection PyPep8Naming
def finishError(error, errorSignal, naturalPairingType=None):
  """
  :param float|numpy.ndarray error:
  :param numpy.ndarray errorSignal:
  :param str|None naturalPairingType: must be "softmax"
  :return: nothing
  """
  make_criterion_class()
  assert naturalPairingType == "softmax"
  assert Task == "train"
  # reformat. see train()
  error = numpy.array([error], dtype="float32")
  errorSignal = errorSignal.transpose()
  errorSignal = errorSignal[:, numpy.newaxis, :]
  errorSignal = numpy.array(errorSignal, dtype="float32")
  assert errorSignal.shape == Criterion.posteriors.shape

  Criterion.error = error
  Criterion.errorSignal = errorSignal
  Criterion.gotErrorSignal.set()


# Keep name for compatibility.
# noinspection PyPep8Naming,PyUnusedLocal
def feedInputAndTarget(features, weights=None, segmentName=None,
                       orthography=None, alignment=None,
                       speaker_name=None, speaker_gender=None,
                       **kwargs):
  """
  :param numpy.ndarray features:
  :param numpy.ndarray|None weights:
  :param str|None segmentName:
  :param str|None orthography:
  :param numpy.ndarray|None alignment:
  :param str|None speaker_name:
  :param str|None speaker_gender:
  :return: nothing
  """
  assert features.shape[0] == InputDim
  targets = {}
  if alignment is not None:
    targets["classes"] = alignment
  if orthography is not None:
    targets["orth"] = orthography
  _train(segmentName, features, targets)


# Keep name for compatibility.
# noinspection PyPep8Naming,PyUnusedLocal
def feedInputAndTargetAlignment(features, targetAlignment, weights=None, segmentName=None):
  """
  :param numpy.ndarray features:
  :param numpy.ndarray targetAlignment:
  :param numpy.ndarray|None weights:
  :param str|None segmentName:
  :return: nothing
  """
  assert features.shape[0] == InputDim
  assert Task == "train"
  _train(segmentName, features, targetAlignment)


# Keep name for compatibility.
# noinspection PyPep8Naming,PyUnusedLocal
def feedInputAndTargetSegmentOrth(features, targetSegmentOrth, weights=None, segmentName=None):
  """
  :param numpy.ndarray features:
  :param str targetSegmentOrth:
  :param numpy.ndarray|None weights:
  :param segmentName:
  :return:
  """
  assert features.shape[0] == InputDim
  assert Task == "train"
  _train(segmentName, features, {"orth": targetSegmentOrth})


# Keep name for compatibility.
# noinspection PyPep8Naming,PyUnusedLocal
def feedInputUnsupervised(features, weights=None, segmentName=None):
  """
  :param numpy.ndarray features:
  :param numpy.ndarray|None weights:
  :param str|None segmentName:
  :return: nothing
  """
  assert features.shape[0] == InputDim
  _train(segmentName, features)


# Keep name for compatibility.
# noinspection PyPep8Naming
def feedInputForwarding(features, weights=None, segmentName=None):
  """
  :param numpy.ndarray features:
  :param numpy.ndarray|None weights:
  :param str|None segmentName:
  :rtype: numpy.ndarray
  """
  assert Task == "forward"
  return feedInput(features, weights=weights, segmentName=segmentName)

# End Sprint PythonTrainer interface. }
# </editor-fold>


def dump_flags():
  """
  Dump some relevant env flags.
  """
  print("CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))
  print("CUDA_LAUNCH_BLOCKING:", os.environ.get("CUDA_LAUNCH_BLOCKING"))

  if BackendEngine.is_theano_selected():
    print("available GPUs:", get_gpu_names())

    # noinspection PyUnresolvedReferences,PyPackageRequirements
    from theano.sandbox import cuda as theano_cuda
    print("CUDA via", theano_cuda.__file__)
    print("CUDA available:", theano_cuda.cuda_available)

    from Util import TheanoFlags
    print("THEANO_FLAGS:", TheanoFlags)


def set_target_mode(mode):
  """
  :param str mode: target mode
  """
  global TargetMode
  assert config, "not initialized"
  TargetMode = mode
  task = "train"
  loss = config.value('loss', None)
  if TargetMode == "criterion-by-sprint":
    assert loss == "sprint", "TargetMode is %s but loss is %s" % (TargetMode, loss)
  elif TargetMode == "target-alignment":
    # RETURNN always expects an alignment, so this is good just as-is.
    # This means that we will not calculate the criterion in Sprint.
    assert loss != "sprint", "invalid loss %s for target mode %s" % (loss, TargetMode)
  elif TargetMode == "forward":
    # Will be handled below.
    task = "forward"
    config.set("extract", ["posteriors"])
  else:
    assert False, "target-mode %s not supported yet..." % TargetMode

  if engine:
    # If we already initialized the engine, the value must not differ,
    # because e.g. Devices will init accordingly.
    orig_task = config.value("task", "train")
    assert orig_task == task

  config.set("task", task)


def _at_exit_handler():
  if not isExited:
    print("SprintInterface[pid %i] atexit handler, exit() was not called, calling it now" % (os.getpid(),))
    exit()
    print("All threads:")
    import Debug
    Debug.dump_all_thread_tracebacks(exclude_self=True)


def _init_base(configfile=None, target_mode=None, epoch=None, sprint_opts=None):
  """
  :param str|None configfile: filename, via init(), this is set
  :param str|None target_mode: "forward" or so. via init(), this is set
  :param int epoch: via init(), this is set
  :param dict[str,str]|None sprint_opts: optional parameters to override values in configfile
  """

  global isInitialized
  isInitialized = True
  # Run through in any case. Maybe just to set targetMode.

  if not getattr(sys, "argv", None):
     # Set some dummy. Some code might want this (e.g. TensorFlow).
     sys.argv = [__file__]

  global config
  if not config:
    # Some subset of what we do in rnn.init().

    rnn.init_better_exchook()
    rnn.init_thread_join_hack()

    if configfile is None:
      configfile = DefaultSprintCrnnConfig
    assert os.path.exists(configfile)
    rnn.init_config(config_filename=configfile, extra_updates={"task": target_mode})
    assert rnn.config
    config = rnn.config
    if sprint_opts is not None:
      config.update(sprint_opts)

    rnn.init_log()
    rnn.returnn_greeting(config_filename=configfile)
    rnn.init_backend_engine()
    rnn.init_faulthandler(sigusr1_chain=True)
    rnn.init_config_json_network()

    global Engine
    if BackendEngine.is_tensorflow_selected():
      # Use TFEngine.Engine class instead of Engine.Engine.
      from TFEngine import Engine
    elif BackendEngine.is_theano_selected():
      from Engine import Engine

    import atexit
    atexit.register(_at_exit_handler)

  if target_mode:
    set_target_mode(target_mode)

  _init_dataset()

  if target_mode and target_mode == "forward" and epoch:
    model_filename = config.value('model', '')
    fns = [EngineBase.epoch_model_filename(model_filename, epoch, is_pretrain) for is_pretrain in [False, True]]
    fn_postfix = ""
    if BackendEngine.is_tensorflow_selected():
      fn_postfix += ".meta"
    fns_existing = [fn for fn in fns if os.path.exists(fn + fn_postfix)]
    assert len(fns_existing) == 1, "%s not found" % fns
    model_epoch_filename = fns_existing[0]
    config.set('load', model_epoch_filename)
    assert EngineBase.get_epoch_model(config)[1] == model_epoch_filename, (
      "%r != %r" % (EngineBase.get_epoch_model(config), model_epoch_filename))

  global engine
  if not engine:
    devices = rnn.init_theano_devices()
    rnn.print_task_properties(devices)
    rnn.init_engine(devices)
    engine = rnn.engine
    assert isinstance(engine, Engine)


def _start_train_thread(epoch=None):
  global config, engine, isInitialized, isTrainThreadStarted
  assert isInitialized, "need to call init() first"
  assert not isTrainThreadStarted
  assert sprintDataset, "need to call initDataset() first"
  assert Task == "train"

  def train_thread_func():
    """
    Train thread.
    """
    # noinspection PyBroadException
    try:
      assert TargetMode
      if TargetMode == "target-alignment":
        pass  # Ok.
      elif TargetMode == "criterion-by-sprint":
        # TODO ...
        raise NotImplementedError
      else:
        raise Exception("target-mode not supported: %s" % TargetMode)

      engine.init_train_from_config(config, train_data=sprintDataset)

      # If some epoch is explicitly specified, it checks whether it matches.
      if epoch is not None:
        assert epoch == engine.start_epoch

      # Do the actual training.
      engine.train()

    except KeyboardInterrupt:  # This happens at forced exit.
      pass

    except BaseException:  # Catch all, even SystemExit. We must stop the main thread then.
      try:
        print("CRNN train failed")
        sys.excepthook(*sys.exc_info())
      finally:
        # Exceptions are fatal. Stop now.
        interrupt_main()

  global trainThread
  trainThread = Thread(target=train_thread_func, name="Sprint CRNN train thread")
  trainThread.daemon = True  # However, at clean exit(), will will join this thread.
  trainThread.start()

  isTrainThreadStarted = True


def _prepare_forwarding():
  assert engine
  assert config
  # Should already be set via setTargetMode().
  assert config.list('extract') == ["posteriors"], "You need to have extract = posteriors in your CRNN config. " + \
                                                   "You have: %s" % config.list('extract')

  # Load network.
  engine.init_network_from_config(config)

  # Copy over net params.
  if BackendEngine.is_theano_selected():
    engine.devices[0].prepare(engine.network)


def _init_dataset():
  global sprintDataset, customDataset
  if sprintDataset:
    return
  assert config
  extra_opts = config.typed_value("sprint_interface_dataset_opts", {})
  assert isinstance(extra_opts, dict)
  sprintDataset = SprintDatasetBase.from_config(config, **extra_opts)
  if config.is_true("sprint_interface_custom_dataset"):
    custom_dataset_func = config.typed_value("sprint_interface_custom_dataset")
    assert callable(custom_dataset_func)
    custom_dataset_opts = custom_dataset_func(sprint_dataset=sprintDataset)
    customDataset = init_dataset(custom_dataset_opts)


def get_final_epoch():
  """
  :rtype: int
  """
  global config, engine
  assert engine
  assert config
  config_num_epochs = engine.config_get_final_epoch(config)
  if engine.is_training:
    assert engine.final_epoch == config_num_epochs
  return config_num_epochs


def _train(segment_name, features, targets=None):
  """
  :param str|None segment_name: full name
  :param numpy.ndarray features: 2d array
  :param numpy.ndarray|dict[str,numpy.ndarray|str]|None targets: 2d or 1d array
  """
  assert engine is not None, "not initialized. call initBase()"
  assert sprintDataset

  if sprintDataset.sprintFinalized:
    return
  sprintDataset.add_new_data(features, targets, segment_name=segment_name)

  # The CRNN train thread started via start() will do the actual training.

  if TargetMode == "criterion-by-sprint":

    # TODO...
    make_criterion_class()

    Criterion.gotPosteriors.clear()

    Criterion.gotPosteriors.wait()
    posteriors = Criterion.posteriors
    assert posteriors is not None

    # posteriors is in format (time,batch,emission)
    # assert posteriors.shape[0] == T
    assert posteriors.shape[1] == 1
    assert OutputDim * MaxSegmentLength == posteriors.shape[2]
    assert len(posteriors.shape) == 3
    # reformat to Sprint expected format (emission,time)
    posteriors = posteriors[:, 0, :]
    posteriors = posteriors.transpose()
    assert posteriors.shape[0] == OutputDim * MaxSegmentLength
    # assert posteriors.shape[1] == T
    assert len(posteriors.shape) == 2

    return posteriors


def features_to_dataset(features, segment_name):
  """
  :param numpy.ndarray features: format (input-feature,time) (via Sprint)
  :param str segment_name:
  :return: (dataset, seq-idx)
  :rtype: (Dataset.Dataset, int)
  """
  assert sprintDataset

  # Features are in Sprint format (feature,time).
  num_time = features.shape[1]
  assert features.shape == (InputDim, num_time)

  if customDataset:
    customDataset.finish_epoch()  # reset
    customDataset.init_seq_order(epoch=1, seq_list=[segment_name])

  # Fill the data for the current segment.
  sprintDataset.shuffle_frames_of_nseqs = 0  # We must not shuffle.
  sprintDataset.init_sprint_epoch(None)  # Reset cache. We don't need old seqs anymore.
  sprintDataset.init_seq_order(epoch=1, seq_list=[segment_name])
  seq = sprintDataset.add_new_data(features, segment_name=segment_name)

  if not customDataset:
    return sprintDataset, seq

  assert seq == 0
  return customDataset, 0


def _forward(segment_name, features):
  """
  :param numpy.ndarray features: format (input-feature,time) (via Sprint)
  :return: format (output-dim,time)
  :rtype: numpy.ndarray
  """
  print("Sprint forward", segment_name, features.shape)
  start_time = time.time()
  assert engine is not None, "not initialized"
  assert sprintDataset

  # Features are in Sprint format (feature,time).
  num_time = features.shape[1]
  assert features.shape == (InputDim, num_time)
  dataset, seq_idx = features_to_dataset(features=features, segment_name=segment_name)

  if BackendEngine.is_theano_selected():
    # Prepare data for device.
    device = engine.devices[0]
    success = assign_dev_data_single_seq(device, dataset=dataset, seq=seq_idx)
    assert success, "failed to allocate & assign data for seq %i, %s" % (seq_idx, segment_name)

    # Do the actual forwarding and collect result.
    device.run("extract")
    result, _ = device.result()
    assert result is not None, "Device crashed."
    assert len(result) == 1
    posteriors = result[0]

  elif BackendEngine.is_tensorflow_selected():
    posteriors = engine.forward_single(dataset=dataset, seq_idx=seq_idx)

  else:
    raise NotImplementedError("unknown backend engine")
  # If we have a sequence training criterion, posteriors might be in format (time,seq|batch,emission).
  if posteriors.ndim == 3:
    assert posteriors.shape == (num_time, 1, OutputDim * MaxSegmentLength)
    posteriors = posteriors[:, 0]
  # Posteriors are in format (time,emission).
  assert posteriors.shape == (num_time, OutputDim * MaxSegmentLength)
  # Reformat to Sprint expected format (emission,time).
  posteriors = posteriors.transpose()
  assert posteriors.shape == (OutputDim * MaxSegmentLength, num_time)
  stats = (numpy.min(posteriors), numpy.max(posteriors), numpy.mean(posteriors), numpy.std(posteriors))
  print("posteriors min/max/mean/std:", stats, "time:", time.time() - start_time)
  if numpy.isinf(posteriors).any() or numpy.isnan(posteriors).any():
    print("posteriors:", posteriors)
    debug_feat_fn = "/tmp/crnn.pid%i.sprintinterface.debug.features.txt" % os.getpid()
    debug_post_fn = "/tmp/crnn.pid%i.sprintinterface.debug.posteriors.txt" % os.getpid()
    print("Wrote to files %s, %s" % (debug_feat_fn, debug_post_fn))
    numpy.savetxt(debug_feat_fn, features)
    numpy.savetxt(debug_post_fn, posteriors)
    assert False, "Error, posteriors contain invalid numbers."

  return posteriors


Criterion = None


def make_criterion_class():
  """
  Make Criterion Theano class.
  """
  global Criterion
  if Criterion:
    return

  # noinspection PyUnresolvedReferences,PyPackageRequirements
  import theano
  # noinspection PyUnresolvedReferences,PyPackageRequirements,PyPep8Naming
  import theano.tensor as T

  # noinspection PyAbstractClass
  class Criterion(theano.Op):
    """
    Criterion Theano Op.
    """
    gotPosteriors = Event()
    gotErrorSignal = Event()
    posteriors = None
    error = None
    errorSignal = None

    def __eq__(self, other):
      return type(self) == type(other)

    def __hash__(self):
      return hash(type(self))

    def __str__(self):
      return self.__class__.__name__

    def make_node(self, posteriors, seq_lengths):
      """
      :param numpy.ndarray posteriors:
      :param numpy.ndarray seq_lengths:
      :return:
      """
      # We get the posteriors here from the Network output function,
      # which should be softmax.
      posteriors = theano.tensor.as_tensor_variable(posteriors)
      seq_lengths = theano.tensor.as_tensor_variable(seq_lengths)
      assert seq_lengths.ndim == 1  # vector of seqs lengths
      return theano.Apply(op=self, inputs=[posteriors, seq_lengths], outputs=[T.fvector(), posteriors.type()])

    # noinspection PyUnusedLocal
    def perform(self, node, inputs, output_storage, params=None):
      """
      :param node:
      :param inputs:
      :param output_storage:
      :param params:
      :return:
      """
      posteriors, seq_lengths = inputs
      num_time_frames = posteriors.shape[0]
      seq_lengths = numpy.array([num_time_frames])  # TODO: fix or so?

      self.__class__.posteriors = posteriors
      self.gotPosteriors.set()

      if numpy.isnan(posteriors).any():
        print('posteriors contain NaN!', file=log.v1)
      if numpy.isinf(posteriors).any():
        print('posteriors contain Inf!', file=log.v1)
        numpy.set_printoptions(threshold=numpy.nan)
        print('posteriors:', posteriors, file=log.v1)

      self.gotErrorSignal.wait()
      loss, errsig = self.error, self.errorSignal
      assert errsig.shape[0] == num_time_frames

      output_storage[0][0] = loss
      output_storage[1][0] = errsig

      print('avg frame loss for segments:', loss.sum() / seq_lengths.sum(), end=" ", file=log.v5)
      print('time-frames:', seq_lengths.sum(), file=log.v5)


def demo():
  """
  Demo for SprintInterface.
  """
  print("Note: Load this module via Sprint python-trainer to really use it.")
  print("We are running a demo now.")
  init(inputDim=493, outputDim=4501, config="",  # hardcoded, just a demo...
       targetMode="criterion-by-sprint", cudaEnabled=False, cudaActiveGpu=-1)
  assert os.path.exists("input-features.npy"), "run Sprint with python-trainer=dump first"
  features = numpy.load("input-features.npy")  # dumped via dump.py
  posteriors = feedInput(features)
  if not os.path.exists("posteriors.npy"):
    numpy.save("posteriors.npy", posteriors)
    print("Saved posteriors.npy. Now run Sprint with python-trainer=dump again.")
    sys.exit()
  old_posteriors = numpy.load("posteriors.npy")
  assert numpy.array_equal(posteriors, old_posteriors)
  error = numpy.load("output-error.npy")  # dumped via dump.py
  error = float(error)
  error_signal = numpy.load("output-error-signal.npy")  # dumped via dump.py
  finishError(error=error, errorSignal=error_signal, naturalPairingType="softmax")
  exit()


if __name__ == "__main__":
  Debug.debug_shell(user_ns=locals(), user_global_ns=globals())