Skip to content

P
pytensor
提交 00189784 authored 作者: khaotik's avatar khaotik
浏览文件
操作

unify cumsum/cumprod in all places

上级 7cff935e
全部展开 隐藏空白字符变更
内嵌 并排
正在显示 包含 69 行增加 和 196 行删除
theano/gpuarray/extra_ops.py
浏览文件 @ 00189784
from __future__ import absolute_import, print_function, division
import os
from theano import Apply, Op
from theano.tensor.extra_ops import CumsumOp, CumprodOp
from theano.tensor.extra_ops import CumOp
from .basic_ops import infer_context_name
try:
from pygpu import gpuarray
......@@ -23,7 +23,7 @@ class GpuCumOp(GpuKernelBase, Op):
__props__ = ('axis', 'mode')
def __init__(self, axis, mode='add'):
self.axis = axis
self.axis = axis if axis else 0
self.mode = mode
def __eq__(self, other):
......@@ -463,17 +463,11 @@ class GpuCumOp(GpuKernelBase, Op):
@register_opt('fast_compile')
@op_lifter([CumsumOp, CumprodOp])
@register_opt2([CumsumOp, CumprodOp], 'fast_compile')
@op_lifter([CumOp])
@register_opt2([CumOp], 'fast_compile')
def local_gpua_cumop(op, ctx_name, inputs, outputs):
if inputs[0].dtype != 'float32':
return False
if isinstance(op, CumsumOp):
mode = 'add'
elif isinstance(op, CumprodOp):
mode = 'mul'
else:
return False
axis = op.axis
x = inputs[0]
if axis is not None and x.ndim > GpuCumOp.SUPPORTED_NDIMS:
......@@ -488,4 +482,4 @@ def local_gpua_cumop(op, ctx_name, inputs, outputs):
if axis is None:
axis = 0
return GpuCumOp(axis, mode)(x)
return GpuCumOp(axis, op.mode)(x)
theano/sandbox/cuda/extra_ops.py
浏览文件 @ 00189784
......@@ -5,7 +5,7 @@ from theano import Op
from theano.gof import local_optimizer
from theano.sandbox.cuda import cuda_available, GpuOp
from theano.sandbox.cuda.basic_ops import gpu_flatten
from theano.tensor.extra_ops import CumsumOp
from theano.tensor.extra_ops import CumOp
if cuda_available:
from theano.sandbox.cuda import CudaNdarrayType
......@@ -13,7 +13,7 @@ if cuda_available:
from theano.sandbox.cuda import register_opt as register_gpu_opt
class GpuCumsum(CumsumOp, GpuOp):
class GpuCumsum(CumOp, GpuOp):
"""
Parameters
......@@ -438,13 +438,16 @@ def values_eq_approx_high_tol(a, b):
@register_gpu_opt()
@local_optimizer([CumsumOp])
@local_optimizer([CumOp])
def use_gpu_cumsum(node):
if type(node.op) is CumsumOp \
if type(node.op) is CumOp \
and node.inputs[0].dtype == 'float32' \
and node.inputs[0].owner \
and isinstance(node.inputs[0].owner.op, HostFromGpu):
if node.op.mode != 'add':
return None
axis = node.op.axis
x = node.inputs[0]
......
theano/tensor/extra_ops.py
浏览文件 @ 00189784
......@@ -242,13 +242,14 @@ def searchsorted(x, v, side='left', sorter=None):
return SearchsortedOp(side=side)(x, v, sorter)
class CumsumOp(theano.Op):
# See function cumsum for docstring
class CumOp(theano.Op):
# See function cumsum/cumprod for docstring
__props__ = ("axis",)
__props__ = ("axis", "mode")
def __init__(self, axis=None):
def __init__(self, axis=None, mode='add'):
self.axis = axis
self.mode = mode
def make_node(self, x):
x = basic.as_tensor_variable(x)
......@@ -264,20 +265,37 @@ class CumsumOp(theano.Op):
def perform(self, node, inputs, output_storage):
x = inputs[0]
z = output_storage[0]
z[0] = np.cumsum(x, axis=self.axis)
z[0] = {'add': np.cumsum, 'mul': np.cumprod}[self.mode](x, axis=self.axis)
def grad(self, inputs, output_gradients):
[gi] = output_gradients
x, = inputs
gi, = output_gradients
if self.mode == 'mul':
fx = cumprod(x, axis=self.axis)
noimpl = NotImplementedError(
'CumOp: unknown gradient for mode %s' % self.mode)
if self.axis is None:
return [cumsum(gi[::-1])[::-1].reshape(inputs[0].shape)]
if self.mode == 'add':
return [cumsum(gi[::-1])[::-1].reshape(x.shape)]
elif self.mode == 'mul':
return [cumsum(
(fx * gi)[::-1])[::-1].reshape(x.shape) / x]
else:
raise noimpl
# We need to reverse the gradients along ``self.axis``,
# compute cumsum, then reverse again
reverse_slicing = [slice(None, None, None)] * gi.ndim
reverse_slicing[self.axis] = slice(None, None, -1)
reverse_slicing = tuple(reverse_slicing)
return [cumsum(gi[reverse_slicing], self.axis)[reverse_slicing]]
if self.mode == 'add':
return [cumsum(gi[reverse_slicing], self.axis)[reverse_slicing]]
elif self.mode == 'mul':
# We need to reverse the gradients along ``self.axis``,
# compute cumsum, then reverse again
return [cumsum(
(fx * gi)[reverse_slicing], self.axis)[reverse_slicing] / x]
else:
raise noimpl
def infer_shape(self, node, shapes):
if self.axis is None:
......@@ -290,6 +308,7 @@ class CumsumOp(theano.Op):
z, = onames
axis = self.axis
fail = sub['fail']
func = dict(mul='CumProd', add='CumSum')[self.mode]
if self.axis is None or (self.axis == 0 and node.inputs[0].ndim == 1):
code = """
......@@ -303,13 +322,13 @@ class CumsumOp(theano.Op):
if (!%(z)s)
%(fail)s;
{
PyObject * t = PyArray_CumSum(
PyObject * t = PyArray_%(func)s(
%(x)s, NPY_MAXDIMS,
PyArray_TYPE((PyArrayObject*) py_%(x)s), %(z)s);
if (!t){
%(fail)s;
}
// Because PyArray_CumSum returns a newly created reference on t.
// Because PyArray_%(func)s returns a newly created reference on t.
Py_XDECREF(t);
}
""" % locals()
......@@ -325,13 +344,13 @@ class CumsumOp(theano.Op):
%(fail)s;
{
PyObject * t = PyArray_CumSum(
PyObject * t = PyArray_%(func)s(
%(x)s, %(axis)s,
PyArray_TYPE((PyArrayObject*) py_%(x)s), %(z)s);
if (!t){
%(fail)s;
}
// Because PyArray_CumSum returns a newly created reference on t.
// Because PyArray_%(func)s returns a newly created reference on t.
Py_XDECREF(t);
}
""" % locals()
......@@ -339,10 +358,10 @@ class CumsumOp(theano.Op):
return code
def c_code_cache_version(self):
return (6,)
return (7,)
def __str__(self):
return "%s{%s}" % (self.__class__.__name__, self.axis)
return "%s{%s, %s}" % (self.__class__.__name__, self.axis, self.mode)
def cumsum(x, axis=None):
......@@ -361,112 +380,7 @@ def cumsum(x, axis=None):
.. versionadded:: 0.7
"""
return CumsumOp(axis=axis)(x)
class CumprodOp(theano.Op):
# See function cumprod for docstring
__props__ = ("axis",)
def __init__(self, axis=None):
self.axis = axis
def make_node(self, x):
x = basic.as_tensor_variable(x)
out_type = x.type()
if self.axis is None:
out_type = theano.tensor.vector(dtype=x.dtype) # Flatten
elif self.axis >= x.ndim or self.axis < -x.ndim:
raise ValueError('axis(={0}) out of bounds'.format(self.axis))
return theano.Apply(self, [x], [out_type])
def perform(self, node, inputs, output_storage):
x = inputs[0]
z = output_storage[0]
z[0] = np.cumprod(x, axis=self.axis)
def grad(self, inputs, output_gradients):
x, = inputs
gi, = output_gradients
fx = cumprod(x, axis=self.axis)
if self.axis is None:
return [cumsum((fx * gi)[::-1])[::-1].reshape(inputs[0].shape) / x]
# We need to reverse the gradients along ``self.axis``,
# compute cumsum, then reverse again
reverse_slicing = [slice(None, None, None)] * gi.ndim
reverse_slicing[self.axis] = slice(None, None, -1)
reverse_slicing = tuple(reverse_slicing)
return [cumsum((fx * gi)[reverse_slicing],
self.axis)[reverse_slicing] / x]
def infer_shape(self, node, shapes):
if self.axis is None:
return [(tensor.prod(shapes[0]),)] # Flatten
return shapes
def c_code(self, node, name, inames, onames, sub):
x, = inames
z, = onames
axis = self.axis
fail = sub['fail']
if self.axis is None or (self.axis == 0 and node.inputs[0].ndim == 1):
code = """
npy_intp shape[1] = { PyArray_SIZE(%(x)s) };
if(!(%(z)s && PyArray_DIMS(%(z)s)[0] == shape[0]))
{
Py_XDECREF(%(z)s);
%(z)s = (PyArrayObject*) PyArray_SimpleNew(1, shape, PyArray_TYPE((PyArrayObject*) py_%(x)s));
}
if (!%(z)s)
%(fail)s;
{
PyObject * t = PyArray_CumProd(
%(x)s, NPY_MAXDIMS,
PyArray_TYPE((PyArrayObject*) py_%(x)s), %(z)s);
if (!t){
%(fail)s;
}
// Because PyArray_CumSum returns a newly created reference on t.
Py_XDECREF(t);
}
""" % locals()
else:
code = """
if(!(%(z)s && PyArray_CompareLists(PyArray_DIMS(%(z)s), PyArray_DIMS(%(x)s), PyArray_NDIM(%(x)s)) ))
{
Py_XDECREF(%(z)s);
%(z)s = (PyArrayObject*) PyArray_SimpleNew(PyArray_NDIM(%(x)s), PyArray_DIMS(%(x)s), PyArray_TYPE((PyArrayObject*) py_%(x)s));
}
if (!%(z)s)
%(fail)s;
{
PyObject * t = PyArray_CumProd(
%(x)s, %(axis)s,
PyArray_TYPE((PyArrayObject*) py_%(x)s), %(z)s);
if (!t){
%(fail)s;
}
// Because PyArray_CumSum returns a newly created reference on t.
Py_XDECREF(t);
}
""" % locals()
return code
def c_code_cache_version(self):
return (4,)
def __str__(self):
return "%s{%s}" % (self.__class__.__name__, self.axis)
return CumOp(axis=axis, mode='add')(x)
def cumprod(x, axis=None):
......@@ -486,7 +400,7 @@ def cumprod(x, axis=None):
.. versionadded:: 0.7
"""
return CumprodOp(axis=axis)(x)
return CumOp(axis=axis, mode='mul')(x)
class DiffOp(theano.Op):
......
theano/tensor/tests/test_extra_ops.py
浏览文件 @ 00189784
from __future__ import absolute_import, print_function, division
from functools import partial
import numpy as np
import numpy
......@@ -7,7 +8,7 @@ import theano
from theano.tests import unittest_tools as utt
from theano.tensor.extra_ops import (SearchsortedOp, searchsorted,
CumsumOp, cumsum, CumprodOp, cumprod,
CumOp, cumsum, cumprod,
CpuContiguous, cpu_contiguous,
bincount, DiffOp, diff, squeeze, compress,
RepeatOp, repeat, Bartlett, bartlett,
......@@ -121,74 +122,33 @@ class TestSearchsortedOp(utt.InferShapeTester):
utt.verify_grad(self.op, [self.a[self.idx_sorted], self.b])
class TestCumsumOp(utt.InferShapeTester):
class TestCumOp(utt.InferShapeTester):
def setUp(self):
super(TestCumsumOp, self).setUp()
self.op_class = CumsumOp
self.op = CumsumOp()
super(TestCumOp, self).setUp()
self.op_class = CumOp
self.op = CumOp()
def test_cumsumOp(self):
def test_cum_op(self):
x = T.tensor3('x')
a = np.random.random((3, 5, 2)).astype(config.floatX)
# Test axis out of bounds
self.assertRaises(ValueError, cumsum, x, axis=3)
self.assertRaises(ValueError, cumsum, x, axis=-4)
f = theano.function([x], cumsum(x))
assert np.allclose(np.cumsum(a), f(a)) # Test axis=None
for axis in range(-len(a.shape), len(a.shape)):
f = theano.function([x], cumsum(x, axis=axis))
assert np.allclose(np.cumsum(a, axis=axis), f(a))
def test_infer_shape(self):
x = T.tensor3('x')
a = np.random.random((3, 5, 2)).astype(config.floatX)
# Test axis=None
self._compile_and_check([x],
[self.op(x)],
[a],
self.op_class)
for axis in range(-len(a.shape), len(a.shape)):
self._compile_and_check([x],
[cumsum(x, axis=axis)],
[a],
self.op_class)
def test_grad(self):
a = np.random.random((3, 5, 2)).astype(config.floatX)
utt.verify_grad(self.op, [a]) # Test axis=None
for axis in range(-len(a.shape), len(a.shape)):
utt.verify_grad(self.op_class(axis=axis), [a], eps=4e-4)
class TestCumprodOp(utt.InferShapeTester):
def setUp(self):
super(TestCumprodOp, self).setUp()
self.op_class = CumprodOp
self.op = CumprodOp()
def test_CumprodOp(self):
x = T.tensor3('x')
a = np.random.random((3, 5, 2)).astype(config.floatX)
# Test axis out of bounds
self.assertRaises(ValueError, cumprod, x, axis=3)
self.assertRaises(ValueError, cumprod, x, axis=-4)
f = theano.function([x], cumprod(x))
assert np.allclose(np.cumprod(a), f(a)) # Test axis=None
f = theano.function([x], [cumsum(x), cumprod(x)])
s, p = f(a)
assert np.allclose(np.cumsum(a), s) # Test axis=None
assert np.allclose(np.cumprod(a), p) # Test axis=None
for axis in range(-len(a.shape), len(a.shape)):
f = theano.function([x], cumprod(x, axis=axis))
assert np.allclose(np.cumprod(a, axis=axis), f(a))
f = theano.function([x], [cumsum(x, axis=axis), cumprod(x, axis=axis)])
s, p = f(a)
assert np.allclose(np.cumsum(a, axis=axis), s)
assert np.allclose(np.cumprod(a, axis=axis), p)
def test_infer_shape(self):
x = T.tensor3('x')
......@@ -202,17 +162,19 @@ class TestCumprodOp(utt.InferShapeTester):
for axis in range(-len(a.shape), len(a.shape)):
self._compile_and_check([x],
[cumprod(x, axis=axis)],
[cumsum(x, axis=axis)],
[a],
self.op_class)
def test_grad(self):
a = np.random.random((3, 5, 2)).astype(config.floatX)
utt.verify_grad(self.op, [a]) # Test axis=None
utt.verify_grad(self.op_class(mode='add'), [a]) # Test axis=None
utt.verify_grad(self.op_class(mode='mul'), [a]) # Test axis=None
for axis in range(-len(a.shape), len(a.shape)):
utt.verify_grad(self.op_class(axis=axis), [a])
utt.verify_grad(self.op_class(axis=axis, mode='add'), [a], eps=4e-4)
utt.verify_grad(self.op_class(axis=axis, mode='mul'), [a], eps=4e-4)
class TestBinCount(utt.InferShapeTester):
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论