Skip to content

P
pytensor
提交 084bd965 authored 作者: Tim Cooijmans's avatar Tim Cooijmans
浏览文件
操作

map_variables: move to scan_utils

上级 c5502dab
隐藏空白字符变更
内嵌 并排
正在显示 包含 287 行增加 和 290 行删除
theano/gof/graph.py
浏览文件 @ 084bd965
......@@ -816,199 +816,6 @@ def clone_get_equiv(inputs, outputs, copy_inputs_and_orphans=True, memo=None):
return memo
def map_variables(replacer, graphs, additional_inputs=[]):
"""Construct new graphs based on 'graphs' with some variables replaced
according to 'replacer'.
:param replacer: function that takes a variable and returns its
replacement.
:param graphs: an iterable of graphs in which to replace variables
:param additional_inputs: an iterable of graph inputs not used in any
of 'graphs' but possibly used in the graphs returned by `replacer`
:return: the new graphs, in the same order as 'graphs'
Example:
.. code-block:: python
tag = "replaceme"
a = tensor.scalar("a")
b = tensor.scalar("b")
c = tensor.scalar("c")
ab = a + b
ab.tag.replacement = a * b
u = ab + c
v, = map_variables(lambda graph:
return getattr(graph.tag, "replacement", graph),
[u])
# v is now equal to a * b + c
"""
from theano.gof.fg import FunctionGraph
from theano.gof.opt import TopoOptimizer, local_optimizer
from theano import clone as the_other_clone
from theano.scan_module.scan_op import Scan
from theano.compile import OpFromGraph
# wrap replacer to avoid replacing things we just put there.
graphs_seen = set()
def wrapped_replacer(graph):
if graph in graphs_seen:
return graph
else:
new_graph = replacer(graph)
graphs_seen.add(new_graph)
return new_graph
graphs = list(graphs)
inputs_ = list(set(inputs(graphs) + list(additional_inputs)))
# perform any desired replacement of input variables. these
# aren't replaced by the local optimizer approach because they are
# not outputs of any Apply node.
new_inputs = list(map(wrapped_replacer, inputs_))
replacements = [(input_, new_input)
for input_, new_input
in zip(inputs_, new_inputs)
if new_input is not input_]
graphs = the_other_clone(graphs, share_inputs=True, replace=replacements)
inputs_ = list(set(inputs(graphs) + list(additional_inputs)))
# clone cached constants or FunctionGraph will complain. this has
# to occur in a separate pass from the replacement above because
# both may suggest different replacements for the same variables.
# since the replacements introduced above may involve cached
# constants, the replacement of said constants has to come after.
cached_constants = [x for x in inputs_ if getattr(x, "cached", False)]
copied_constants, _ = clone(cached_constants, [], copy_inputs=True)
replacements = list(zip(cached_constants, copied_constants))
inputs_ = list(set(inputs_) - set(cached_constants)) + list(copied_constants)
graphs = the_other_clone(graphs, share_inputs=True, replace=replacements)
fg = FunctionGraph(inputs_, graphs, clone=False)
nodes_seen = set()
@local_optimizer(None)
def local_transform(node):
if node in nodes_seen:
return False
if isinstance(node.op, (Scan, OpFromGraph)):
# recurse on the inner graph
(new_inner_inputs,
new_outer_inputs,
new_inner_outputs) = _map_variables_inner(
wrapped_replacer,
inner_inputs=node.op.inputs,
outer_inputs=node.inputs,
inner_outputs=node.op.outputs)
# reinstantiate the op
if isinstance(node.op, Scan):
new_op = Scan(new_inner_inputs,
new_inner_outputs,
node.op.info,
# FIXME: infer this someday?
typeConstructor=None)
elif isinstance(node.op, OpFromGraph):
new_op = OpFromGraph(new_inner_inputs,
new_inner_outputs,
**node.op.kwargs)
# make a new node to replace the old one
new_node = new_op.make_node(*new_outer_inputs)
nodes_seen.add(new_node)
return new_node.outputs
else:
nodes_seen.add(node)
return list(map(wrapped_replacer, node.outputs))
topo_transform = TopoOptimizer(local_transform, 'out_to_in')
topo_transform.optimize(fg)
new_graphs = fg.outputs
theano.printing.debugprint(new_graphs)
fg.disown()
return new_graphs
def _map_variables_inner(replacer, inner_inputs, outer_inputs, inner_outputs):
# the replacements returned by the replacer may involve variables
# that are already owned by the outer fgraph (`fg` in the caller)
# and so cannot be added to the inner fgraph (`fg` in the
# recursive call). wrap the replacer to catch these before they
# are added.
# additionally, some of these may be fgraph inputs or shared
# variables, which we cannot directly use inside the inner graph.
# we need to create inner inputs to access them through.
outer_to_inner = dict(zip(outer_inputs, inner_inputs))
extra_inner_inputs = []
extra_outer_inputs = []
from theano.scan_module import scan_utils
from itertools import chain
from theano import gof
def inner_replacer(graph):
new_graph = replacer(graph)
other_inputs = []
constants = []
for input_ in gof.graph.inputs([new_graph]):
if isinstance(input_, gof.Variable):
if isinstance(input_, Constant):
constants.append(input_)
else:
other_inputs.append(input_)
# foreign inputs are fgraph inputs and shared variables that we need
# to access through inner inputs
foreign_inputs = list(set(other_inputs) - set(outer_to_inner.values()))
# skip further processing if there is nothing to do
#if not constants and not foreign_inputs:
# return new_graph
replacements = []
# constants just need to be replaced by copies that the inner
# `fg` can take ownership of
for input_ in constants:
new_input = input_.clone()
new_input.name = "%s_copiedd" % new_input.name
replacements.append((input_, new_input))
for outer_input in foreign_inputs:
# if this foreign input is not already available
# as an inner input, connect it through a new
# inner input
if outer_input not in outer_to_inner.keys():
inner_input = scan_utils.safe_new(outer_input, tag="_copy")
outer_to_inner[outer_input] = inner_input
extra_inner_inputs.append(inner_input)
extra_outer_inputs.append(outer_input)
# the inner FunctionGraph wants to know its inputs
# beforehand, but we don't always know. so add them
# as we discover them.
graph.owner.fgraph.add_input(inner_input)
replacements.extend(outer_to_inner.items())
new_graph, = theano.clone([new_graph],
share_inputs=True,
replace=replacements)
return new_graph
new_inner_outputs = map_variables(inner_replacer, inner_outputs)
new_inner_inputs = list(chain(inner_inputs, extra_inner_inputs))
new_outer_inputs = list(chain(outer_inputs, extra_outer_inputs))
return new_inner_inputs, new_outer_inputs, new_inner_outputs
def general_toposort(r_out, deps, debug_print=False,
compute_deps_cache=None, deps_cache=None):
"""
......
theano/gof/tests/test_graph.py
浏览文件 @ 084bd965
......@@ -11,7 +11,7 @@ from theano import (
from theano.gof.graph import (
Apply,
as_string, clone, general_toposort, inputs, io_toposort,
is_same_graph, Variable, map_variables)
is_same_graph, Variable)
from theano.gof.op import Op
from theano.gof.type import Type
from theano.sandbox.cuda.var import (
......@@ -158,101 +158,6 @@ class TestClone(X):
assert self.str(inputs(node.outputs), node.outputs) == ["MyOp(MyOp(R1, R2), R5)"]
#################
# map_variables #
#################
class TestMapVariables(X):
@staticmethod
def replacer(graph):
return getattr(graph.tag, "replacement", graph)
def test_leaf(self):
a = tensor.scalar("a")
b = tensor.scalar("b")
c = tensor.scalar("c")
b.tag.replacement = c
u = a + b
v, = map_variables(self.replacer, [u])
assert u.owner.inputs == [a, b]
assert v.owner.inputs == [a, c]
def test_opfromgraph(self):
from theano import OpFromGraph, function
import itertools
a = tensor.scalar()
b = tensor.scalar()
r = a + b
r.tag.replacement = a - b
c = tensor.scalar()
d = tensor.scalar()
u = OpFromGraph([a, b], [r])(c, d)
v, = map_variables(self.replacer, [u])
f = function([c, d], [u, v])
for m, n in itertools.combinations(range(10), 2):
assert f(m, n) == [m + n, m - n]
def test_scan(self):
from theano import shared, scan, function
x = tensor.vector('x')
# we will insert a subgraph involving these variables into the inner
# graph of scan. since they were not previously in the inner graph,
# they are like non_sequences to scan(). scan() infers these and
# imports them into the inner graph properly, and map_variables()
# should do this as well.
outer = tensor.scalar("outer")
shared = shared(1, name="shared")
constant = tensor.constant(1, name="constant")
# z will equal 1 so multiplying by it doesn't change any values
z = outer * (shared + constant)
def step(x, a):
r = a + x
r.tag.replacement = z * (a - x)
return r
s, _ = scan(step, sequences=x,
outputs_info=[numpy.array(0.)])
# ensure z is owned by the outer graph so map_variables() will need to
# jump through additional hoops to placate FunctionGraph.
t = z * s
s2, = map_variables(self.replacer, [t])
t2 = z * s2
f = function([x, outer], [t, t2])
rval = f(x=numpy.array([1, 2, 3]), outer=0.5)
assert numpy.array_equal(rval, [[ 1, 3, 6], [-1, -3, -6]])
def test_leaf_inside_scan(self):
import numpy
from theano import function, scan
x = tensor.vector('x')
y = tensor.scalar('y')
z = tensor.scalar('z')
y.tag.replacement = z
s, _ = scan(lambda x: x * y, sequences=x)
s2, = map_variables(self.replacer, [s])
f = function([x, y, z], [s, s2])
assert numpy.array_equal(
f(x=numpy.array([1, 2, 3]), y=1, z=2),
[[ 1, 2, 3],
[ 2, 4, 6]])
############
# toposort #
############
......
theano/scan_module/scan_utils.py
浏览文件 @ 084bd965
......@@ -256,6 +256,201 @@ def clone(output,
return outs
def map_variables(replacer, graphs, additional_inputs=[]):
"""Construct new graphs based on 'graphs' with some variables replaced
according to 'replacer'.
:param replacer: function that takes a variable and returns its
replacement.
:param graphs: an iterable of graphs in which to replace variables
:param additional_inputs: an iterable of graph inputs not used in any
of 'graphs' but possibly used in the graphs returned by `replacer`
:return: the new graphs, in the same order as 'graphs'
Example:
.. code-block:: python
tag = "replaceme"
a = tensor.scalar("a")
b = tensor.scalar("b")
c = tensor.scalar("c")
ab = a + b
ab.tag.replacement = a * b
u = ab + c
v, = map_variables(lambda graph:
return getattr(graph.tag, "replacement", graph),
[u])
# v is now equal to a * b + c
"""
# wrap replacer to avoid replacing things we just put there.
graphs_seen = set()
def wrapped_replacer(graph):
if graph in graphs_seen:
return graph
else:
new_graph = replacer(graph)
graphs_seen.add(new_graph)
return new_graph
graphs = list(graphs)
inputs_ = list(set(gof.graph.inputs(graphs) + list(additional_inputs)))
# perform any desired replacement of input variables. these
# aren't replaced by the local optimizer approach because they are
# not outputs of any Apply node.
new_inputs = list(map(wrapped_replacer, inputs_))
replacements = [(input_, new_input)
for input_, new_input
in zip(inputs_, new_inputs)
if new_input is not input_]
graphs = clone(graphs, share_inputs=True, replace=replacements)
inputs_ = list(set(gof.graph.inputs(graphs) + list(additional_inputs)))
# clone cached constants or FunctionGraph will complain. this has
# to occur in a separate pass from the replacement above because
# both may suggest different replacements for the same variables.
# since the replacements introduced above may involve cached
# constants, the replacement of said constants has to come after.
cached_constants = [x for x in inputs_ if getattr(x, "cached", False)]
copied_constants = clone(cached_constants, share_inputs=False)
replacements = list(zip(cached_constants, copied_constants))
inputs_ = list(set(inputs_) - set(cached_constants)) + list(copied_constants)
graphs = clone(graphs, share_inputs=True, replace=replacements)
fg = gof.fg.FunctionGraph(inputs_, graphs, clone=False)
nodes_seen = set()
@gof.opt.local_optimizer(None)
def local_transform(node):
if node in nodes_seen:
return False
# importing Scan into module scope would be circular
from theano.scan_module.scan_op import Scan
from theano.compile import OpFromGraph
if isinstance(node.op, (Scan, OpFromGraph)):
# recurse on the inner graph
(new_inner_inputs,
new_outer_inputs,
new_inner_outputs) = _map_variables_inner(
wrapped_replacer,
inner_inputs=node.op.inputs,
outer_inputs=node.inputs,
inner_outputs=node.op.outputs)
# reinstantiate the op
if isinstance(node.op, Scan):
new_op = Scan(new_inner_inputs,
new_inner_outputs,
node.op.info,
# FIXME: infer this someday?
typeConstructor=None)
elif isinstance(node.op, OpFromGraph):
new_op = OpFromGraph(new_inner_inputs,
new_inner_outputs,
**node.op.kwargs)
# make a new node to replace the old one
new_node = new_op.make_node(*new_outer_inputs)
nodes_seen.add(new_node)
return new_node.outputs
else:
nodes_seen.add(node)
return list(map(wrapped_replacer, node.outputs))
topo_transform = gof.opt.TopoOptimizer(local_transform, 'out_to_in')
topo_transform.optimize(fg)
new_graphs = fg.outputs
theano.printing.debugprint(new_graphs)
fg.disown()
return new_graphs
def _map_variables_inner(replacer, inner_inputs, outer_inputs, inner_outputs):
# the replacements returned by the replacer may involve variables
# that are already owned by the outer fgraph (`fg` in the caller)
# and so cannot be added to the inner fgraph (`fg` in the
# recursive call). wrap the replacer to catch these before they
# are added.
# additionally, some of these may be fgraph inputs or shared
# variables, which we cannot directly use inside the inner graph.
# we need to create inner inputs to access them through.
# TODO: handle potential updates of newly introduced shared variables.
outer_to_inner = dict(zip(outer_inputs, inner_inputs))
extra_inner_inputs = []
extra_outer_inputs = []
from theano.scan_module import scan_utils
from itertools import chain
from theano import gof
def inner_replacer(graph):
new_graph = replacer(graph)
other_inputs = []
constants = []
for input_ in gof.graph.inputs([new_graph]):
if isinstance(input_, gof.Variable):
if isinstance(input_, gof.Constant):
constants.append(input_)
else:
other_inputs.append(input_)
# foreign inputs are fgraph inputs and shared variables that we need
# to access through inner inputs
foreign_inputs = list(set(other_inputs) - set(outer_to_inner.values()))
# skip further processing if there is nothing to do
#if not constants and not foreign_inputs:
# return new_graph
replacements = []
# constants just need to be replaced by copies that the inner
# `fg` can take ownership of
for input_ in constants:
new_input = input_.clone()
new_input.name = "%s_copiedd" % new_input.name
replacements.append((input_, new_input))
for outer_input in foreign_inputs:
# if this foreign input is not already available
# as an inner input, connect it through a new
# inner input
if outer_input not in outer_to_inner.keys():
inner_input = scan_utils.safe_new(outer_input, tag="_copy")
outer_to_inner[outer_input] = inner_input
extra_inner_inputs.append(inner_input)
extra_outer_inputs.append(outer_input)
# the inner FunctionGraph wants to know its inputs
# beforehand, but we don't always know. so add them
# as we discover them.
graph.owner.fgraph.add_input(inner_input)
replacements.extend(outer_to_inner.items())
new_graph, = theano.clone([new_graph],
share_inputs=True,
replace=replacements)
return new_graph
new_inner_outputs = map_variables(inner_replacer, inner_outputs)
new_inner_inputs = list(chain(inner_inputs, extra_inner_inputs))
new_outer_inputs = list(chain(outer_inputs, extra_outer_inputs))
return new_inner_inputs, new_outer_inputs, new_inner_outputs
def get_updates_and_outputs(ls):
"""
This function tries to recognize the updates OrderedDict, the
......
theano/scan_module/tests/test_scan_utils.py
浏览文件 @ 084bd965
import itertools
import numpy
import theano
from theano.scan_module.scan_utils import equal_computations
from theano import tensor
from theano.scan_module.scan_utils import equal_computations, map_variables
from theano.tensor.type_other import NoneConst
......@@ -11,3 +14,90 @@ def test_equal_compuations():
max_argmax1 = theano.tensor.max_and_argmax(m)
max_argmax2 = theano.tensor.max_and_argmax(m)
assert equal_computations(max_argmax1, max_argmax2)
#################
# map_variables #
#################
class TestMapVariables(object):
@staticmethod
def replacer(graph):
return getattr(graph.tag, "replacement", graph)
def test_leaf(self):
a = tensor.scalar("a")
b = tensor.scalar("b")
c = tensor.scalar("c")
b.tag.replacement = c
u = a + b
v, = map_variables(self.replacer, [u])
assert u.owner.inputs == [a, b]
assert v.owner.inputs == [a, c]
def test_opfromgraph(self):
a = tensor.scalar()
b = tensor.scalar()
r = a + b
r.tag.replacement = a - b
c = tensor.scalar()
d = tensor.scalar()
u = theano.OpFromGraph([a, b], [r])(c, d)
v, = map_variables(self.replacer, [u])
f = theano.function([c, d], [u, v])
for m, n in itertools.combinations(range(10), 2):
assert f(m, n) == [m + n, m - n]
def test_scan(self):
x = tensor.vector('x')
# we will insert a subgraph involving these variables into the inner
# graph of scan. since they were not previously in the inner graph,
# they are like non_sequences to scan(). scan() infers these and
# imports them into the inner graph properly, and map_variables()
# should do this as well.
outer = tensor.scalar("outer")
shared = theano.shared(1, name="shared")
constant = tensor.constant(1, name="constant")
# z will equal 1 so multiplying by it doesn't change any values
z = outer * (shared + constant)
def step(x, a):
r = a + x
r.tag.replacement = z * (a - x)
return r
s, _ = theano.scan(step, sequences=x,
outputs_info=[numpy.array(0.)])
# ensure z is owned by the outer graph so map_variables() will need to
# jump through additional hoops to placate FunctionGraph.
t = z * s
s2, = map_variables(self.replacer, [t])
t2 = z * s2
f = theano.function([x, outer], [t, t2])
rval = f(x=numpy.array([1, 2, 3]), outer=0.5)
assert numpy.array_equal(rval, [[ 1, 3, 6], [-1, -3, -6]])
def test_leaf_inside_scan(self):
x = tensor.vector('x')
y = tensor.scalar('y')
z = tensor.scalar('z')
y.tag.replacement = z
s, _ = theano.scan(lambda x: x * y, sequences=x)
s2, = map_variables(self.replacer, [s])
f = theano.function([x, y, z], [s, s2])
assert numpy.array_equal(
f(x=numpy.array([1, 2, 3]), y=1, z=2),
[[ 1, 2, 3],
[ 2, 4, 6]])
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论