提交 0a7a5538 authored 作者: Matthew Rocklin's avatar Matthew Rocklin

Added shape_of_variables function

Given an env and the shape of all inputs returns a dictionary holding the shapes of all intermediate and output variables
上级 3775cef2
import numpy
from theano.tensor.utils import hash_from_ndarray, hash_from_dict
import theano
from theano.tensor.utils import (hash_from_ndarray, hash_from_dict,
shape_of_variables)
def test_hash_from_ndarray():
......@@ -46,3 +48,23 @@ def test_hash_from_dict():
# List are not hashable. So they are transformed into tuple.
assert hash_from_dict({0: (0,)}) == hash_from_dict({0: [0]})
def test_shape_of_variables_simple():
x = theano.tensor.matrix('x')
y = x+x
env = theano.Env([x], [y])
assert shape_of_variables(env, {x: (5, 5)}) == {x: (5, 5), y: (5, 5)}
x = theano.tensor.matrix('x')
y = theano.tensor.dot(x, x.T)
env = theano.Env([x], [y])
shapes = shape_of_variables(env, {x: (5, 1)})
assert shapes[x] == (5, 1)
assert shapes[y] == (5, 5)
def test_shape_of_variables_subtensor():
x = theano.tensor.matrix('x')
subx = x[1:]
env = theano.Env([x], [subx])
shapes = shape_of_variables(env, {x: (10, 10)})
assert shapes[subx] == (9, 10)
import numpy
import theano
from theano.gof.cc import hash_from_code
......@@ -43,3 +44,45 @@ def hash_from_dict(d):
second_part += [v]
tuple_items = tuple(first_part + second_part)
return hash(tuple_items)
def shape_of_variables(env, input_shapes):
"""
Compute the numeric shape of all intermediate variables given input shapes
Inputs:
env - the theano.Env in question
input_shapes - a dict mapping input to shape
Outputs:
shapes - a dict mapping variable to shape
WARNING : This modifies the env. Not pure.
>>> import theano
>>> x = theano.tensor.matrix('x')
>>> y = x[512:]; y.name = 'y'
>>> env = theano.Env([x], [y])
>>> shape_of_variables(env, {x: (1024, 1024)})
{y: (512, 1024), x: (1024, 1024)}
"""
if not hasattr(env, 'shape_feature'):
env.extend(theano.tensor.opt.ShapeFeature())
input_dims = [dimension for inp in env.inputs
for dimension in env.shape_feature.shape_of[inp]]
output_dims = [dimension for shape in env.shape_feature.shape_of.values()
for dimension in shape]
compute_shapes = theano.function(input_dims, output_dims)
numeric_input_dims = [dim for inp in env.inputs
for dim in input_shapes[inp]]
numeric_output_dims = compute_shapes(*numeric_input_dims)
sym_to_num_dict = dict(zip(output_dims, numeric_output_dims))
return {var: tuple(sym_to_num_dict[sym]
for sym in env.shape_feature.shape_of[var])
for var in env.shape_feature.shape_of}
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论