Skip to content

P
pytensor
提交 5f2ae036 authored 作者: Caglar's avatar Caglar
浏览文件
操作

changed the order in the if condition and remove the while loop.

上级 6fe5f212
隐藏空白字符变更
内嵌 并排
正在显示 包含 162 行增加 和 171 行删除
theano/scan_module/scan_opt.py
浏览文件 @ 5f2ae036
...@@ -206,7 +206,10 @@ def remove_constants_and_unused_inputs_scan(node): ...@@ -206,7 +206,10 @@ def remove_constants_and_unused_inputs_scan(node):
# This is a global opt for historical reason # This is a global opt for historical reason
# It should be possible to change it to a local opt. # It should be possible to change it to a local opt.
class PushOutNonSeqScan(gof.Optimizer): class PushOutNonSeqScan(gof.Optimizer):
"""
A global optimizer for pushing out the variables from scan that are not
used inside the scan.
"""
def __init__(self): def __init__(self):
gof.Optimizer.__init__(self) gof.Optimizer.__init__(self)
...@@ -220,24 +223,31 @@ class PushOutNonSeqScan(gof.Optimizer): ...@@ -220,24 +223,31 @@ class PushOutNonSeqScan(gof.Optimizer):
self.process_node(fgraph, node) self.process_node(fgraph, node)
def process_node(self, fgraph, node): def process_node(self, fgraph, node):
"""
Important note: This function uses set and dictionary data structure.
By default they are not ordered for efficiency reasons. Take care and make
sure of changing them to Ordered versions if you need to iterate over those
variables.
"""
# this flag tells if there was any change during the last iterations # this flag tells if there was any change during the last iterations
changed = True
clean_inputs, clean_outputs = scan_utils.reconstruct_graph( clean_inputs, clean_outputs = scan_utils.reconstruct_graph(
node.op.inputs, node.op.outputs) node.op.inputs, node.op.outputs)
local_fgraph = gof.FunctionGraph(clean_inputs, clean_outputs, clone=False) local_fgraph = gof.FunctionGraph(clean_inputs,
clean_outputs,
clone=False)
local_fgraph_topo = local_fgraph.toposort() local_fgraph_topo = local_fgraph.toposort()
local_fgraph_outs_set = set(local_fgraph.outputs) local_fgraph_outs_set = set(local_fgraph.outputs)
local_fgraph_outs_map = dict([(v, k) for k,v in enumerate(local_fgraph.outputs)]) local_fgraph_outs_map = dict([(v, k) for k,v in enumerate(local_fgraph.outputs)])
max_iterations = 2 * len(local_fgraph_topo) + 3
counts = 0
to_remove_set = set() to_remove_set = set()
to_remove_add = to_remove_set.add to_remove_add = to_remove_set.add
to_replace_set = set() to_replace_set = set()
to_replace_add = to_replace_set.add to_replace_add = to_replace_set.add
to_replace_map = {} to_replace_map = OrderedDict()
nto_replace = 0 nto_replace = 0
def add_to_replace(y, nto_replace): def add_to_replace(y, nto_replace):
to_replace_add(y) to_replace_add(y)
to_replace_map[y] = nto_replace to_replace_map[y] = nto_replace
...@@ -245,6 +255,7 @@ class PushOutNonSeqScan(gof.Optimizer): ...@@ -245,6 +255,7 @@ class PushOutNonSeqScan(gof.Optimizer):
replace_with_in = [] replace_with_in = []
replace_with_out = [] replace_with_out = []
op = node.op op = node.op
# Construct the list of non_sequences to simplify a few things # Construct the list of non_sequences to simplify a few things
inner_non_seqs = op.inner_non_seqs(clean_inputs) inner_non_seqs = op.inner_non_seqs(clean_inputs)
...@@ -259,65 +270,55 @@ class PushOutNonSeqScan(gof.Optimizer): ...@@ -259,65 +270,55 @@ class PushOutNonSeqScan(gof.Optimizer):
assert len(inner_non_seqs) == len(outer_non_seqs) assert len(inner_non_seqs) == len(outer_non_seqs)
assert len(inner_seqs) == len(outer_seqs) assert len(inner_seqs) == len(outer_seqs)
while changed and counts < max_iterations: for nd in local_fgraph_topo:
counts += 1 if (# we haven't already looked at this node
changed = False nd not in to_remove_set and
all([((x in inner_non_seqs_set) or
for nd in local_fgraph_topo: (x.owner in to_remove_set) or
if (all([(x in inner_non_seqs_set) or isinstance(x, tensor.Constant))
(x.owner in to_remove_set) or for x in nd.inputs]) and
isinstance(x, tensor.Constant) # we can do this because the assumption is that a
for x in nd.inputs]) and # viewOp or deepCopyOp will be just at the end of the
# we can do this because the assumption is that a # function and not somewhere in the middle ..
# viewOp or deepCopyOp will be just at the end of the not isinstance(nd.op, theano.compile.ViewOp) and
# function and not somewhere in the middle .. not isinstance(nd.op, theano.compile.DeepCopyOp)):
not isinstance(nd.op, theano.compile.ViewOp) and
not isinstance(nd.op, theano.compile.DeepCopyOp) and # We have a candidate node to removable
# and we didn't already looked at this node # Step 1. Reconstruct it on outside
nd not in to_remove_set): to_remove_add(nd)
outside_ins = []
# We have a candidate node to removable for x in nd.inputs:
# Step 1. Reconstruct it on outside if x in inner_non_seqs_set:
to_remove_add(nd) _idx = inner_non_seqs_map[x]
outside_ins = [] outside_ins.append(outer_non_seqs[_idx])
for x in nd.inputs: elif x in to_replace_set:
if x in inner_non_seqs_set: outside_ins.append(replace_with_out[to_replace_map[x]])
_idx = inner_non_seqs_map[x] elif isinstance(x, theano.Constant):
outside_ins.append(outer_non_seqs[_idx]) outside_ins.append(x.clone())
elif x in to_replace_set: else:
outside_ins.append(replace_with_out[to_replace_map[x]]) raise Exception(
elif isinstance(x, theano.Constant): ('Error in the `scan_pushout_non_seq_'
outside_ins.append(x.clone()) 'operations`. The optimization tries '
else: 'to move some computation fron scan '
raise Exception( 'which is not allowed to move. Report '
('Error in the `scan_pushout_non_seq_' 'this on theano-users list'), x)
'operations`. The optimization tries ' outside_ins = [x.type.filter_variable(y) for x, y in
'to move some computation fron scan ' zip(nd.inputs, outside_ins)]
'which is not allowed to move. Report '
'this on theano-users list'), x) # Do not call make_node for test_value
outside_ins = [x.type.filter_variable(y) for x, y in nw_outer_node = nd.op(*outside_ins,
zip(nd.inputs, outside_ins)] **dict(return_list=True))[0].owner
# Do not call make_node for test_value # Step 2. Create variables for replacements
nw_outer_node = nd.op(*outside_ins, for idx, y in enumerate(nd.outputs):
**dict(return_list=True))[0].owner y_place_holder = scan_utils.safe_new(y, '_replace')
nto_replace = add_to_replace(y, nto_replace)
# Step 2. Create variables for replacements replace_with_in.append(y_place_holder)
for idx, y in enumerate(nd.outputs): assert isinstance(y, type(nw_outer_node.outputs[idx]))
y_place_holder = scan_utils.safe_new(y, '_replace') replace_with_out.append(nw_outer_node.outputs[idx])
nto_replace = add_to_replace(y, nto_replace)
replace_with_in.append(y_place_holder)
assert isinstance(y, type(nw_outer_node.outputs[idx]))
replace_with_out.append(nw_outer_node.outputs[idx])
changed = True
if counts >= max_iterations:
raise Exception('Error in the `scan_pushout_non_seq_operations`.'
' The optimization exhausted the maximal number '
'of iterations allowed!')
# We need to check all candidate replacements and choose those that # We need to check all candidate replacements and choose those that
# make sense for us # make sense for us
# Step 1. which elements of `to_replace` are used by remaining # Step 1. which elements of `to_replace` are used by remaining
# components of the inner function # components of the inner function
clean_to_replace = [] clean_to_replace = []
...@@ -330,11 +331,11 @@ class PushOutNonSeqScan(gof.Optimizer): ...@@ -330,11 +331,11 @@ class PushOutNonSeqScan(gof.Optimizer):
to_keep_set = set(to_keep) to_keep_set = set(to_keep)
for out, idx in to_replace_map.items(): for out, idx in to_replace_map.items():
if (out in to_keep_set if (# If types are different, conversion Op will be inserted,
and out.owner not in existent_nodes_set # and it may trigger an infinite loop.
# If types are different, conversion Op will be inserted, replace_with_in[idx].type == out.type and
# and it may trigger an infinite loop. out in to_keep_set and
and replace_with_in[idx].type == out.type): out.owner not in existent_nodes_set):
clean_to_replace.append(out) clean_to_replace.append(out)
clean_replace_with_in.append(replace_with_in[idx]) clean_replace_with_in.append(replace_with_in[idx])
clean_replace_with_out.append(replace_with_out[idx]) clean_replace_with_out.append(replace_with_out[idx])
...@@ -356,6 +357,7 @@ class PushOutNonSeqScan(gof.Optimizer): ...@@ -356,6 +357,7 @@ class PushOutNonSeqScan(gof.Optimizer):
_op_outs = scan_utils.clone(clean_outputs, _op_outs = scan_utils.clone(clean_outputs,
replace=givens) replace=givens)
_op_ins = clean_inputs + nw_inner _op_ins = clean_inputs + nw_inner
op_ins, op_outs = scan_utils.reconstruct_graph(_op_ins, _op_outs) op_ins, op_outs = scan_utils.reconstruct_graph(_op_ins, _op_outs)
# Reconstruct node # Reconstruct node
...@@ -398,7 +400,10 @@ class PushOutNonSeqScan(gof.Optimizer): ...@@ -398,7 +400,10 @@ class PushOutNonSeqScan(gof.Optimizer):
# This is a global opt for historical reason # This is a global opt for historical reason
# It should be possible to change it to a local opt. # It should be possible to change it to a local opt.
class PushOutSeqScan(gof.Optimizer): class PushOutSeqScan(gof.Optimizer):
"""
A global optimizer for pushing out the input variables that are not being
used inside the scan and provided in the sequences.
"""
def __init__(self): def __init__(self):
gof.Optimizer.__init__(self) gof.Optimizer.__init__(self)
...@@ -412,8 +417,13 @@ class PushOutSeqScan(gof.Optimizer): ...@@ -412,8 +417,13 @@ class PushOutSeqScan(gof.Optimizer):
self.process_node(fgraph, node) self.process_node(fgraph, node)
def process_node(self, fgraph, node): def process_node(self, fgraph, node):
"""
Important note: This function uses set and dictionary data structure.
By default they are not ordered for efficiency reasons. Take care and make
sure of changing them to Ordered versions if you need to iterate over those
variables.
"""
# this flag tells if there was any change during the last iterations # this flag tells if there was any change during the last iterations
changed = True
clean_inputs, clean_outputs = scan_utils.reconstruct_graph( clean_inputs, clean_outputs = scan_utils.reconstruct_graph(
node.op.inputs, node.op.outputs) node.op.inputs, node.op.outputs)
...@@ -422,14 +432,11 @@ class PushOutSeqScan(gof.Optimizer): ...@@ -422,14 +432,11 @@ class PushOutSeqScan(gof.Optimizer):
local_fgraph_outs_set = set(local_fgraph.outputs) local_fgraph_outs_set = set(local_fgraph.outputs)
local_fgraph_outs_map = dict([(v,k) for k,v in enumerate(local_fgraph.outputs)]) local_fgraph_outs_map = dict([(v,k) for k,v in enumerate(local_fgraph.outputs)])
max_iterations = 2 * len(local_fgraph_topo) + 3
counts = 0
to_remove_set = set() to_remove_set = set()
to_remove_add = to_remove_set.add to_remove_add = to_remove_set.add
to_replace_set = set() to_replace_set = set()
to_replace_add = to_replace_set.add to_replace_add = to_replace_set.add
to_replace_map = {} to_replace_map = OrderedDict()
nto_replace = 0 nto_replace = 0
def add_to_replace(y, nto_replace): def add_to_replace(y, nto_replace):
...@@ -455,102 +462,90 @@ class PushOutSeqScan(gof.Optimizer): ...@@ -455,102 +462,90 @@ class PushOutSeqScan(gof.Optimizer):
assert len(inner_non_seqs) == len(outer_non_seqs) assert len(inner_non_seqs) == len(outer_non_seqs)
assert len(inner_seqs) == len(outer_seqs) assert len(inner_seqs) == len(outer_seqs)
while changed and counts < max_iterations: for nd in local_fgraph_topo:
counts += 1 if (nd not in to_remove_set and
changed = False all([(x in inner_non_seqs_set) or
(x.owner in to_remove_set) or
for nd in local_fgraph_topo: isinstance(x, tensor.Constant) or
if (isinstance(nd.op, theano.tensor.Elemwise) and (x in inner_seqs_set) for x in nd.inputs]) and
all([(x in inner_non_seqs_set) or isinstance(nd.op, theano.tensor.Elemwise)):
(x.owner in to_remove_set) or
isinstance(x, tensor.Constant) or to_remove_add(nd)
(x in inner_seqs_set) outside_ins = []
for x in nd.inputs]) and depends_on_seqs = False
nd not in to_remove_set):
to_remove_add(nd) for x in nd.inputs:
outside_ins = [] if x in inner_non_seqs_set:
depends_on_seqs = False _idx = inner_non_seqs_map[x]
outside_ins.append(outer_non_seqs[_idx])
for x in nd.inputs: elif x in inner_seqs_set:
if x in inner_non_seqs_set: outside_ins.append(outer_seqs[inner_seqs_map[x]])
_idx = inner_non_seqs_map[x] depends_on_seqs = True
outside_ins.append(outer_non_seqs[_idx])
elif x in inner_seqs_set:
outside_ins.append(outer_seqs[inner_seqs_map[x]])
depends_on_seqs = True
elif x in to_replace_set:
outside_ins.append(replace_with_out[
to_replace_map[x]])
depends_on_seqs = True
elif isinstance(x, theano.Constant):
outside_ins.append(x.clone())
else:
raise Exception(
('Error in the `scan_pushout_seq_'
'operations`. The optimization tries '
'to move some computation fron scan '
'which is not allowed to move. Report '
'this on theano-users list'), x)
if not depends_on_seqs:
# Removing this node from the inner graph of scan
# should be handled by the PushOutNonSeqScan
# optimization. The current optimization only tries
# to pull sequence-dependant computation out of
# scan.
continue
# Do not call make_node for test_value
nw_outer_node = nd.op(*outside_ins,
**dict(return_list=True))[0].owner
# Step 2. Create variables for replacements
for idx, y in enumerate(nd.outputs):
y_place_holder = scan_utils.safe_new(y, '_replace')
nto_replace = add_to_replace(y, nto_replace)
replace_with_in.append(y_place_holder)
replace_with_out.append(nw_outer_node.outputs[idx])
changed = True
elif (isinstance(nd.op, theano.tensor.DimShuffle) and
(nd.inputs[0] in inner_seqs_set or
nd.inputs[0].owner in to_remove_set) and
not nd in to_remove_set):
to_remove_add(nd)
x = nd.inputs[0]
if x in inner_seqs_set:
outside_ins = outer_seqs[inner_seqs_map[x]]
elif x in to_replace_set: elif x in to_replace_set:
outside_ins = replace_with_out[to_replace_map[x]] outside_ins.append(replace_with_out[
new_ord = (0,) to_replace_map[x]])
for old_ord in nd.op.new_order: depends_on_seqs = True
if (old_ord == 'x'): elif isinstance(x, theano.Constant):
new_ord += (old_ord,) outside_ins.append(x.clone())
else: else:
new_ord += (old_ord + 1,) raise Exception(
new_outer = outside_ins.dimshuffle(new_ord) ('Error in the `scan_pushout_seq_'
y = nd.outputs[0] 'operations`. The optimization tries '
'to move some computation fron scan '
'which is not allowed to move. Report '
'this on theano-users list'), x)
if not depends_on_seqs:
# Removing this node from the inner graph of scan
# should be handled by the PushOutNonSeqScan
# optimization. The current optimization only tries
# to pull sequence-dependant computation out of
# scan.
continue
# Do not call make_node for test_value
nw_outer_node = nd.op(*outside_ins,
**dict(return_list=True))[0].owner
# Step 2. Create variables for replacements
for idx, y in enumerate(nd.outputs):
y_place_holder = scan_utils.safe_new(y, '_replace') y_place_holder = scan_utils.safe_new(y, '_replace')
nto_replace = add_to_replace(y, nto_replace) nto_replace = add_to_replace(y, nto_replace)
replace_with_in.append(y_place_holder) replace_with_in.append(y_place_holder)
replace_with_out.append(new_outer) replace_with_out.append(nw_outer_node.outputs[idx])
if hasattr(new_outer.tag, "test_value"): elif (nd not in to_remove_set and
new_sh = new_outer.tag.test_value.shape isinstance(nd.op, theano.tensor.DimShuffle) and
ref_sh = (outside_ins.tag.test_value.shape[0],) (nd.inputs[0] in inner_seqs_set or
ref_sh += nd.outputs[0].tag.test_value.shape nd.inputs[0].owner in to_remove_set)):
assert new_sh == ref_sh
to_remove_add(nd)
changed = True x = nd.inputs[0]
if x in inner_seqs_set:
outside_ins = outer_seqs[inner_seqs_map[x]]
elif x in to_replace_set:
outside_ins = replace_with_out[to_replace_map[x]]
new_ord = (0,)
for old_ord in nd.op.new_order:
if (old_ord == 'x'):
new_ord += (old_ord,)
else:
new_ord += (old_ord + 1,)
new_outer = outside_ins.dimshuffle(new_ord)
y = nd.outputs[0]
y_place_holder = scan_utils.safe_new(y, '_replace')
nto_replace = add_to_replace(y, nto_replace)
replace_with_in.append(y_place_holder)
replace_with_out.append(new_outer)
if hasattr(new_outer.tag, "test_value"):
new_sh = new_outer.tag.test_value.shape
ref_sh = (outside_ins.tag.test_value.shape[0],)
ref_sh += nd.outputs[0].tag.test_value.shape
assert new_sh == ref_sh
if counts >= max_iterations:
raise Exception('Error in the `scan_pushout_seq_operations`.'
' The optimization exhausted the maximal number '
'of iterations allowed!')
# We need to check all candidate replacements and choose those that # We need to check all candidate replacements and choose those that
# make sense for us # make sense for us
# Step 1. which elements of `to_replace` are used by remaining # Step 1. which elements of `to_replace` are used by remaining
# components of the inner function # components of the inner function
clean_to_replace = [] clean_to_replace = []
...@@ -565,10 +560,10 @@ class PushOutSeqScan(gof.Optimizer): ...@@ -565,10 +560,10 @@ class PushOutSeqScan(gof.Optimizer):
for out, idx in to_replace_map.items(): for out, idx in to_replace_map.items():
if (out in to_keep_set if (out in to_keep_set
and out.owner not in existent_nodes_set and out.owner not in existent_nodes_set
# If types are different, conversion Op will be inserted, # If types are different, conversion Op will be inserted,
# and it may trigger an infinite loop. # and it may trigger an infinite loop.
and replace_with_in[idx].type == out.type): and replace_with_in[idx].type == out.type):
clean_to_replace.append(out) clean_to_replace.append(out)
clean_replace_with_in.append(replace_with_in[idx]) clean_replace_with_in.append(replace_with_in[idx])
...@@ -684,10 +679,8 @@ class PushOutScanOutput(gof.Optimizer): ...@@ -684,10 +679,8 @@ class PushOutScanOutput(gof.Optimizer):
new_scan_node = None new_scan_node = None
local_fgraph_topo = local_fgraph.toposort() local_fgraph_topo = local_fgraph.toposort()
for nd in local_fgraph_topo: for nd in local_fgraph_topo:
if (isinstance(nd.op, theano.tensor.Dot) and if (isinstance(nd.op, theano.tensor.Dot) and
nd.out in args.inner_out_nit_sot): nd.out in args.inner_out_nit_sot):
""" """
The following optimization involves pushing out, after the The following optimization involves pushing out, after the
scan, a Dot whose output is nitsot (not feed back to the inner scan, a Dot whose output is nitsot (not feed back to the inner
...@@ -857,9 +850,7 @@ class PushOutScanOutput(gof.Optimizer): ...@@ -857,9 +850,7 @@ class PushOutScanOutput(gof.Optimizer):
outer_var = scan_args.outer_out_sit_sot[idx] outer_var = scan_args.outer_out_sit_sot[idx]
if len(outer_var.clients) == 1: if len(outer_var.clients) == 1:
client = outer_var.clients[0][0] client = outer_var.clients[0][0]
if (client != 'output' and if (client != 'output' and
isinstance(client.op, theano.tensor.Subtensor)): isinstance(client.op, theano.tensor.Subtensor)):
lst = theano.tensor.subtensor.get_idx_list( lst = theano.tensor.subtensor.get_idx_list(
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论