karolmajek/Axial-DeepLab-SWideRNet

1# coding=utf-82# Copyright 2021 The Deeplab2 Authors.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16"""Provide utility functions to write simple tests."""17import functools18 19import numpy as np20import tensorflow as tf21 22 23NORMALIZATION_LAYERS = (24    tf.keras.layers.experimental.SyncBatchNormalization,25    tf.keras.layers.BatchNormalization26)27 28 29def create_strategy():30  """Returns a strategy based on available devices.31 32  Does NOT work with local_multiworker_tpu_test tests!33  """34  tpus = tf.config.list_logical_devices(device_type='TPU')35  gpus = tf.config.list_logical_devices(device_type='GPU')36  if tpus:37    resolver = tf.distribute.cluster_resolver.TPUClusterResolver('')38    tf.config.experimental_connect_to_cluster(resolver)39    tf.tpu.experimental.initialize_tpu_system(resolver)40    return tf.distribute.TPUStrategy(resolver)41  elif gpus:42    return tf.distribute.OneDeviceStrategy('/gpu:0')43  else:44    return tf.distribute.OneDeviceStrategy('/cpu:0')45 46 47def test_all_strategies(func):48  """Decorator to test CPU, GPU and TPU strategies."""49  @functools.wraps(func)50  def decorator(self):51    strategy = create_strategy()52    return func(self, strategy)53  return decorator54 55 56def create_test_input(batch, height, width, channels):57  """Creates test input tensor."""58  return tf.convert_to_tensor(59      np.tile(60          np.reshape(61              np.reshape(np.arange(height), [height, 1]) +62              np.reshape(np.arange(width), [1, width]),63              [1, height, width, 1]),64          [batch, 1, 1, channels]), dtype=tf.float32)65