DebugHelpers.py

"""
This file is going to be imported by Debug.debug_shell() and available as interactive commands.
"""

import sys
import numpy
import h5py
from Util import BackendEngine

if BackendEngine.is_theano_selected():
  # noinspection PyUnresolvedReferences,PyPackageRequirements
  import theano
  # noinspection PyUnresolvedReferences,PyPackageRequirements
  import theano.tensor as tt
  # noinspection PyUnresolvedReferences,PyPackageRequirements
  import theano.sandbox.cuda as cuda
  from TheanoUtil import make_var_tuple
  from Network import LayerNetwork
else:
  theano = None


def find_obj_in_stack(cls, stack=None, all_threads=True):
  """
  :param type cls:
  :param types.FrameType|traceback.FrameSummary stack:
  :param bool all_threads:
  :return: obj
  """
  if all_threads:
    assert stack is None
    # noinspection PyProtectedMember
    for tid, stack in sys._current_frames().items():
      obj = find_obj_in_stack(cls=cls, stack=stack, all_threads=False)
      if obj is not None:
        return obj
    return None

  assert not all_threads
  if stack is None:
    # noinspection PyProtectedMember
    stack = sys._getframe()
    assert stack, "could not get stack"

  import inspect
  isframe = inspect.isframe

  _tb = stack
  while _tb is not None:
    if isframe(_tb):
      f = _tb
    else:
      f = _tb.tb_frame

    for obj in f.f_locals.values():
      if isinstance(obj, cls):
        return obj

    if isframe(_tb):
      _tb = _tb.f_back
    else:
      _tb = _tb.tb_next

  return None


_device = None


def get_device():
  """
  :rtype: Device.Device
  """
  global _device
  if _device:
    return _device
  from Device import Device
  _device = find_obj_in_stack(Device)
  return _device


if BackendEngine.is_theano_selected():
  def compute(var, trainnet=True):
    """
    :param theano.Variable var: variable which we should compute the value of
    :param bool trainnet: whether to make givens based on dev.trainnet or dev.testnet
    :return: the computed value
    :rtype: numpy.ndarray

    This expects to calculate some value of the trainnet or testnet of the current Device.
    """
    dev = get_device()
    assert dev, "no Device instance found"
    if trainnet:
      network = dev.trainnet
    else:
      network = dev.testnet
    givens = dev.make_givens(network)
    if isinstance(var, list):
      outputs = var
    else:
      outputs = [var]
    func = theano.function(inputs=[dev.block_start, dev.block_end],
                           outputs=outputs,
                           givens=givens,
                           on_unused_input='warn',
                           name="debug compute")
    batch_dim = dev.y["data"].get_value(borrow=True, return_internal_type=True).shape[1]
    batch_start = 0
    batch_end = batch_dim
    result = func(batch_start, batch_end)
    if not isinstance(var, list):
      result = result[0]
    return result


  class DebugNn:
    """
    Debug helper for NN.
    """

    def __init__(self, filename):
      self.network = LayerNetwork.from_hdf(filename, mask="unity", train_flag=False, eval_flag=True)
      self.f_forwarder = None

    def compile_forwarder(self):
      """
      Compile function for forwarding.
      """
      network = self.network
      data_keys = list(sorted(network.j.keys()))
      # All input seqs expected to have same length.
      givens = [(network.j[k], tt.ones(network.y["data"].shape[:2], dtype="int8")) for k in data_keys]
      self.f_forwarder = theano.function(
        inputs=[network.y["data"]],
        outputs=(
          [network.output["output"].output] +
          [layer.output for name, layer in sorted(network.output.items()) if name != "output"]),
        givens=givens,
        on_unused_input='warn',
        name="forwarder")

    def forward(self, data, output_index=0):
      """
      :param numpy.ndarray data: shape (time,in-dim)
      :param int output_index:
      :return: shape (time,out-dim)
      :rtype: numpy.ndarray
      """
      assert data.ndim == 2
      data = data[:, None, :]  # add batch-dim
      assert self.f_forwarder
      res = self.f_forwarder(data)
      res = make_var_tuple(res)[output_index]
      assert res.ndim == 3
      assert res.shape[1] == 1
      res = res[:, 0]
      return res


class SimpleHdf:
  """
  Simple HDF writer.
  """

  def __init__(self, filename):
    self.hdf = h5py.File(filename)
    self.seq_tag_to_idx = {name: i for (i, name) in enumerate(self.hdf["seqTags"])}
    self.num_seqs = len(self.hdf["seqTags"])
    assert self.num_seqs == len(self.seq_tag_to_idx), "not unique seq tags"
    seq_lens = self.hdf["seqLengths"]
    if seq_lens.ndim == 2:
      seq_lens = seq_lens[:, 0]
    assert self.num_seqs == len(seq_lens)
    self.seq_starts = [0] + list(numpy.cumsum(seq_lens))
    total_len = self.seq_starts[-1]
    inputs_len = self.hdf["inputs"].shape[0]
    assert total_len == inputs_len, "time-dim does not match: %i vs %i" % (total_len, inputs_len)
    assert self.seq_starts[-1] == self.hdf["targets/data/classes"].shape[0]

  def get_seq_tags(self):
    """
    :rtype: list[str]
    """
    return self.hdf["seqTags"]

  def get_data(self, seq_idx):
    """
    :param int seq_idx:
    :rtype: numpy.ndarray
    """
    seq_t0, seq_t1 = self.seq_starts[seq_idx:seq_idx + 2]
    return self.hdf["inputs"][seq_t0:seq_t1]

  def get_targets(self, seq_idx):
    """
    :param int seq_idx:
    :rtype: numpy.ndarray
    """
    seq_t0, seq_t1 = self.seq_starts[seq_idx:seq_idx + 2]
    return self.hdf["targets/data/classes"][seq_t0:seq_t1]

  def get_data_dict(self, seq_idx):
    """
    :param int seq_idx:
    :rtype: dict[str,numpy.ndarray]
    """
    return {"data": self.get_data(seq_idx), "classes": self.get_targets(seq_idx)}