karolmajek/maxdeeplab
0
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"""Tests for deeplabv3plus."""17 18import numpy as np19import tensorflow as tf20 21from deeplab2 import common22from deeplab2 import config_pb223from deeplab2.model.decoder import deeplabv3plus24from deeplab2.utils import test_utils25 26 27def _create_deeplabv3plus_model(high_level_feature_name, low_level_feature_name,28 low_level_channels_project,29 aspp_output_channels, decoder_output_channels,30 atrous_rates, num_classes, **kwargs):31 decoder_options = config_pb2.DecoderOptions(32 feature_key=high_level_feature_name,33 decoder_channels=decoder_output_channels,34 aspp_channels=aspp_output_channels,35 atrous_rates=atrous_rates)36 deeplabv3plus_options = config_pb2.ModelOptions.DeeplabV3PlusOptions(37 low_level=config_pb2.LowLevelOptions(38 feature_key=low_level_feature_name,39 channels_project=low_level_channels_project),40 num_classes=num_classes)41 return deeplabv3plus.DeepLabV3Plus(decoder_options, deeplabv3plus_options,42 **kwargs)43 44 45class Deeplabv3PlusTest(tf.test.TestCase):46 47 def test_deeplabv3plus_feature_key_not_present(self):48 deeplabv3plus_decoder = _create_deeplabv3plus_model(49 high_level_feature_name='not_in_features_dict',50 low_level_feature_name='in_feature_dict',51 low_level_channels_project=128,52 aspp_output_channels=64,53 decoder_output_channels=64,54 atrous_rates=[6, 12, 18],55 num_classes=80)56 input_dict = dict()57 input_dict['in_feature_dict'] = tf.random.uniform(shape=(2, 65, 65, 32))58 59 with self.assertRaises(KeyError):60 _ = deeplabv3plus_decoder(input_dict)61 62 def test_deeplabv3plus_output_shape(self):63 list_of_num_classes = [2, 19, 133]64 for num_classes in list_of_num_classes:65 deeplabv3plus_decoder = _create_deeplabv3plus_model(66 high_level_feature_name='high',67 low_level_feature_name='low',68 low_level_channels_project=128,69 aspp_output_channels=64,70 decoder_output_channels=128,71 atrous_rates=[6, 12, 18],72 num_classes=num_classes)73 input_dict = dict()74 input_dict['high'] = tf.random.uniform(shape=(2, 65, 65, 32))75 input_dict['low'] = tf.random.uniform(shape=(2, 129, 129, 16))76 expected_shape = [2, 129, 129, num_classes]77 78 logit_tensor = deeplabv3plus_decoder(input_dict)79 self.assertListEqual(80 logit_tensor[common.PRED_SEMANTIC_LOGITS_KEY].shape.as_list(),81 expected_shape)82 83 def test_deeplabv3plus_feature_extraction_consistency(self):84 deeplabv3plus_decoder = _create_deeplabv3plus_model(85 high_level_feature_name='high',86 low_level_feature_name='low',87 low_level_channels_project=128,88 aspp_output_channels=96,89 decoder_output_channels=64,90 atrous_rates=[6, 12, 18],91 num_classes=80)92 input_dict = dict()93 input_dict['high'] = tf.random.uniform(shape=(2, 65, 65, 32))94 input_dict['low'] = tf.random.uniform(shape=(2, 129, 129, 16))95 96 reference_logits_tensor = deeplabv3plus_decoder(97 input_dict, training=False)98 logits_tensor_to_compare = deeplabv3plus_decoder(input_dict, training=False)99 100 np.testing.assert_equal(101 reference_logits_tensor[common.PRED_SEMANTIC_LOGITS_KEY].numpy(),102 logits_tensor_to_compare[common.PRED_SEMANTIC_LOGITS_KEY].numpy())103 104 def test_deeplabv3plus_pool_size_setter(self):105 deeplabv3plus_decoder = _create_deeplabv3plus_model(106 high_level_feature_name='high',107 low_level_feature_name='low',108 low_level_channels_project=128,109 aspp_output_channels=96,110 decoder_output_channels=64,111 atrous_rates=[6, 12, 18],112 num_classes=80)113 pool_size = (10, 10)114 deeplabv3plus_decoder.set_pool_size(pool_size)115 116 self.assertTupleEqual(deeplabv3plus_decoder._aspp._aspp_pool._pool_size,117 pool_size)118 119 @test_utils.test_all_strategies120 def test_deeplabv3plus_sync_bn(self, strategy):121 input_dict = dict()122 input_dict['high'] = tf.random.uniform(shape=(2, 65, 65, 32))123 input_dict['low'] = tf.random.uniform(shape=(2, 129, 129, 16))124 with strategy.scope():125 for bn_layer in test_utils.NORMALIZATION_LAYERS:126 deeplabv3plus_decoder = _create_deeplabv3plus_model(127 high_level_feature_name='high',128 low_level_feature_name='low',129 low_level_channels_project=128,130 aspp_output_channels=96,131 decoder_output_channels=64,132 atrous_rates=[6, 12, 18],133 num_classes=80,134 bn_layer=bn_layer)135 _ = deeplabv3plus_decoder(input_dict)136 137 def test_deeplabv3plus_pool_size_resetter(self):138 deeplabv3plus_decoder = _create_deeplabv3plus_model(139 high_level_feature_name='high',140 low_level_feature_name='low',141 low_level_channels_project=128,142 aspp_output_channels=96,143 decoder_output_channels=64,144 atrous_rates=[6, 12, 18],145 num_classes=80)146 pool_size = (None, None)147 deeplabv3plus_decoder.reset_pooling_layer()148 149 self.assertTupleEqual(deeplabv3plus_decoder._aspp._aspp_pool._pool_size,150 pool_size)151 152 def test_deeplabv3plus_ckpt_items(self):153 deeplabv3plus_decoder = _create_deeplabv3plus_model(154 high_level_feature_name='high',155 low_level_feature_name='low',156 low_level_channels_project=128,157 aspp_output_channels=96,158 decoder_output_channels=64,159 atrous_rates=[6, 12, 18],160 num_classes=80)161 ckpt_dict = deeplabv3plus_decoder.checkpoint_items162 self.assertIn(common.CKPT_DEEPLABV3PLUS_ASPP, ckpt_dict)163 self.assertIn(common.CKPT_DEEPLABV3PLUS_PROJECT_CONV_BN_ACT, ckpt_dict)164 self.assertIn(common.CKPT_DEEPLABV3PLUS_FUSE, ckpt_dict)165 self.assertIn(common.CKPT_SEMANTIC_LAST_LAYER, ckpt_dict)166 167 168if __name__ == '__main__':169 tf.test.main()170