From fb3e3dbd5ede384173ba5fee299174ef948171ef Mon Sep 17 00:00:00 2001 From: han Date: Tue, 7 Sep 2021 20:24:45 +0900 Subject: [PATCH] Create pr --- timm/models/pr | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 timm/models/pr diff --git a/timm/models/pr b/timm/models/pr new file mode 100644 index 00000000..8393cc44 --- /dev/null +++ b/timm/models/pr @@ -0,0 +1,43 @@ +def resize_pos_embed(posemb, posemb_new, num_tokens=1, gs_new=()): + # Rescale the grid of position embeddings when loading from state_dict. Adapted from + # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224 + _logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape) + ntok_new = posemb_new.shape[1] + if num_tokens: + posemb_tok, posemb_grid = posemb[:, :num_tokens], posemb[0, num_tokens:] + ntok_new -= num_tokens + else: + posemb_tok, posemb_grid = posemb[:, :0], posemb[0] + gs_old = int(math.sqrt(len(posemb_grid))) + if not len(gs_new): # backwards compatibility + gs_new = [int(math.sqrt(ntok_new))] * 2 + assert len(gs_new) >= 2 + _logger.info('Position embedding grid-size from %s to %s', [gs_old, gs_old], gs_new) + posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2) + posemb_grid = F.interpolate(posemb_grid, size=gs_new, mode='bicubic', align_corners=False) + posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new[0] * gs_new[1], -1) + posemb = torch.cat([posemb_tok, posemb_grid], dim=1) + return posemb + +from einops import rearrange + +def interpolate(raw_param: torch.tensor, pretrained_weight: torch.tensor, token_num: int = 1): + # Rescale the grid of position embeddings when loading from state_dict. Adapted from + # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224 + _logger.info('Resized position embedding: %s to %s', pretrained_weight.shape, raw_param.shape) + + tokens, positional_embedding = pretrained_weight[:, :token_num], pretrained_weight[:, token_num:] + raw_dim, pretrained_dim = get_dim(raw_param), get_dim(positional_embedding) + _logger.info('Position embedding grid-size from %s to %s', [pretrained_dim, pretrained_dim], raw_dim) + + grid_positional_embedding = rearrange(positional_embedding, '1 (h w) d -> 1 d h w', h=pretrained_dim) + resized_grid_positional_embedding = F.interpolate(grid_positional_embedding, size=[raw_dim, raw_dim], mode='bicubic', align_corners=False) + positional_embedding = rearrange(resized_grid_positional_embedding, '1 d h w -> 1 (h w) d') + + pretrained_weight = torch.cat([tokens, positional_embedding], dim=1) + + return pretrained_weight + + +def get_dim(weight): + return int(math.sqrt(float(weight.size(1))))