thutil.py

'''
Copyright (C) 2014 New York University

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
'''
import sys
import time
import numpy as np
import operator
import types
import ipdb
import inspect
import traceback

import theano
import theano.tensor as T

from theano import Op, Apply

from theano.tensor.shared_randomstreams import RandomStreams
from theano.tensor.nnet import conv
from theano.gof import local_optimizer

from common import imgutil, logutil

_log = logutil.getLogger()

use_gpu = theano.config.device.startswith('gpu')

checkgrad = False

if use_gpu:
    from theano.sandbox.cuda import GpuOp, gpu_from_host, host_from_gpu, \
                                    CudaNdarrayType, CudaNdarray
    from theano.sandbox.cuda.basic_ops import gpu_contiguous

class Eval(object):
    def __init__(self, globals=None, locals=None):
        self.globals = globals or {}
        self.locals = locals or sys._getframe(1).f_locals

    def __getitem__(self, key):
        return eval(key, self.globals, self.locals)

def c_contiguous(x):
    if x.is_c_contiguous():
        return x
    return x.copy()

def isvalid(x):
    return T.all(T.logical_not(T.logical_or(T.isnan(x), T.isinf(x))))

def maximum(x, y):
    if checkgrad:
        return x + y
    return T.maximum(x, y)

def minimum(x, y):
    if checkgrad:
        return x + y
    return T.minimum(x, y)

def named(x, name):
    x.name = name
    return x

def test_value(x):
    if isinstance(x, np.ndarray):
        return x
    return theano.gof.op.get_test_value(x)

def test_shape(x):
    return tuple(test_value(x.shape))

def theano_function(*vars_by_pos, **kwargs):
    '''theano function decorator'''
    mode = kwargs.pop('mode', 'FAST_RUN')
    check_valid = kwargs.pop('check_valid', False)
    checks = kwargs.pop('checks', ())
    vars_by_name = kwargs
    def compile_func(f):
        argnames = f.func_code.co_varnames[:f.func_code.co_argcount]
        if any([a in vars_by_name for a in argnames[:len(vars_by_pos)]]):
            raise ValueError('Argument supplied twice to %s' % f.func_name)
        varspec = dict(vars_by_name)
        varspec.update(zip(argnames[:len(vars_by_pos)], vars_by_pos))
        argvars = []
        for name in argnames:
            spec = varspec[name]
            if isinstance(spec, (tuple, list)):
                (var, test_val) = spec
            else:
                var = spec
                test_val = None
            assert isinstance(var, T.Variable)
            var.name = name
            if test_val is not None:
                var.tag.test_value = test_val
            argvars.append(var)
        return function(argvars, f(*argvars),
                        check_valid=check_valid,
                        checks=checks,
                        mode=mode)
    return compile_func

def function(inputs, outputs=None, check_valid=False, checks=(), **kwargs):
    input_names = None
    output_names = None
    if isinstance(inputs, dict):
        if inputs:
            (input_names, inputs) = zip(*inputs.iteritems())
        else:
            (input_names, inputs) = ((), ())
    if isinstance(outputs, dict):
        if outputs:
            (output_names, outputs) = zip(*outputs.iteritems())
        else:
            (output_names, outputs) = ((), ())

    if check_valid or checks:
        updates = kwargs.setdefault('updates', {})
        asserts = [assert_(c, 'check failed: %s' % c) for c in checks]

        if check_valid:
            if outputs:
                if not isinstance(outputs, (list, tuple)):
                    outputs = [outputs]
                asserts += (assert_(isvalid(x),
                                    'output invalid: %d (%s)' % (i, x.name))
                            for (i, x) in enumerate(outputs))

            if updates:
                asserts += (assert_(isvalid(xnew),
                                    'update invalid: variable %s' % str(x))
                               for (x, xnew) in updates.iteritems())

        checks_passed = theano.shared(np.int8(1), name='checks_passed')
        updates[checks_passed] = \
            T.all(T.as_tensor_variable(asserts)).astype('int8')

        f = _CheckedFunction(inputs, outputs, **kwargs)
    else:
        f = theano.function(inputs, outputs, **kwargs)
        if hasattr(f.fn, 'clear_storage'):
            f.clear_storage = f.fn.clear_storage
        else:
            _log.warn('Function %s has no clear_storage: disabling', f.fn)
            f.clear_storage = lambda: None

    if input_names is not None or output_names is not None:
        return NamedInputOutputFunction(input_names, output_names, f)
    return f

class NamedInputOutputFunction(object):
    def __init__(self, input_names, output_names, f):
        self.input_names = input_names
        self.output_names = output_names
        self.f = f

        if output_names:
            class _NamedOutputs(object):
                __slots__ = output_names

                def __init__(self, vals):
                    [setattr(self, k, v) for (k,v) in zip(self.__slots__, vals)]
                
                def __eq__(self, other):
                    return type(self) == type(other) and \
                           self.items() == other.items()

                def __getitem__(self, k):
                    return getattr(self, k)

                def iteritems(self):
                    return ((s, self[s]) for s in self.__slots__)

                __iter__ = iteritems

                def items(self):
                    return list(self.iteritems())

            self._NamedOutputs = _NamedOutputs

        if hasattr(f.fn, 'clear_storage'):
            self.clear_storage = f.fn.clear_storage
        else:
            _log.warn('Function %s has no clear_storage: disabling', f.fn)
            self.clear_storage = lambda: None

    def __call__(self, *args, **kwargs):
        inputs = args
        if self.input_names:
            assert not inputs, \
                   'theano function with kw args cannot take positional args'
            inputs = [kwargs[k] for k in self.input_names]

        outputs = self.f(*inputs)

        if self.output_names:
            outputs = self._NamedOutputs(outputs)

        return outputs

class _CheckedFunction(object):
    def __init__(self, inputs, outputs, **kwargs):
        self.f = theano.function(inputs, outputs,
                                 inplace_updates=False,
                                 **kwargs)
        self.dbg_kwargs = dict(kwargs)
        self.dbg_kwargs.update(inputs=inputs,
                               outputs=outputs,
                               inplace_updates=False,
                               mode='DEBUG_MODE')
        self.f_dbg = None
        self.fn = self.f.fn
        self.clear_storage = self.f.fn.clear_storage

    def __call__(self, *args, **kwargs):
        try:
            return self.f(*args, **kwargs)
        except AssertionError:
            _log.exception('assertion failed in function %s' % self.f.name)
            if self.f_dbg is None:
                _log.info('creating debug function for %s' % self.f.name)
                self.f_dbg = theano.function(**self.dbg_kwargs)
            _log.error('calling debug function for %s' % self.f.name)
            self.f_dbg(*args, **kwargs)
            _log.error('debug version seems to have passed' % self.f.name)
            raise

class Assert(theano.Op):
    view_map = {0: [0]}

    def __init__(self, msg=None):
        self.msg = msg

    def __eq__(self, other):
        return (type(self) == type(other) and
                self.msg == other.msg)

    def __hash__(self):
        return reduce(operator.xor, map(hash, (type(self), self.msg)))

    def make_node(self, input):
        output = T.as_tensor_variable(input).type()
        return theano.Apply(self, (input,), (output,))

    def make_gpu_node(self, input):
        return Assert_GPU(self.msg)(input)

    def infer_shape(self, node, input_shapes):
        return input_shapes

    def perform(self, node, (input,), (output,)):
        assert np.all(input), self.msg
        output[0] = input

    def grad(self, inputs, doutputs):
        return (None,)

def assert_(cond, msg=None):
    return Assert(msg)(cond)

class Constant(theano.Op):
    def __init__(self, ninputs):
        self.view_map = dict((i,[i]) for i in xrange(ninputs))

    def __eq__(self, other):
        return (type(self) == type(other) and
                len(self.view_map) == len(other.view_map))

    def __hash__(self):
        return reduce(operator.xor,
                      map(hash, (type(self), len(self.view_map))))

    def make_node(self, *inputs):
        outputs = tuple([T.as_tensor_variable(inp).type() for inp in inputs])
        return theano.Apply(self, inputs, outputs)

    def make_gpu_node(self, *inputs):
        return Constant_GPU(len(inputs))(*inputs)

    def infer_shape(self, node, input_shapes):
        return input_shapes

    def perform(self, node, inputs, outputs):
        for (inp, out) in zip(inputs, outputs):
            out[0] = inp

    def grad(self, inputs, doutputs):
        return [T.DisconnectedType()() for _ in inputs]

def constant(*inputs):
    return Constant(len(inputs))(*inputs)


class _BreakpointVars(object):
    def __init__(self, th_vars, py_vars):
        self.th_vars = th_vars
        self.py_vars = py_vars

    def __getattr__(self, k):
        if k in self.th_vars:
            return self.th_vars[k]
        if k in self.py_vars:
            return self.py_vars[k]
        return object.__getattr__(self, k)
    
    def __repr__(self):
        s = []
        s.append('Theano runtime variables:')
        s += ('%-16s  %s' % (k, str(v.shape))
              for (k, v) in sorted(self.th_vars.items(), key=lambda (k,v): k))
        s.append('')
        s.append('Python creation-time variables:')
        s.append(', '.join(sorted(self.py_vars.keys())))
        s.append('')
        return '\n'.join(s)

class Breakpoint(theano.Op):
    view_map = {0: [0]}

    global_breakpoint_enable = False

    def __init__(self, var_names, cond, tb, py_vars,
                       breakpoint_grad, is_grad=False):
        self.var_names = var_names
        self.cond = cond
        self.tb = tb
        self.py_vars = py_vars
        self.nvars = len(var_names)
        self.breakpoint_grad = breakpoint_grad
        self.is_grad = is_grad

    def __eq__(self, other):
        return (type(self) == type(other) and
                self.var_names == other.var_names and
                self.cond == other.cond and
                self.tb == other.tb)

    def __hash__(self):
        return reduce(operator.xor, map(hash, (
                   type(self), self.var_names, self.cond, self.tb)))

    def make_node(self, *inputs):
        output = T.as_tensor_variable(inputs[0]).type()
        return theano.Apply(self, inputs, (output,))

    def make_gpu_node(self, *inputs):
        return Breakpoint_GPU(
                    self.var_names, self.cond, self.tb, self.py_vars,
                    self.breakpoint_grad, self.is_grad)(*inputs)

    def infer_shape(self, node, input_shapes):
        return (input_shapes[0],)

    def perform(self, node, inputs, (output,)):
        output[0] = inputs[0]
        if not Breakpoint.global_breakpoint_enable:
            return
        x = inputs[0]
        if not isinstance(x, np.ndarray):
            x = np.array(x)
        if self.cond(x):
            vars = _BreakpointVars(
                        dict(zip(self.var_names, map(np.array, inputs[1:]))),
                        self.py_vars)
            if self.is_grad:
                place = 'theano gradient eval'
            else:
                place = 'theano eval'
            print >> sys.stderr, 'Breakpoint in %s, created at' % place
            print >> sys.stderr, '  ...'
            traceback.print_list(self.tb[-4:], sys.stderr)
            ipdb.set_trace()
            pass    # in theano breakpoint

    def grad(self, inputs, (doutput,)):
        if self.breakpoint_grad:
            doutput = Breakpoint(self.var_names, self.cond,
                                 self.tb, self.py_vars, True, True) \
                          (doutput, *inputs[1:])
        return [doutput] + [T.DisconnectedType()() for _ in xrange(self.nvars)]

_theano_types = (theano.tensor.basic.TensorConstant,
                 theano.tensor.basic.TensorVariable,
                 theano.compile.SharedVariable,
                 )

def is_theano_var(x):
    return isinstance(x, _theano_types)

def breakpoint(output, vars=None, cond=lambda v: True, grad=True):
    tb = tuple(traceback.extract_stack()[:-1])
    py_vars = {}
    if type(vars) not in (tuple, list, dict, types.NoneType):
        raise ValueError('vars keyword arg must be None, dict, list or tuple')
    if not isinstance(vars, dict):
        frame_locals = inspect.stack()[1][0].f_locals
        if vars is not None:
            frame_locals = dict((name, val)
                                for (name, val) in frame_locals.iteritems()
                                if name in vars or val in vars)
        vars = frame_locals
    assert isinstance(vars, dict)
    th_vars = dict((name, val) for (name, val) in vars.iteritems()
                               if isinstance(val, _theano_types))
    py_vars = dict((name, val) for (name, val) in vars.iteritems()
                               if name not in th_vars)
    (th_var_names, th_var_vals) = zip(*th_vars.iteritems())
    return Breakpoint(th_var_names, cond, tb, py_vars, grad) \
                     (output, *th_var_vals)

def enable_breakpoints(enable=True):
    Breakpoint.global_breakpoint_enable = enable

def cross(x, y, axis=None):
    ndim = x.ndim
    assert x.ndim == y.ndim
    if axis is None:
        axis = ndim - 1
    def _getindexslice(a, i):
        return a[tuple([slice(i,i+1) if d == axis else slice(None)
                        for d in xrange(ndim)])]
    x0 = _getindexslice(x, 0)
    x1 = _getindexslice(x, 1)
    x2 = _getindexslice(x, 2)
    y0 = _getindexslice(y, 0)
    y1 = _getindexslice(y, 1)
    y2 = _getindexslice(y, 2)

    res = T.concatenate((x1*y2 - x2*y1,
                         x2*y0 - x0*y2,
                         x0*y1 - x1*y0), axis=axis)
    return res


if use_gpu:

    class Constant_GPU(Constant, GpuOp):
        def make_node(self, *inputs):
            outputs = tuple([inp.type() for inp in inputs])
            return theano.Apply(self, inputs, outputs)

    class Assert_GPU(Assert, GpuOp):
        def make_node(self, input):
            output = input.type()
            return theano.Apply(self, (input,), (output,))

        def perform(self, node, (input,), (output,)):
            assert np.all(np.array(input))
            output[0] = input

    class Breakpoint_GPU(Breakpoint, GpuOp):
        def make_node(self, *inputs):
            output = inputs[0].type()
            return theano.Apply(self, inputs, (output,))

    @theano.sandbox.cuda.opt.register_opt()
    @theano.gof.local_optimizer(None)
    def local_gpu_togpu(node):
        if node.op == gpu_from_host:
            host_input = node.inputs[0]
            if host_input.owner and \
                    hasattr(host_input.owner.op, 'make_gpu_node'):
                try:
                    gpu_inputs = map(gpu_from_host, host_input.owner.inputs)
                except TypeError:
                    return False
                return [host_input.owner.op.make_gpu_node(*gpu_inputs)]
        elif hasattr(node.op, 'make_gpu_node') and \
                all([x.owner and x.owner.op == host_from_gpu
                     for x in node.inputs]):
            gpu_inputs = [x.owner.inputs[0] for x in node.inputs]
            return [host_from_gpu(node.op.make_gpu_node(*gpu_inputs))]
        return False

    @theano.sandbox.cuda.opt.register_opt()
    @theano.gof.local_optimizer([Breakpoint])
    def local_gpu_togpu_breakpoint(node):
        if isinstance(node.op, Breakpoint):
            result_input = node.inputs[0]
            if result_input.owner and result_input.owner.op == host_from_gpu:
                gpu_inputs = [x.owner.inputs[0]
                                if x.owner and x.owner.op == host_from_gpu
                                else x
                              for x in node.inputs]
                return [host_from_gpu(node.op.make_gpu_node(*gpu_inputs))]
        return False