# Copyright (c) OpenMMLab. All rights reserved. import sys import numpy as np import torch from mmcv.cnn.bricks.drop import build_dropout from mmcv.cnn.bricks.transformer import FFN, PatchEmbed from mmengine.model import BaseModule, ModuleList from mmengine.model.weight_init import trunc_normal_ from torch import nn from torch.autograd import Function as Function from mmpretrain.models.backbones.base_backbone import BaseBackbone from mmpretrain.registry import MODELS from ..utils import (MultiheadAttention, build_norm_layer, resize_pos_embed, to_2tuple) class RevBackProp(Function): """Custom Backpropagation function to allow (A) flushing memory in forward and (B) activation recomputation reversibly in backward for gradient calculation. Inspired by https://github.com/RobinBruegger/RevTorch/blob/master/revtorch/revtorch.py """ @staticmethod def forward( ctx, x, layers, buffer_layers, # List of layer ids for int activation to buffer ): """Reversible Forward pass. Any intermediate activations from `buffer_layers` are cached in ctx for forward pass. This is not necessary for standard usecases. Each reversible layer implements its own forward pass logic. """ buffer_layers.sort() x1, x2 = torch.chunk(x, 2, dim=-1) intermediate = [] for layer in layers: x1, x2 = layer(x1, x2) if layer.layer_id in buffer_layers: intermediate.extend([x1.detach(), x2.detach()]) if len(buffer_layers) == 0: all_tensors = [x1.detach(), x2.detach()] else: intermediate = [torch.LongTensor(buffer_layers), *intermediate] all_tensors = [x1.detach(), x2.detach(), *intermediate] ctx.save_for_backward(*all_tensors) ctx.layers = layers return torch.cat([x1, x2], dim=-1) @staticmethod def backward(ctx, dx): """Reversible Backward pass. Any intermediate activations from `buffer_layers` are recovered from ctx. Each layer implements its own loic for backward pass (both activation recomputation and grad calculation). """ d_x1, d_x2 = torch.chunk(dx, 2, dim=-1) # retrieve params from ctx for backward x1, x2, *int_tensors = ctx.saved_tensors # no buffering if len(int_tensors) != 0: buffer_layers = int_tensors[0].tolist() else: buffer_layers = [] layers = ctx.layers for _, layer in enumerate(layers[::-1]): if layer.layer_id in buffer_layers: x1, x2, d_x1, d_x2 = layer.backward_pass( y1=int_tensors[buffer_layers.index(layer.layer_id) * 2 + 1], y2=int_tensors[buffer_layers.index(layer.layer_id) * 2 + 2], d_y1=d_x1, d_y2=d_x2, ) else: x1, x2, d_x1, d_x2 = layer.backward_pass( y1=x1, y2=x2, d_y1=d_x1, d_y2=d_x2, ) dx = torch.cat([d_x1, d_x2], dim=-1) del int_tensors del d_x1, d_x2, x1, x2 return dx, None, None class RevTransformerEncoderLayer(BaseModule): """Reversible Transformer Encoder Layer. This module is a building block of Reversible Transformer Encoder, which support backpropagation without storing activations. The residual connection is not applied to the FFN layer. Args: embed_dims (int): The feature dimension. num_heads (int): Parallel attention heads. feedforward_channels (int): The hidden dimension for FFNs. drop_rate (float): Probability of an element to be zeroed. Default: 0.0 attn_drop_rate (float): The drop out rate for attention layer. Default: 0.0 drop_path_rate (float): stochastic depth rate. Default 0.0 num_fcs (int): The number of linear in FFN Default: 2 qkv_bias (bool): enable bias for qkv if True. Default: True act_cfg (dict): The activation config for FFNs. Default: dict(type='GELU') norm_cfg (dict): Config dict for normalization layer. Default: dict(type='LN') layer_id (int): The layer id of current layer. Used in RevBackProp. Default: 0 init_cfg (dict or list[dict], optional): Initialization config dict. """ def __init__(self, embed_dims: int, num_heads: int, feedforward_channels: int, drop_rate: float = 0., attn_drop_rate: float = 0., drop_path_rate: float = 0., num_fcs: int = 2, qkv_bias: bool = True, act_cfg: dict = dict(type='GELU'), norm_cfg: dict = dict(type='LN'), layer_id: int = 0, init_cfg=None): super(RevTransformerEncoderLayer, self).__init__(init_cfg=init_cfg) self.drop_path_cfg = dict(type='DropPath', drop_prob=drop_path_rate) self.embed_dims = embed_dims self.ln1 = build_norm_layer(norm_cfg, self.embed_dims) self.attn = MultiheadAttention( embed_dims=embed_dims, num_heads=num_heads, attn_drop=attn_drop_rate, proj_drop=drop_rate, qkv_bias=qkv_bias) self.ln2 = build_norm_layer(norm_cfg, self.embed_dims) self.ffn = FFN( embed_dims=embed_dims, feedforward_channels=feedforward_channels, num_fcs=num_fcs, ffn_drop=drop_rate, act_cfg=act_cfg, add_identity=False) self.layer_id = layer_id self.seeds = {} def init_weights(self): super(RevTransformerEncoderLayer, self).init_weights() for m in self.ffn.modules(): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) nn.init.normal_(m.bias, std=1e-6) def seed_cuda(self, key): """Fix seeds to allow for stochastic elements such as dropout to be reproduced exactly in activation recomputation in the backward pass.""" # randomize seeds # use cuda generator if available if (hasattr(torch.cuda, 'default_generators') and len(torch.cuda.default_generators) > 0): # GPU device_idx = torch.cuda.current_device() seed = torch.cuda.default_generators[device_idx].seed() else: # CPU seed = int(torch.seed() % sys.maxsize) self.seeds[key] = seed torch.manual_seed(self.seeds[key]) def forward(self, x1, x2): """ Implementation of Reversible TransformerEncoderLayer ` x = x + self.attn(self.ln1(x)) x = self.ffn(self.ln2(x), identity=x) ` """ self.seed_cuda('attn') # attention output f_x2 = self.attn(self.ln1(x2)) # apply droppath on attention output self.seed_cuda('droppath') f_x2_dropped = build_dropout(self.drop_path_cfg)(f_x2) y1 = x1 + f_x2_dropped # free memory if self.training: del x1 # ffn output self.seed_cuda('ffn') g_y1 = self.ffn(self.ln2(y1)) # apply droppath on ffn output torch.manual_seed(self.seeds['droppath']) g_y1_dropped = build_dropout(self.drop_path_cfg)(g_y1) # final output y2 = x2 + g_y1_dropped # free memory if self.training: del x2 return y1, y2 def backward_pass(self, y1, y2, d_y1, d_y2): """Activation re-compute with the following equation. x2 = y2 - g(y1), g = FFN x1 = y1 - f(x2), f = MSHA """ # temporarily record intermediate activation for G # and use them for gradient calculation of G with torch.enable_grad(): y1.requires_grad = True torch.manual_seed(self.seeds['ffn']) g_y1 = self.ffn(self.ln2(y1)) torch.manual_seed(self.seeds['droppath']) g_y1 = build_dropout(self.drop_path_cfg)(g_y1) g_y1.backward(d_y2, retain_graph=True) # activate recomputation is by design and not part of # the computation graph in forward pass with torch.no_grad(): x2 = y2 - g_y1 del g_y1 d_y1 = d_y1 + y1.grad y1.grad = None # record F activation and calculate gradients on F with torch.enable_grad(): x2.requires_grad = True torch.manual_seed(self.seeds['attn']) f_x2 = self.attn(self.ln1(x2)) torch.manual_seed(self.seeds['droppath']) f_x2 = build_dropout(self.drop_path_cfg)(f_x2) f_x2.backward(d_y1, retain_graph=True) # propagate reverse computed activations at the # start of the previous block with torch.no_grad(): x1 = y1 - f_x2 del f_x2, y1 d_y2 = d_y2 + x2.grad x2.grad = None x2 = x2.detach() return x1, x2, d_y1, d_y2 class TwoStreamFusion(nn.Module): """A general constructor for neural modules fusing two equal sized tensors in forward. Args: mode (str): The mode of fusion. Options are 'add', 'max', 'min', 'avg', 'concat'. """ def __init__(self, mode: str): super().__init__() self.mode = mode if mode == 'add': self.fuse_fn = lambda x: torch.stack(x).sum(dim=0) elif mode == 'max': self.fuse_fn = lambda x: torch.stack(x).max(dim=0).values elif mode == 'min': self.fuse_fn = lambda x: torch.stack(x).min(dim=0).values elif mode == 'avg': self.fuse_fn = lambda x: torch.stack(x).mean(dim=0) elif mode == 'concat': self.fuse_fn = lambda x: torch.cat(x, dim=-1) else: raise NotImplementedError def forward(self, x): # split the tensor into two halves in the channel dimension x = torch.chunk(x, 2, dim=2) return self.fuse_fn(x) @MODELS.register_module() class RevVisionTransformer(BaseBackbone): """Reversible Vision Transformer. A PyTorch implementation of : `Reversible Vision Transformers `_ # noqa: E501 Args: arch (str | dict): Vision Transformer architecture. If use string, choose from 'small', 'base', 'large', 'deit-tiny', 'deit-small' and 'deit-base'. If use dict, it should have below keys: - **embed_dims** (int): The dimensions of embedding. - **num_layers** (int): The number of transformer encoder layers. - **num_heads** (int): The number of heads in attention modules. - **feedforward_channels** (int): The hidden dimensions in feedforward modules. Defaults to 'base'. img_size (int | tuple): The expected input image shape. Because we support dynamic input shape, just set the argument to the most common input image shape. Defaults to 224. patch_size (int | tuple): The patch size in patch embedding. Defaults to 16. in_channels (int): The num of input channels. Defaults to 3. drop_rate (float): Probability of an element to be zeroed. Defaults to 0. drop_path_rate (float): stochastic depth rate. Defaults to 0. qkv_bias (bool): Whether to add bias for qkv in attention modules. Defaults to True. norm_cfg (dict): Config dict for normalization layer. Defaults to ``dict(type='LN')``. final_norm (bool): Whether to add a additional layer to normalize final feature map. Defaults to True. out_type (str): The type of output features. Please choose from - ``"cls_token"``: The class token tensor with shape (B, C). - ``"featmap"``: The feature map tensor from the patch tokens with shape (B, C, H, W). - ``"avg_featmap"``: The global averaged feature map tensor with shape (B, C). - ``"raw"``: The raw feature tensor includes patch tokens and class tokens with shape (B, L, C). Defaults to ``"avg_featmap"``. with_cls_token (bool): Whether concatenating class token into image tokens as transformer input. Defaults to False. frozen_stages (int): Stages to be frozen (stop grad and set eval mode). -1 means not freezing any parameters. Defaults to -1. interpolate_mode (str): Select the interpolate mode for position embeding vector resize. Defaults to "bicubic". patch_cfg (dict): Configs of patch embeding. Defaults to an empty dict. layer_cfgs (Sequence | dict): Configs of each transformer layer in encoder. Defaults to an empty dict. fusion_mode (str): The fusion mode of transformer layers. Defaults to 'concat'. no_custom_backward (bool): Whether to use custom backward. Defaults to False. init_cfg (dict, optional): Initialization config dict. Defaults to None. """ arch_zoo = { **dict.fromkeys( ['s', 'small'], { 'embed_dims': 768, 'num_layers': 8, 'num_heads': 8, 'feedforward_channels': 768 * 3, }), **dict.fromkeys( ['b', 'base'], { 'embed_dims': 768, 'num_layers': 12, 'num_heads': 12, 'feedforward_channels': 3072 }), **dict.fromkeys( ['l', 'large'], { 'embed_dims': 1024, 'num_layers': 24, 'num_heads': 16, 'feedforward_channels': 4096 }), **dict.fromkeys( ['h', 'huge'], { # The same as the implementation in MAE # 'embed_dims': 1280, 'num_layers': 32, 'num_heads': 16, 'feedforward_channels': 5120 }), **dict.fromkeys( ['deit-t', 'deit-tiny'], { 'embed_dims': 192, 'num_layers': 12, 'num_heads': 3, 'feedforward_channels': 192 * 4 }), **dict.fromkeys( ['deit-s', 'deit-small'], { 'embed_dims': 384, 'num_layers': 12, 'num_heads': 6, 'feedforward_channels': 384 * 4 }), **dict.fromkeys( ['deit-b', 'deit-base'], { 'embed_dims': 768, 'num_layers': 12, 'num_heads': 12, 'feedforward_channels': 768 * 4 }), } num_extra_tokens = 0 # The official RevViT doesn't have class token OUT_TYPES = {'raw', 'cls_token', 'featmap', 'avg_featmap'} def __init__(self, arch='base', img_size=224, patch_size=16, in_channels=3, drop_rate=0., drop_path_rate=0., qkv_bias=True, norm_cfg=dict(type='LN', eps=1e-6), final_norm=True, out_type='avg_featmap', with_cls_token=False, frozen_stages=-1, interpolate_mode='bicubic', patch_cfg=dict(), layer_cfgs=dict(), fusion_mode='concat', no_custom_backward=False, init_cfg=None): super(RevVisionTransformer, self).__init__(init_cfg) if isinstance(arch, str): arch = arch.lower() assert arch in set(self.arch_zoo), \ f'Arch {arch} is not in default archs {set(self.arch_zoo)}' self.arch_settings = self.arch_zoo[arch] else: essential_keys = { 'embed_dims', 'num_layers', 'num_heads', 'feedforward_channels' } assert isinstance(arch, dict) and essential_keys <= set(arch), \ f'Custom arch needs a dict with keys {essential_keys}' self.arch_settings = arch self.embed_dims = self.arch_settings['embed_dims'] self.num_layers = self.arch_settings['num_layers'] self.img_size = to_2tuple(img_size) self.no_custom_backward = no_custom_backward # Set patch embedding _patch_cfg = dict( in_channels=in_channels, input_size=img_size, embed_dims=self.embed_dims, conv_type='Conv2d', kernel_size=patch_size, stride=patch_size, ) _patch_cfg.update(patch_cfg) self.patch_embed = PatchEmbed(**_patch_cfg) self.patch_resolution = self.patch_embed.init_out_size num_patches = self.patch_resolution[0] * self.patch_resolution[1] # Set out type if out_type not in self.OUT_TYPES: raise ValueError(f'Unsupported `out_type` {out_type}, please ' f'choose from {self.OUT_TYPES}') self.out_type = out_type # Set cls token if with_cls_token: self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims)) self.num_extra_tokens = 1 elif out_type != 'cls_token': self.cls_token = None self.num_extra_tokens = 0 else: raise ValueError( 'with_cls_token must be True when `out_type="cls_token"`.') # Set position embedding self.interpolate_mode = interpolate_mode self.pos_embed = nn.Parameter( torch.zeros(1, num_patches + self.num_extra_tokens, self.embed_dims)) self._register_load_state_dict_pre_hook(self._prepare_pos_embed) self.drop_after_pos = nn.Dropout(p=drop_rate) # stochastic depth decay rule dpr = np.linspace(0, drop_path_rate, self.num_layers) self.layers = ModuleList() if isinstance(layer_cfgs, dict): layer_cfgs = [layer_cfgs] * self.num_layers for i in range(self.num_layers): _layer_cfg = dict( embed_dims=self.embed_dims, num_heads=self.arch_settings['num_heads'], feedforward_channels=self. arch_settings['feedforward_channels'], drop_rate=drop_rate, drop_path_rate=dpr[i], qkv_bias=qkv_bias, layer_id=i, norm_cfg=norm_cfg) _layer_cfg.update(layer_cfgs[i]) self.layers.append(RevTransformerEncoderLayer(**_layer_cfg)) # fusion operation for the final output self.fusion_layer = TwoStreamFusion(mode=fusion_mode) self.frozen_stages = frozen_stages self.final_norm = final_norm if final_norm: self.ln1 = build_norm_layer(norm_cfg, self.embed_dims * 2) # freeze stages only when self.frozen_stages > 0 if self.frozen_stages > 0: self._freeze_stages() def init_weights(self): super(RevVisionTransformer, self).init_weights() if not (isinstance(self.init_cfg, dict) and self.init_cfg['type'] == 'Pretrained'): trunc_normal_(self.pos_embed, std=0.02) def _prepare_pos_embed(self, state_dict, prefix, *args, **kwargs): name = prefix + 'pos_embed' if name not in state_dict.keys(): return ckpt_pos_embed_shape = state_dict[name].shape if self.pos_embed.shape != ckpt_pos_embed_shape: from mmengine.logging import MMLogger logger = MMLogger.get_current_instance() logger.info( f'Resize the pos_embed shape from {ckpt_pos_embed_shape} ' f'to {self.pos_embed.shape}.') ckpt_pos_embed_shape = to_2tuple( int(np.sqrt(ckpt_pos_embed_shape[1] - self.num_extra_tokens))) pos_embed_shape = self.patch_embed.init_out_size state_dict[name] = resize_pos_embed(state_dict[name], ckpt_pos_embed_shape, pos_embed_shape, self.interpolate_mode, self.num_extra_tokens) @staticmethod def resize_pos_embed(*args, **kwargs): """Interface for backward-compatibility.""" return resize_pos_embed(*args, **kwargs) def _freeze_stages(self): # freeze position embedding self.pos_embed.requires_grad = False # set dropout to eval model self.drop_after_pos.eval() # freeze patch embedding self.patch_embed.eval() for param in self.patch_embed.parameters(): param.requires_grad = False # freeze cls_token if self.cls_token is not None: self.cls_token.requires_grad = False # freeze layers for i in range(1, self.frozen_stages + 1): m = self.layers[i - 1] m.eval() for param in m.parameters(): param.requires_grad = False # freeze the last layer norm if self.frozen_stages == len(self.layers) and self.final_norm: self.ln1.eval() for param in self.ln1.parameters(): param.requires_grad = False def forward(self, x): B = x.shape[0] x, patch_resolution = self.patch_embed(x) if self.cls_token is not None: cls_token = self.cls_token.expand(B, -1, -1) x = torch.cat((cls_token, x), dim=1) x = x + resize_pos_embed( self.pos_embed, self.patch_resolution, patch_resolution, mode=self.interpolate_mode, num_extra_tokens=self.num_extra_tokens) x = self.drop_after_pos(x) x = torch.cat([x, x], dim=-1) # forward with different conditions if not self.training or self.no_custom_backward: # in eval/inference model executing_fn = RevVisionTransformer._forward_vanilla_bp else: # use custom backward when self.training=True. executing_fn = RevBackProp.apply x = executing_fn(x, self.layers, []) if self.final_norm: x = self.ln1(x) x = self.fusion_layer(x) return (self._format_output(x, patch_resolution), ) @staticmethod def _forward_vanilla_bp(hidden_state, layers, buffer=[]): """Using reversible layers without reversible backpropagation. Debugging purpose only. Activated with self.no_custom_backward """ # split into ffn state(ffn_out) and attention output(attn_out) ffn_out, attn_out = torch.chunk(hidden_state, 2, dim=-1) del hidden_state for _, layer in enumerate(layers): attn_out, ffn_out = layer(attn_out, ffn_out) return torch.cat([attn_out, ffn_out], dim=-1) def _format_output(self, x, hw): if self.out_type == 'raw': return x if self.out_type == 'cls_token': return x[:, 0] patch_token = x[:, self.num_extra_tokens:] if self.out_type == 'featmap': B = x.size(0) # (B, N, C) -> (B, H, W, C) -> (B, C, H, W) return patch_token.reshape(B, *hw, -1).permute(0, 3, 1, 2) if self.out_type == 'avg_featmap': return patch_token.mean(dim=1)