# coding=utf-8 # Copyright 2022 The Deeplab2 Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """MOAT: This file contains the implementation of MOAT [1]. [1] MOAT: Alternating Mobile Convolution and Attention Brings Strong Vision Models, arXiv: 2210.01820. Chenglin Yang, Siyuan Qiao, Qihang Yu, Xiaoding Yuan, Yukun Zhu, Alan Yuille, Hartwig Adam, Liang-Chieh Chen. """ import copy import re from typing import Optional import tensorflow as tf from absl import logging from deeplab2.model.layers.moat_blocks import MBConvBlock from deeplab2.model.layers.moat_blocks import MOATBlock from deeplab2.utils.hparam_configs import Config from deeplab2.utils.hparam_configs import create_config_from_dict # This handles the invalid name scope of tf.keras.sequential # used in stem layers. _STEM_LAYER_NAME_SCOPE = 'moat/stem/' # This is for loading the exponential moving average of variables # in the checkpoint. _EMA_VARIABLE_NAME_POSTFIX = '/ExponentialMovingAverage' # The position embedding size at stride 16 and 32. # The input size changes, but the number of learnable parameters of position # embedding does not change. The position embeddings are interpolated for # different input sizes. _STRIDE_16_POSITION_EMBEDDING_SIZE = 14 _STRIDE_32_POSITION_EMBEDDING_SIZE = 7 class MOAT(tf.keras.Model): """MOAT backbone.""" def _retrieve_config(self, config): """Retrieves the config of MOAT. Args: config: A dictionary containing the following keys. -stem_size: A list of integers, the list length is the number of stem layers, the list element is the output channels of the stem layer. -block_type: A list of strings, the list length is the stage number. The list element is the block type for a stage. The supported block types are 'mbconv' and 'moat'. -num_blocks: A list of integers, the list length is the stage number. The list element is the number of blocks in a stage. -hidden_size: A list of integers, the list length is the stage number. The list element is the output channels of the blocks in a stage. -stage_stride: A list of intergers, the stride of first block in each stage. -expansion_rate: An integer, expansion rate of MBConv. -se_ratio: An integer, expansion ratio of SE in MBConv block. -head_size: An integer, feature channels per head in Attention. -window_size: A list of list of integers, specifying the window size for each stage. The length should be the number of stages. For example: [None, None, [14, 16], [7, 8]] means the last two stages uses 14 by 16 and 7 by 8 windows, respectively. -position_embedding_size: A list of integers, specifying the position embedding sizes for all stages. If the feature maps have larger sizes, the position embeddings will be interpolated to match them. -use_checkpointing_for_attention: A boolean, specifying whether to use checkpointing for attention. -global_attention_at_end_of_moat_stage: A boolean, specifying whether to use global attention for the last block if the stage consists of moat blocks. -relative_position_embedding_type: A string, type of relative position embedding in Attention. Currently, only '2d_multi_head' is supported. If None, no relative position embedding will be used. -ln_epsilon: A float, epsilon for layer normalization in Attention. -pool_size: An integer, kernel size for pooling in shortcut branch in MBConv block and MOAT block. -survival_prob: A float, 1 - drop_path_rate. -kernel_initializer: Initializer for the kernel weights matrix. -bias_initializer: Initializer for the bias vector. -name: A string, model name. -build_classification_head_with_class_num: An integer, number of classes. If None, no classification head will be built. Returns: A Config class: hparams_config.Config. Raises: ValueError: If the lengths of block_type, num_blocks and hidden_size are not the same. ValueError: If the element of block_type is not one of ['mbconv', 'moat']. """ required_keys = ['stem_size', 'block_type', 'num_blocks', 'hidden_size'] optional_keys = { 'stage_stride': [2, 2, 2, 2], 'expansion_rate': 4, 'se_ratio': 0.25, 'head_size': 32, 'window_size': [None, None, [14, 14], [7, 7]], 'position_embedding_size': [ None, None, _STRIDE_16_POSITION_EMBEDDING_SIZE, _STRIDE_32_POSITION_EMBEDDING_SIZE], 'use_checkpointing_for_attention': False, 'global_attention_at_end_of_moat_stage': False, 'relative_position_embedding_type': '2d_multi_head', 'ln_epsilon': 1e-5, 'pool_size': 2, 'survival_prob': None, 'kernel_initializer': tf.random_normal_initializer(stddev=0.02), 'bias_initializer': tf.zeros_initializer, 'build_classification_head_with_class_num': None, } config = create_config_from_dict(config, required_keys, optional_keys) stage_number = len(config.block_type) if (len(config.num_blocks) != stage_number or len(config.hidden_size) != stage_number): raise ValueError('The lengths of block_type, num_blocks and hidden_size ', 'should be the same.') return config def _local_config(self, config, idx, exclude_regex=None): """Gets stage-wise config from backbone-wise config.""" config = copy.deepcopy(config) for key in config.__dict__: if isinstance(config[key], (list, tuple)): if exclude_regex is None or not re.search(exclude_regex, key): config[key] = config[key][idx] return config def __init__(self, **config): super().__init__(name='moat') self._config = self._retrieve_config(config) def _build_stem(self): stem_layers = [] for i in range(len(self._config.stem_size)): conv_layer = tf.keras.layers.Conv2D( filters=self._config.stem_size[i], kernel_size=3, strides=2 if i == 0 else 1, padding='same', kernel_initializer=self._config.kernel_initializer, bias_initializer=self._config.bias_initializer, use_bias=True, name='conv_{}'.format(i)) stem_layers.append(conv_layer) if i < len(self._config.stem_size) - 1: stem_layers.append(self._config.norm_class(name='norm_{}'.format(i))) stem_layers.append(tf.keras.layers.Activation( self._config.activation, name='act_{}'.format(i))) # The name scope of tf.keras.Sequential is invalid, see error handling # in the part of loading checkpoints in function get_model. self._stem = tf.keras.Sequential( layers=stem_layers, name='stem') def _build_block(self, local_block_config): if local_block_config.block_type == 'mbconv': block = MBConvBlock(**local_block_config) elif local_block_config.block_type == 'moat': block = MOATBlock(**local_block_config) else: raise ValueError('Unsupported block_type: {}'.format( local_block_config.block_type)) return block def _build_classification_head(self): self._final_layer_norm = tf.keras.layers.LayerNormalization( epsilon=self._config.ln_epsilon, name='final_layer_norm') self._logits_head = tf.keras.layers.Conv2D( filters=self._config.build_classification_head_with_class_num, kernel_size=1, strides=1, kernel_initializer=self._config.kernel_initializer, bias_initializer=self._config.bias_initializer, padding='same', use_bias=True, name='logits_head') def _adjust_survival_rate(self, local_block_config, block_id, total_num_blocks): survival_prob = self._config.survival_prob if survival_prob is not None: drop_rate = 1.0 - survival_prob survival_prob = 1.0 - drop_rate * block_id / total_num_blocks local_block_config = local_block_config.replace( survival_prob=survival_prob) return local_block_config def build(self, input_shape) -> None: norm_class = tf.keras.layers.experimental.SyncBatchNormalization self._config.norm_class = norm_class self._config.activation = tf.nn.gelu self._build_stem() self._blocks = [] total_num_blocks = sum(self._config.num_blocks) for stage_id in range(len(self._config.block_type)): stage_config = self._local_config(self._config, stage_id, '^stem.*') stage_blocks = [] for local_block_id in range(stage_config.num_blocks): local_block_config = copy.deepcopy(stage_config) block_stride = 1 if local_block_id == 0: block_stride = self._config.stage_stride[stage_id] local_block_config = local_block_config.replace( block_stride=block_stride) block_id = sum(self._config.num_blocks[:stage_id]) + local_block_id local_block_config = self._adjust_survival_rate( local_block_config, block_id, total_num_blocks) block_name = 'block_{:0>2d}_{:0>2d}'.format(stage_id, local_block_id) local_block_config.name = block_name if (local_block_id == stage_config.num_blocks - 1 and self._config.block_type[stage_id] == 'moat' and self._config.global_attention_at_end_of_moat_stage): local_block_config.window_size = None block = self._build_block(local_block_config) stage_blocks.append(block) self._blocks.append(stage_blocks) if self._config.build_classification_head_with_class_num is not None: self._build_classification_head() def call(self, inputs, training=False, mask=None): endpoints = {} output = self._stem(inputs, training=training) endpoints['stage1'] = output endpoints['res1'] = self._config.activation(output) for stage_id, stage_blocks in enumerate(self._blocks): for block in stage_blocks: output = block(output, training=training) endpoints['stage{}'.format(stage_id + 2)] = output endpoints['res{}'.format(stage_id + 2)] = self._config.activation(output) if self._config.build_classification_head_with_class_num is None: return endpoints else: reduce_axes = list(range(1, output.shape.rank - 1)) output = tf.reduce_mean(output, axis=reduce_axes, keepdims=True) output = self._final_layer_norm(output) output = self._logits_head(output, training=training) logits = tf.squeeze(output, axis=[1, 2]) return logits moat0_config = Config( stem_size=[64, 64], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 3, 7, 2], hidden_size=[96, 192, 384, 768], ) moat1_config = Config( stem_size=[64, 64], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 6, 14, 2], hidden_size=[96, 192, 384, 768], ) moat2_config = Config( stem_size=[128, 128], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 6, 14, 2], hidden_size=[128, 256, 512, 1024], ) moat3_config = Config( stem_size=[160, 160], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 12, 28, 2], hidden_size=[160, 320, 640, 1280], ) moat4_config = Config( stem_size=[256, 256], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 12, 28, 2], hidden_size=[256, 512, 1024, 2048], ) tiny_moat0_config = Config( stem_size=[32, 32], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 3, 7, 2], hidden_size=[32, 64, 128, 256], ) tiny_moat1_config = Config( stem_size=[40, 40], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 3, 7, 2], hidden_size=[40, 80, 160, 320], ) tiny_moat2_config = Config( stem_size=[56, 56], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 3, 7, 2], hidden_size=[56, 112, 224, 448], ) tiny_moat3_config = Config( stem_size=[80, 80], block_type=['mbconv', 'mbconv', 'moat', 'moat'], num_blocks=[2, 3, 7, 2], hidden_size=[80, 160, 320, 640], ) no_relative_pe = Config( relative_position_embedding_type=None, ) def get_model( name: str, input_shape, window_size=None, survival_rate: Optional[float] = None, pool_size: Optional[int] = 3, override_config=None, pretrained_weights_path: Optional[str] = None, remove_position_embedding: Optional[bool] = None, return_config: Optional[bool] = False, strict_loading: Optional[bool] = False, use_checkpointing_for_attention: Optional[bool] = False, global_attention_at_end_of_moat_stage: Optional[bool] = False, ) -> ...: """Gets MOAT according to name. Args: name: A string, the model name. It must be the key of the global dictionary _CHECKPOINTS_PATH. input_shape: A list of integers with length 3, the input shape to initialize the model. window_size: A list of lists of integers, specifying the window size for each stage. The length should be the number of stages. For example: [None, None, [14, 16], [7, 8]] means the last two stages uses 14 by 16 and 7 by 8 windows, respectively. Global attention will be used if it is None. survival_rate: A float, this equals 1 - drop path rate. pool_size: An integer, kernel size for pooling in shortcut branch. For classification, pool size 2x2 saves model flops. For downstream tasks, pool size 3x3 is preferred for better feature alignments. override_config: A dictionary overriding the model config. pretrained_weights_path: A string, specifying the path to load the pretrained checkpoint. If None, no pretrained checkpoint will be loaded. remove_position_embedding: A boolean, specifying whether to remove position embedding. If None, the usage of position embedding will depend on the model name. This flag can enable the model that is pretrained with position embedding to be finetuned without position embedding in downstream tasks. return_config: A boolean, whether to return config or model. strict_loading: A boolean, if True, all variables in the checkpoint must be found in the model. use_checkpointing_for_attention: A boolean, specifying whether to use checkpointing for attention. global_attention_at_end_of_moat_stage: A boolean, specifying whether to use global attention for the last block if the stage consists of moat blocks. Returns: config: Config to build moat. This is returned when return_config is True. moat: moat model. This is returned when return_config is False. Raises: ValueError: If the length of input_shape does not equal 3. """ if len(input_shape) != 3: raise ValueError('The input shape should be a list of length 3.') config = _get_moat_config_via_model_name(name) config.window_size = window_size config.pool_size = pool_size config.use_checkpointing_for_attention = use_checkpointing_for_attention config.global_attention_at_end_of_moat_stage = ( global_attention_at_end_of_moat_stage) if survival_rate is not None: config.survival_prob = survival_rate if override_config is not None: config.update(override_config) if remove_position_embedding: config.update(no_relative_pe) if return_config: return config moat = MOAT(**config) moat(tf.keras.Input(shape=input_shape)) if pretrained_weights_path: moat = _load_moat_pretrained_checkpoint( moat, pretrained_weights_path, strict_loading) return moat def _get_moat_config_via_model_name(name): """Gets the moat config according to name.""" if name in ['moat0_pretrain_224_1k', 'moat0_pretrain_224_no_pe_1k']: config = copy.deepcopy(moat0_config) config.survival_prob = 0.8 elif name in ['moat1_pretrain_224_1k', 'moat1_pretrain_224_no_pe_1k']: config = copy.deepcopy(moat1_config) config.survival_prob = 0.7 elif name in ['moat2_finetune_384_22k', 'moat2_finetune_384_no_pe_22k']: config = copy.deepcopy(moat2_config) config.survival_prob = 0.7 elif name in ['moat3_finetune_512_22k', 'moat3_finetune_512_no_pe_22k']: config = copy.deepcopy(moat3_config) config.survival_prob = 0.4 elif name in ['moat4_finetune_512_22k', 'moat4_finetune_512_no_pe_22k']: config = copy.deepcopy(moat4_config) config.survival_prob = 0.3 elif name in ['tiny_moat0_pretrain_256_1k', 'tiny_moat0_pretrain_256_no_pe_1k']: config = copy.deepcopy(tiny_moat0_config) config.survival_prob = 1.0 elif name in ['tiny_moat1_pretrain_256_1k', 'tiny_moat1_pretrain_256_no_pe_1k']: config = copy.deepcopy(tiny_moat1_config) config.survival_prob = 1.0 elif name in ['tiny_moat2_pretrain_256_1k', 'tiny_moat2_pretrain_256_no_pe_1k']: config = copy.deepcopy(tiny_moat2_config) config.survival_prob = 1.0 elif name in ['tiny_moat3_pretrain_256_1k', 'tiny_moat3_pretrain_256_no_pe_1k']: config = copy.deepcopy(tiny_moat3_config) config.survival_prob = 0.9 else: raise ValueError('Not supported moat name %s' % name) if 'no_pe' in name: config.update(no_relative_pe) return config def _load_moat_pretrained_checkpoint( moat, path, strict_loading=False): """Loads the TF1 pretrained weights for MOAT.""" checkpoint_reader = tf.train.load_checkpoint(path) variable_to_shape_map = checkpoint_reader.get_variable_to_shape_map() model_var_name = sorted([var.name for var in moat.trainable_variables]) ckpt_var_name = list(sorted(variable_to_shape_map.keys())) # This for loop ensures all moat variables can be found in the checkpoint. for var in moat.trainable_variables: name_to_find = var.name # We change the TF2 variable name for loading TF1 checkpoint. name_to_find = name_to_find.replace(':0', '') if name_to_find not in ckpt_var_name: if _STEM_LAYER_NAME_SCOPE + name_to_find in ckpt_var_name: # This handles the invalid name scope of tf.keras.sequential. logging.info('Add name scope %s to the variable: %s', _STEM_LAYER_NAME_SCOPE, name_to_find) name_to_find = _STEM_LAYER_NAME_SCOPE + name_to_find else: raise ValueError('Variable name %s is not found in the checkpoint' % var.name) if var.shape != variable_to_shape_map[name_to_find]: raise ValueError('Inconsistent shape for var: %s ' % var.name, 'should be %s' % var.shape, 'but found %s' % variable_to_shape_map[name_to_find]) var_mean_before_loading = tf.reduce_mean(var) var.assign(tf.Variable(checkpoint_reader.get_tensor(name_to_find))) var_mean_after_loading = tf.reduce_mean(var) if tf.math.equal(var_mean_before_loading, var_mean_after_loading).numpy(): raise ValueError('Loading error for the variable: %s' % name_to_find) logging.info('All variables in moat are found in the checkpoint.') # This for loop finds all checkpoint variables that are not found in moat. for var_name in variable_to_shape_map.keys(): name_in_ckpt = var_name.replace(_EMA_VARIABLE_NAME_POSTFIX, '') name_in_ckpt = name_in_ckpt + ':0' if (name_in_ckpt not in model_var_name and name_in_ckpt.replace(_STEM_LAYER_NAME_SCOPE, '') not in model_var_name): # This handles the invalid name scope of tf.keras.sequential. if strict_loading: raise ValueError('In ckpt but not in model_var_name', name_in_ckpt) else: logging.info('In ckpt but not in model_var_name: %s', name_in_ckpt) if strict_loading: logging.info('All variables in checkpoint are found in moat.') return moat