提交 aeb8c035 authored 作者: Xavier Bouthillier's avatar Xavier Bouthillier

Fix optimizations

上级 76b71018
......@@ -84,7 +84,19 @@ class SparseBlockGemv(Op):
return Apply(self, [o, W, h, inputIdx, outputIdx], [output])
def perform(self, node, inp, out_):
raise NotImplementedError('Optimization of SparseBlockGemv failed.')
o, W, h, iIdx, oIdx = inp[:5]
if not self.inplace:
o = o.copy()
for b in range(o.shape[0]):
for j in range(o.shape[1]):
outputIdx = oIdx[b, j]
for i in range(h.shape[1]):
inputIdx = iIdx[b, i]
w = W[inputIdx, outputIdx]
o[b, j, :] += numpy.dot(h[b, i], w)
out_[0][0] = o
def grad(self, inputs, grads):
o, W, h, inputIdx, outputIdx = inputs
......@@ -160,50 +172,6 @@ class SparseBlockOuter(Op):
return Apply(self, [o, x, y, xIdx, yIdx, alpha],
[output])
def perform(self, node, inp, out_):
raise NotImplementedError('Optimization of SparseBlockOuter failed.')
def grad(self, inputs, output_gradients):
raise NotImplementedError("SparseBlockOuter has no gradient "
"implemented")
class CpuSparseBlockGemv(SparseBlockGemv):
"""
CPU version of SparseBlockGemv. Check SparseBlockGemv's docstring for more
information.
This should not be directly called since the interface is subject
to change without notice. Use the sandbox.blocksparse.sparse_block_dot()
function for a stable interface.
"""
def perform(self, node, inp, out_):
o, W, h, iIdx, oIdx = inp[:5]
if not self.inplace:
o = o.copy()
for b in range(o.shape[0]):
for j in range(o.shape[1]):
outputIdx = oIdx[b, j]
for i in range(h.shape[1]):
inputIdx = iIdx[b, i]
w = W[inputIdx, outputIdx]
o[b, j, :] += numpy.dot(h[b, i], w)
out_[0][0] = o
class CpuSparseBlockOuter(SparseBlockOuter):
"""
CPU version of SparseBlockOuter. See SparseBlockOuter's docstring for more
information.
This op should not be called directly since its interface is
subject to change without notice. It is involved in the gradient
of GpuSparseBlockGemv. The gradient is not implemented.
"""
def perform(self, node, inp, out_):
o, x, y, xIdx, yIdx, alpha = inp[:6]
......@@ -223,11 +191,6 @@ sparse_block_gemv_inplace = SparseBlockGemv(True)
sparse_block_outer = SparseBlockOuter(False)
sparse_block_outer_inplace = SparseBlockOuter(True)
cpu_sparse_block_gemv = CpuSparseBlockGemv(False)
cpu_sparse_block_gemv_inplace = CpuSparseBlockGemv(True)
cpu_sparse_block_outer = CpuSparseBlockOuter(False)
cpu_sparse_block_outer_inplace = CpuSparseBlockOuter(True)
def sparse_block_dot(W, h, inputIdx, b, outputIdx):
"""
......
差异被折叠。
......@@ -2,170 +2,42 @@
Optimizations addressing the ops in sandbox root directory
"""
import bisect
import logging
from theano.compile import optdb
from theano.gof import local_optimizer, EquilibriumDB
from theano.tensor.opt import register_specialize
from theano import compile # to register the optimizer built by this file
from theano import gof
from theano.sandbox.blocksparse import (
SparseBlockGemv,
SparseBlockOuter,
sparse_block_gemv,
sparse_block_outer,
sparse_block_gemv_inplace,
sparse_block_outer_inplace,
CpuSparseBlockGemv,
CpuSparseBlockOuter)
_logger = logging.getLogger('theano.sandbox.opt')
def _db_exists(db, db_name):
"""
Tests whether the full path from `db_name[0]` down to
`db_name[-1]` exists.
Parameters
----------
db: `theano.gof.optdb.DB`
A dataset of optimisations or sub-datasets.
db_name: list or tuple of strings
Names of datasets from given one `db[db_name[0]]` down
to the dataset of interest where to register.
ex: ['level_1_dataset', 'level_2_dataset']
"""
if len(db_name) == 1:
return db_name[0] in db._names
return db_name[0] in db._names and _db_exists(db[db_name[0]], db_name[1:])
def _db_register(db, db_name, *args):
"""
Registers an object in last datasets given in db_name. `db_name[-1]`
is deep in the hierarchy of `db`.
Parameters
----------
db: `theano.gof.optdb.DB`
A dataset of optimisations or sub-datasets.
db_name: list or tuple of strings
Names of datasets from given one `db[db_name[0]]` down
to the dataset of interest where to register.
ex: ['level_1_dataset', 'level_2_dataset']
"""
if len(db_name) == 0:
return db.register(*args)
return _db_register(db[db_name[0]], db_name[1:], *args)
def _db_positions(db, db_name, positions=()):
"""
Returns the list of positions of all databases from `db_name[0]`
down to `db_name[-1]`. The path is hierarchical, hence `db_name[0]`
is in `db`, `db_name[1]` is in `db[db_name[0]]`, etc.
Parameters
----------
db: `theano.gof.optdb.DB`
A dataset of optimisations or sub-datasets.
db_name: list or tuple of strings
Names of datasets from given one `db[db_name[0]]` down
to the dataset of interests.
ex: ['level_1_dataset', 'level_2_dataset']
"""
if len(db_name) == 0:
return positions
db_position = db.__position__.get(db_name[0], 0.)
return _db_positions(db[db_name[0]], db_name[1:],
positions + (db_position, ))
def register_meta_opt(op_class, db_name, position, *args):
"""
Registers a given optimization under given database name and saves
optimization information in `op_class.registered_opts`.
Parameters
----------
op_class: `theano.gof.Op`
A meta Op which have multiple implementations available
for optimization.
db_name: string, list or tuple of strings
A string if optimization is inserted in `theano.compile.optdb`
directly. List is used to insert an optimization deep inside a
hierarchy of optimization databases.
position: int or float
Position of the optimisation in the target dataset.
(Position in deep database if not optdb)
*args
Arguments to register the optimization.
"""
if isinstance(db_name, str):
db_name = [db_name]
def call(local_meta_opt):
if not _db_exists(optdb, db_name):
# TODO: Would another default DB be better?
_db_register(optdb, db_name[:-2],
db_name[-1], EquilibriumDB(), position, *args)
_db_register(optdb, db_name,
local_meta_opt.__name__, local_meta_opt, *args)
positions = _db_positions(optdb, db_name)
idx = bisect.bisect_left((positions, local_meta_opt),
op_class.registered_opts)
op_class.registered_opts.insert(idx,
(positions, local_meta_opt.__name__))
return local_meta_opt
return call
@register_meta_opt(SparseBlockGemv, ["meta_cpu"], 51.0,
"fast_run", "fast_compile")
@local_optimizer([SparseBlockGemv])
def cpu_sparse_block_gemv_opt(node):
"""
SparseBlockGemv -> CpuSparseBlockGemv
"""
return [CpuSparseBlockGemv(node.op.inplace)(*node.inputs)]
@register_meta_opt(SparseBlockOuter, ["meta_cpu"], 51.0,
"fast_run", "fast_compile")
@local_optimizer([SparseBlockOuter])
def cpu_sparse_block_outer_opt(node):
"""
SparseBlockOuter -> CpuSparseBlockOuter
"""
return [CpuSparseBlockOuter(node.op.inplace)(*node.inputs)]
sparse_block_outer_inplace)
@register_specialize
@local_optimizer([sparse_block_gemv], inplace=True)
def local_inplace_block_sparse_gemv(node):
@gof.local_optimizer([SparseBlockGemv], inplace=True)
def local_inplace_sparse_block_gemv(node):
"""
SparseBlockGemv(inplace=False) -> SparseBlockGemv(inplace=True)
"""
return [sparse_block_gemv_inplace(*node.inputs)]
if isinstance(node.op, SparseBlockGemv) and not node.op.inplace:
new_node = sparse_block_gemv_inplace(*node.inputs)
return [new_node]
return False
compile.optdb.register('local_inplace_sparse_block_gemv',
gof.TopoOptimizer(
local_inplace_sparse_block_gemv,
failure_callback=gof.TopoOptimizer.warn_inplace),
60, 'fast_run', 'inplace') # DEBUG
@register_specialize
@local_optimizer([sparse_block_outer], inplace=True)
def local_inplace_block_sparse_outer(node):
@gof.local_optimizer([SparseBlockOuter], inplace=True)
def local_inplace_sparse_block_outer(node):
"""
SparseBlockOuter(inplace=False) -> SparseBlockOuter(inplace=True)
"""
return [sparse_block_outer_inplace(*node.inputs)]
if isinstance(node.op, SparseBlockOuter) and not node.op.inplace:
new_node = sparse_block_outer_inplace(*node.inputs)
return [new_node]
return False
compile.optdb.register('local_inplace_sparse_block_outer',
gof.TopoOptimizer(
local_inplace_sparse_block_outer,
failure_callback=gof.TopoOptimizer.warn_inplace),
60, 'fast_run', 'inplace') # DEBUG
......@@ -11,7 +11,7 @@ from theano import tensor
import theano.tests.unittest_tools as utt
from theano.sandbox.blocksparse import sparse_block_dot, \
cpu_sparse_block_gemv, cpu_sparse_block_outer
sparse_block_gemv, sparse_block_outer
class BlockSparse_Gemv_and_Outer(unittest.TestCase):
......@@ -24,8 +24,8 @@ class BlockSparse_Gemv_and_Outer(unittest.TestCase):
self.mode = theano.compile.get_default_mode().excluding(
'constant_folding'
)
self.gemv_op = cpu_sparse_block_gemv
self.outer_op = cpu_sparse_block_outer
self.gemv_op = sparse_block_gemv
self.outer_op = sparse_block_outer
@staticmethod
def gemv_data():
......
import theano
from theano import tensor
from theano.sandbox.blocksparse import CpuSparseBlockGemv, \
CpuSparseBlockOuter, sparse_block_dot
from theano.sandbox.blocksparse import sparse_block_dot
def test_blocksparse_cpu_gemv_opt():
def test_blocksparse_inplace_gemv_opt():
b = tensor.fmatrix()
W = tensor.ftensor4()
h = tensor.ftensor3()
......@@ -15,10 +14,13 @@ def test_blocksparse_cpu_gemv_opt():
f = theano.function([W, h, iIdx, b, oIdx], o)
assert isinstance(f.maker.fgraph.toposort()[-1].op, CpuSparseBlockGemv)
if theano.config.mode == "FAST_COMPILE":
assert not f.maker.fgraph.toposort()[-1].op.inplace
else:
assert f.maker.fgraph.toposort()[-1].op.inplace
def test_blocksparse_cpu_outer_opt():
def test_blocksparse_inplace_outer_opt():
b = tensor.fmatrix()
W = tensor.ftensor4()
h = tensor.ftensor3()
......@@ -32,4 +34,7 @@ def test_blocksparse_cpu_outer_opt():
f = theano.function([W, h, iIdx, b, oIdx],
[o, tensor.grad(o.sum(), wrt=W)])
assert isinstance(f.maker.fgraph.toposort()[-1].op, CpuSparseBlockOuter)
if theano.config.mode == "FAST_COMPILE":
assert not f.maker.fgraph.toposort()[-1].op.inplace
else:
assert f.maker.fgraph.toposort()[-1].op.inplace
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论