karolmajek/maxdeeplab

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 deeplabv3."""17 18import numpy as np19import tensorflow as tf20 21from deeplab2 import common22from deeplab2 import config_pb223from deeplab2.model.decoder import deeplabv324from deeplab2.utils import test_utils25 26 27def _create_deeplabv3_model(feature_key, decoder_channels, aspp_channels,28                            atrous_rates, num_classes, **kwargs):29  decoder_options = config_pb2.DecoderOptions(30      feature_key=feature_key,31      decoder_channels=decoder_channels,32      aspp_channels=aspp_channels,33      atrous_rates=atrous_rates)34  deeplabv3_options = config_pb2.ModelOptions.DeeplabV3Options(35      num_classes=num_classes)36  return deeplabv3.DeepLabV3(decoder_options, deeplabv3_options, **kwargs)37 38 39class Deeplabv3Test(tf.test.TestCase):40 41  def test_deeplabv3_feature_key_not_present(self):42    deeplabv3_decoder = _create_deeplabv3_model(43        feature_key='not_in_features_dict',44        aspp_channels=64,45        decoder_channels=48,46        atrous_rates=[6, 12, 18],47        num_classes=80)48    input_dict = dict()49    input_dict['not_the_same_key'] = tf.random.uniform(shape=(2, 65, 65, 32))50 51    with self.assertRaises(KeyError):52      _ = deeplabv3_decoder(input_dict)53 54  def test_deeplabv3_output_shape(self):55    list_of_num_classes = [2, 19, 133]56    for num_classes in list_of_num_classes:57      deeplabv3_decoder = _create_deeplabv3_model(58          feature_key='not_used',59          aspp_channels=64,60          decoder_channels=48,61          atrous_rates=[6, 12, 18],62          num_classes=num_classes)63      input_tensor = tf.random.uniform(shape=(2, 65, 65, 32))64      expected_shape = [2, 65, 65, num_classes]65 66      logit_tensor = deeplabv3_decoder(input_tensor)67      self.assertListEqual(68          logit_tensor[common.PRED_SEMANTIC_LOGITS_KEY].shape.as_list(),69          expected_shape)70 71  @test_utils.test_all_strategies72  def test_sync_bn(self, strategy):73    input_tensor = tf.random.uniform(shape=(2, 65, 65, 32))74    with strategy.scope():75      for bn_layer in test_utils.NORMALIZATION_LAYERS:76        deeplabv3_decoder = _create_deeplabv3_model(77            feature_key='not_used',78            aspp_channels=64,79            decoder_channels=48,80            atrous_rates=[6, 12, 18],81            num_classes=19,82            bn_layer=bn_layer)83        _ = deeplabv3_decoder(input_tensor)84 85  def test_deeplabv3_feature_extraction_consistency(self):86    deeplabv3_decoder = _create_deeplabv3_model(87        aspp_channels=64,88        decoder_channels=48,89        atrous_rates=[6, 12, 18],90        num_classes=80,91        feature_key='feature_key')92    input_tensor = tf.random.uniform(shape=(2, 65, 65, 32))93    input_dict = dict()94    input_dict['feature_key'] = input_tensor95 96    reference_logits_tensor = deeplabv3_decoder(input_tensor, training=False)97    logits_tensor_to_compare = deeplabv3_decoder(input_dict, training=False)98 99    np.testing.assert_equal(100        reference_logits_tensor[common.PRED_SEMANTIC_LOGITS_KEY].numpy(),101        logits_tensor_to_compare[common.PRED_SEMANTIC_LOGITS_KEY].numpy())102 103  def test_deeplabv3_pool_size_setter(self):104    deeplabv3_decoder = _create_deeplabv3_model(105        feature_key='not_used',106        aspp_channels=64,107        decoder_channels=48,108        atrous_rates=[6, 12, 18],109        num_classes=80)110    pool_size = (10, 10)111    deeplabv3_decoder.set_pool_size(pool_size)112 113    self.assertTupleEqual(deeplabv3_decoder._aspp._aspp_pool._pool_size,114                          pool_size)115 116  def test_deeplabv3_pool_size_resetter(self):117    deeplabv3_decoder = _create_deeplabv3_model(118        feature_key='not_used',119        aspp_channels=64,120        decoder_channels=48,121        atrous_rates=[6, 12, 18],122        num_classes=80)123    pool_size = (None, None)124    deeplabv3_decoder.reset_pooling_layer()125 126    self.assertTupleEqual(deeplabv3_decoder._aspp._aspp_pool._pool_size,127                          pool_size)128 129  def test_deeplabv3_ckpt_items(self):130    deeplabv3_decoder = _create_deeplabv3_model(131        feature_key='not_used',132        aspp_channels=64,133        decoder_channels=48,134        atrous_rates=[6, 12, 18],135        num_classes=80)136    ckpt_dict = deeplabv3_decoder.checkpoint_items137    self.assertIn(common.CKPT_DEEPLABV3_ASPP, ckpt_dict)138    self.assertIn(common.CKPT_DEEPLABV3_CLASSIFIER_CONV_BN_ACT, ckpt_dict)139    self.assertIn(common.CKPT_SEMANTIC_LAST_LAYER, ckpt_dict)140 141 142if __name__ == '__main__':143  tf.test.main()144