# Copyright (C) 2022-present Naver Corporation. All rights reserved. # Licensed under CC BY-NC-SA 4.0 (non-commercial use only). # -------------------------------------------------------- # CroCo model during pretraining # -------------------------------------------------------- import torch import torch.nn as nn torch.backends.cuda.matmul.allow_tf32 = True # for gpu >= Ampere and pytorch >= 1.12 from functools import partial from models.blocks import Block, DecoderBlock, PatchEmbed from models.pos_embed import get_2d_sincos_pos_embed, RoPE2D from models.masking import RandomMask class CroCoNet(nn.Module): def __init__(self, img_size=224, # input image size patch_size=16, # patch_size mask_ratio=0.9, # ratios of masked tokens enc_embed_dim=768, # encoder feature dimension enc_depth=12, # encoder depth enc_num_heads=12, # encoder number of heads in the transformer block dec_embed_dim=512, # decoder feature dimension dec_depth=8, # decoder depth dec_num_heads=16, # decoder number of heads in the transformer block mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), norm_im2_in_dec=True, # whether to apply normalization of the 'memory' = (second image) in the decoder pos_embed='cosine', # positional embedding (either cosine or RoPE100) ): super(CroCoNet, self).__init__() # patch embeddings (with initialization done as in MAE) self._set_patch_embed(img_size, patch_size, enc_embed_dim) # mask generations self._set_mask_generator(self.patch_embed.num_patches, mask_ratio) self.pos_embed = pos_embed if pos_embed=='cosine': # positional embedding of the encoder enc_pos_embed = get_2d_sincos_pos_embed(enc_embed_dim, int(self.patch_embed.num_patches**.5), n_cls_token=0) self.register_buffer('enc_pos_embed', torch.from_numpy(enc_pos_embed).float()) # positional embedding of the decoder dec_pos_embed = get_2d_sincos_pos_embed(dec_embed_dim, int(self.patch_embed.num_patches**.5), n_cls_token=0) self.register_buffer('dec_pos_embed', torch.from_numpy(dec_pos_embed).float()) # pos embedding in each block self.rope = None # nothing for cosine elif pos_embed.startswith('RoPE'): # eg RoPE100 self.enc_pos_embed = None # nothing to add in the encoder with RoPE self.dec_pos_embed = None # nothing to add in the decoder with RoPE if RoPE2D is None: raise ImportError("Cannot find cuRoPE2D, please install it following the README instructions") freq = float(pos_embed[len('RoPE'):]) self.rope = RoPE2D(freq=freq) else: raise NotImplementedError('Unknown pos_embed '+pos_embed) # transformer for the encoder self.enc_depth = enc_depth self.enc_embed_dim = enc_embed_dim self.enc_blocks = nn.ModuleList([ Block(enc_embed_dim, enc_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, rope=self.rope) for i in range(enc_depth)]) self.enc_norm = norm_layer(enc_embed_dim) # masked tokens self._set_mask_token(dec_embed_dim) # decoder self._set_decoder(enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec) # prediction head self._set_prediction_head(dec_embed_dim, patch_size) # initializer weights self.initialize_weights() def _set_patch_embed(self, img_size=224, patch_size=16, enc_embed_dim=768): self.patch_embed = PatchEmbed(img_size, patch_size, 3, enc_embed_dim) def _set_mask_generator(self, num_patches, mask_ratio): self.mask_generator = RandomMask(num_patches, mask_ratio) def _set_mask_token(self, dec_embed_dim): self.mask_token = nn.Parameter(torch.zeros(1, 1, dec_embed_dim)) def _set_decoder(self, enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec): self.dec_depth = dec_depth self.dec_embed_dim = dec_embed_dim # transfer from encoder to decoder self.decoder_embed = nn.Linear(enc_embed_dim, dec_embed_dim, bias=True) # transformer for the decoder self.dec_blocks = nn.ModuleList([ DecoderBlock(dec_embed_dim, dec_num_heads, mlp_ratio=mlp_ratio, qkv_bias=True, norm_layer=norm_layer, norm_mem=norm_im2_in_dec, rope=self.rope) for i in range(dec_depth)]) # final norm layer self.dec_norm = norm_layer(dec_embed_dim) def _set_prediction_head(self, dec_embed_dim, patch_size): self.prediction_head = nn.Linear(dec_embed_dim, patch_size**2 * 3, bias=True) def initialize_weights(self): # patch embed self.patch_embed._init_weights() # mask tokens if self.mask_token is not None: torch.nn.init.normal_(self.mask_token, std=.02) # linears and layer norms self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): # we use xavier_uniform following official JAX ViT: torch.nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def _encode_image(self, image, do_mask=False, return_all_blocks=False): """ image has B x 3 x img_size x img_size do_mask: whether to perform masking or not return_all_blocks: if True, return the features at the end of every block instead of just the features from the last block (eg for some prediction heads) """ # embed the image into patches (x has size B x Npatches x C) # and get position if each return patch (pos has size B x Npatches x 2) x, pos = self.patch_embed(image) # add positional embedding without cls token if self.enc_pos_embed is not None: x = x + self.enc_pos_embed[None,...] # apply masking B,N,C = x.size() if do_mask: masks = self.mask_generator(x) x = x[~masks].view(B, -1, C) posvis = pos[~masks].view(B, -1, 2) else: B,N,C = x.size() masks = torch.zeros((B,N), dtype=bool) posvis = pos # now apply the transformer encoder and normalization if return_all_blocks: out = [] for blk in self.enc_blocks: x = blk(x, posvis) out.append(x) out[-1] = self.enc_norm(out[-1]) return out, pos, masks else: for blk in self.enc_blocks: x = blk(x, posvis) x = self.enc_norm(x) return x, pos, masks def _decoder(self, feat1, pos1, masks1, feat2, pos2, return_all_blocks=False): """ return_all_blocks: if True, return the features at the end of every block instead of just the features from the last block (eg for some prediction heads) masks1 can be None => assume image1 fully visible """ # encoder to decoder layer visf1 = self.decoder_embed(feat1) f2 = self.decoder_embed(feat2) # append masked tokens to the sequence B,Nenc,C = visf1.size() if masks1 is None: # downstreams f1_ = visf1 else: # pretraining Ntotal = masks1.size(1) f1_ = self.mask_token.repeat(B, Ntotal, 1).to(dtype=visf1.dtype) f1_[~masks1] = visf1.view(B * Nenc, C) # add positional embedding if self.dec_pos_embed is not None: f1_ = f1_ + self.dec_pos_embed f2 = f2 + self.dec_pos_embed # apply Transformer blocks out = f1_ out2 = f2 if return_all_blocks: _out, out = out, [] for blk in self.dec_blocks: _out, out2 = blk(_out, out2, pos1, pos2) out.append(_out) out[-1] = self.dec_norm(out[-1]) else: for blk in self.dec_blocks: out, out2 = blk(out, out2, pos1, pos2) out = self.dec_norm(out) return out def patchify(self, imgs): """ imgs: (B, 3, H, W) x: (B, L, patch_size**2 *3) """ p = self.patch_embed.patch_size[0] assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0 h = w = imgs.shape[2] // p x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p)) x = torch.einsum('nchpwq->nhwpqc', x) x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3)) return x def unpatchify(self, x, channels=3): """ x: (N, L, patch_size**2 *channels) imgs: (N, 3, H, W) """ patch_size = self.patch_embed.patch_size[0] h = w = int(x.shape[1]**.5) assert h * w == x.shape[1] x = x.reshape(shape=(x.shape[0], h, w, patch_size, patch_size, channels)) x = torch.einsum('nhwpqc->nchpwq', x) imgs = x.reshape(shape=(x.shape[0], channels, h * patch_size, h * patch_size)) return imgs def forward(self, img1, img2): """ img1: tensor of size B x 3 x img_size x img_size img2: tensor of size B x 3 x img_size x img_size out will be B x N x (3*patch_size*patch_size) masks are also returned as B x N just in case """ # encoder of the masked first image feat1, pos1, mask1 = self._encode_image(img1, do_mask=True) # encoder of the second image feat2, pos2, _ = self._encode_image(img2, do_mask=False) # decoder decfeat = self._decoder(feat1, pos1, mask1, feat2, pos2) # prediction head out = self.prediction_head(decfeat) # get target target = self.patchify(img1) return out, mask1, target