static_shape.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # ==============================================================================
  15. """Helper functions to access TensorShape values.
  16. The rank 4 tensor_shape must be of the form [batch_size, height, width, depth].
  17. """
  18. def get_batch_size(tensor_shape):
  19. """Returns batch size from the tensor shape.
  20. Args:
  21. tensor_shape: A rank 4 TensorShape.
  22. Returns:
  23. An integer representing the batch size of the tensor.
  24. """
  25. tensor_shape.assert_has_rank(rank=4)
  26. return tensor_shape[0].value
  27. def get_height(tensor_shape):
  28. """Returns height from the tensor shape.
  29. Args:
  30. tensor_shape: A rank 4 TensorShape.
  31. Returns:
  32. An integer representing the height of the tensor.
  33. """
  34. tensor_shape.assert_has_rank(rank=4)
  35. return tensor_shape[1].value
  36. def get_width(tensor_shape):
  37. """Returns width from the tensor shape.
  38. Args:
  39. tensor_shape: A rank 4 TensorShape.
  40. Returns:
  41. An integer representing the width of the tensor.
  42. """
  43. tensor_shape.assert_has_rank(rank=4)
  44. return tensor_shape[2].value
  45. def get_depth(tensor_shape):
  46. """Returns depth from the tensor shape.
  47. Args:
  48. tensor_shape: A rank 4 TensorShape.
  49. Returns:
  50. An integer representing the depth of the tensor.
  51. """
  52. tensor_shape.assert_has_rank(rank=4)
  53. return tensor_shape[3].value