better dims/rank checking
@@ -10,7 +10,9 @@ def check_spatial_dims(input_tensor, criteria):
criteria can be a lambda function
e.g. lambda x : x > 10 checks whether each dim is greater than 10
"""
- spatial_dims = input_tensor.get_shape()[1:-1].as_list()
+ input_shape = input_tensor.get_shape()
+ input_shape.with_rank_at_least(3)
+ spatial_dims = input_shape[1:-1].as_list()
all_dims_satisfied = np.all([criteria(x) for x in spatial_dims])
if not all_dims_satisfied:
import inspect
@@ -18,7 +20,6 @@ def check_spatial_dims(input_tensor, criteria):
" not compatible, please tune "
"the input window sizes: {}".format(
inspect.getsource(criteria)))
- else:
return all_dims_satisfied
@@ -26,10 +27,12 @@ def infer_spatial_rank(input_tensor):
"""
e.g. given an input tensor [Batch, X, Y, Z, Feature] the spatial rank is 3
"""
- dims = input_tensor.get_shape().ndims - 2
- assert dims > 0, "input tensor should have at least one spatial dim, " \
- "in addition to batch and channel dims"
- return dims
+ input_shape = input_tensor.get_shape()
+ input_shape.with_rank_at_least(3)
+ #dims = input_tensor.get_shape().ndims - 2
+ #assert dims > 0, "input tensor should have at least one spatial dim, " \
+ # "in addition to batch and channel dims"
+ return input_shape.ndims - 2
def trivial_kernel(kernel_shape):
