1# Copyright 2018 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"""Utilities for helping test ops.""" 16 17import numpy as np 18from six.moves import range 19 20 21def ConvertBetweenDataFormats(x, data_format_src, data_format_dst): 22 """Converts 4D tensor between data formats.""" 23 24 valid_data_formats = ["NHWC", "NCHW", "HWNC", "HWCN"] 25 if data_format_src not in valid_data_formats: 26 raise ValueError("data_format_src must be of %s, got %s." % 27 (valid_data_formats, data_format_src)) 28 if data_format_dst not in valid_data_formats: 29 raise ValueError("data_format_dst must be of %s, got %s." % 30 (valid_data_formats, data_format_dst)) 31 if len(x.shape) != 4: 32 raise ValueError("x must be 4D, got shape %s." % x.shape) 33 34 if data_format_src == data_format_dst: 35 return x 36 37 dim_map = {d: i for i, d in enumerate(data_format_src)} 38 transpose_dims = [dim_map[d] for d in data_format_dst] 39 return np.transpose(x, transpose_dims) 40 41 42def PermuteDimsBetweenDataFormats(dims, data_format_src, data_format_dst): 43 """Get new shape for converting between data formats.""" 44 45 valid_data_formats = ["NHWC", "NCHW", "HWNC", "HWCN"] 46 if data_format_src not in valid_data_formats: 47 raise ValueError("data_format_src must be of %s, got %s." % 48 (valid_data_formats, data_format_src)) 49 if data_format_dst not in valid_data_formats: 50 raise ValueError("data_format_dst must be of %s, got %s." % 51 (valid_data_formats, data_format_dst)) 52 if len(dims) != 4: 53 raise ValueError("dims must be of length 4, got %s." % dims) 54 55 if data_format_src == data_format_dst: 56 return dims 57 58 dim_map = {d: i for i, d in enumerate(data_format_src)} 59 permuted_dims = [dims[dim_map[d]] for d in data_format_dst] 60 return permuted_dims 61 62 63_JIT_WARMUP_ITERATIONS = 10 64 65 66def RunWithWarmup(sess, op_to_run, feed_dict, options=None, run_metadata=None): 67 """Runs a graph a few times to ensure that its clusters are compiled.""" 68 for _ in range(0, _JIT_WARMUP_ITERATIONS): 69 sess.run(op_to_run, feed_dict, options=options) 70 return sess.run( 71 op_to_run, feed_dict, options=options, run_metadata=run_metadata) 72