wehub-resource-sync
551 lines
25 KiB
Python
551 lines
25 KiB
Python
from fairseq.models import FairseqEncoder, register_model, FairseqEncoderDecoderModel, register_model_architecture
|
|
from fairseq.models.transformer import TransformerDecoder, Embedding, TransformerModel
|
|
from fairseq.models.transformer import base_architecture as base_transformer
|
|
from fairseq.models.fairseq_encoder import EncoderOut
|
|
from torch.nn import Parameter
|
|
from fairseq import utils
|
|
from torch import Tensor
|
|
|
|
import torch
|
|
from torch.hub import load_state_dict_from_url
|
|
|
|
from timm.models import create_model
|
|
|
|
from functools import partial
|
|
import logging
|
|
import argparse
|
|
from typing import Dict, Optional, Tuple
|
|
from collections import OrderedDict
|
|
import os
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_MAX_TARGET_POSITIONS = 1024
|
|
|
|
from argparse import Namespace
|
|
from omegaconf import DictConfig
|
|
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
|
|
|
|
try:
|
|
from .unilm_models import UniLMDecoder
|
|
except:
|
|
from unilm_models import UniLMDecoder
|
|
|
|
@register_model('DeiT_TR')
|
|
@register_model('TrOCR')
|
|
class TrOCRModel(FairseqEncoderDecoderModel):
|
|
|
|
def load_state_dict(
|
|
self,
|
|
state_dict,
|
|
strict=True,
|
|
model_cfg: Optional[DictConfig] = None,
|
|
args: Optional[Namespace] = None,
|
|
):
|
|
|
|
if model_cfg is None and args is not None:
|
|
logger.warn("using 'args' is deprecated, please update your code to use dataclass config")
|
|
model_cfg = convert_namespace_to_omegaconf(args).model
|
|
|
|
self.upgrade_state_dict(state_dict)
|
|
|
|
from fairseq.checkpoint_utils import prune_state_dict
|
|
|
|
new_state_dict = prune_state_dict(state_dict, model_cfg)
|
|
if not model_cfg.ape:
|
|
model_seq_len = self.state_dict()['encoder.deit.pos_embed'].shape[1]
|
|
ckpt_seq_len = new_state_dict['encoder.deit.pos_embed'].shape[1]
|
|
if model_seq_len != ckpt_seq_len and getattr(args, "adapt_encoder_pos_embed", None):
|
|
logger.warning('Load from encoder.deit {:d} seq len to {:d}'.format(ckpt_seq_len, model_seq_len))
|
|
if model_seq_len <= ckpt_seq_len:
|
|
new_state_dict['encoder.deit.pos_embed'] = new_state_dict['encoder.deit.pos_embed'][:, :model_seq_len, :]
|
|
else:
|
|
t = self.state_dict()['encoder.deit.pos_embed']
|
|
t[:, :ckpt_seq_len, :] = new_state_dict['encoder.deit.pos_embed']
|
|
new_state_dict['encoder.deit.pos_embed'] = t
|
|
|
|
# if hasattr(model_cfg, "reset_dictionary") and model_cfg.reset_dictionary:
|
|
# logger.info('Reset token embed weights and output projection during loading pretrained models')
|
|
# del new_state_dict['decoder.embed_tokens.weight']
|
|
# del new_state_dict['decoder.output_projection.weight']
|
|
|
|
return super().load_state_dict(new_state_dict, strict=False)
|
|
|
|
@staticmethod
|
|
def add_args(parser):
|
|
TransformerModel.add_args(parser)
|
|
parser.add_argument(
|
|
'--deit-arch', type=str,
|
|
help='the arch name for the DeiT encoder'
|
|
)
|
|
parser.add_argument(
|
|
'--ape', action='store_true',
|
|
help='if use absolute_pos_embed'
|
|
)
|
|
parser.set_defaults(ape=False)
|
|
parser.add_argument(
|
|
'--mask-ratio', default=0.0, type=float,
|
|
help='the mask ratio for the encoder output masking.'
|
|
)
|
|
parser.add_argument(
|
|
'--only-keep-pretrained-decoder-structure', action='store_true',
|
|
help='if only keep the pretrained decoder structure'
|
|
)
|
|
parser.add_argument(
|
|
'--only-keep-pretrained-encoder-structure', action='store_true',
|
|
help='if only keep the pretrained encoder structure'
|
|
)
|
|
|
|
@staticmethod
|
|
def read_args_from_roberta(roberta_args: argparse.Namespace):
|
|
# TODO: this would become easier if encoder/decoder where using a similar
|
|
# TransformerConfig object
|
|
args = argparse.Namespace(**vars(roberta_args))
|
|
attr_map = [
|
|
("encoder_attention_heads", "decoder_attention_heads"),
|
|
("encoder_embed_dim", "decoder_embed_dim"),
|
|
("encoder_embed_dim", "decoder_output_dim"),
|
|
("encoder_normalize_before", "decoder_normalize_before"),
|
|
("encoder_layers_to_keep", "decoder_layers_to_keep"),
|
|
("encoder_ffn_embed_dim", "decoder_ffn_embed_dim"),
|
|
("encoder_layerdrop", "decoder_layerdrop"),
|
|
("encoder_layers", "decoder_layers"),
|
|
("encoder_learned_pos", "decoder_learned_pos"),
|
|
# should this be set from here ?
|
|
("max_positions", "max_target_positions"),
|
|
]
|
|
for k1, k2 in attr_map:
|
|
setattr(args, k2, getattr(roberta_args, k1))
|
|
|
|
args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
|
|
args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
|
|
args.share_decoder_input_output_embed = not roberta_args.untie_weights_roberta
|
|
return args
|
|
|
|
@classmethod
|
|
def build_model(cls, args, task):
|
|
encoder = TrOCREncoder(
|
|
args = args,
|
|
dictionary = task.source_dictionary
|
|
)
|
|
|
|
args.encoder_embed_dim = encoder.deit.embed_dim
|
|
|
|
if getattr(args, "max_target_positions", None) is None:
|
|
args.max_target_positions = DEFAULT_MAX_TARGET_POSITIONS
|
|
|
|
if getattr(args, "decoder_pretrained", None) == None or getattr(args, "decoder_pretrained", None).upper() == 'None':
|
|
logger.info('Decoder is randomly initialized.')
|
|
decoder_embed_tokens = cls.build_embedding(
|
|
args, task.target_dictionary, args.decoder_embed_dim, args.decoder_embed_path
|
|
)
|
|
decoder = TransformerDecoder(
|
|
args = args,
|
|
dictionary=task.target_dictionary,
|
|
embed_tokens=decoder_embed_tokens,
|
|
no_encoder_attn=False
|
|
)
|
|
|
|
elif getattr(args, "decoder_pretrained", None).startswith('roberta2'):
|
|
logger.info('Using the learned pos embedding version loading roberta.')
|
|
decoder_embed_tokens = cls.build_embedding(
|
|
args, task.target_dictionary, args.decoder_embed_dim, args.decoder_embed_path
|
|
)
|
|
|
|
pretrained_model = getattr(args, "decoder_pretrained", None)
|
|
specified = pretrained_model.find('-')!=-1
|
|
|
|
if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] != '0':
|
|
torch.distributed.barrier()
|
|
if specified:
|
|
pretrained_model = pretrained_model.replace('-', '.')
|
|
logger.info('Load pre-trained decoder parameters from {}'.format(pretrained_model))
|
|
roberta = torch.hub.load('pytorch/fairseq:main', pretrained_model)
|
|
elif args.decoder_layers == 6:
|
|
logger.info('Load pre-trained decoder parameters from roberta.base')
|
|
roberta = torch.hub.load('pytorch/fairseq:main', 'roberta.base')
|
|
elif args.decoder_layers == 12:
|
|
logger.info('Load pre-trained decoder parameters from roberta.large')
|
|
roberta = torch.hub.load('pytorch/fairseq:main', 'roberta.large')
|
|
else:
|
|
raise AttributeError('Cannot determined the pre-trained model')
|
|
|
|
if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] == '0':
|
|
torch.distributed.barrier()
|
|
|
|
roberta.model.args.encoder_layers = args.decoder_layers
|
|
roberta.model.args.fp16 = args.fp16
|
|
roberta_args = TrOCRModel.read_args_from_roberta(roberta.model.args)
|
|
roberta_args.encoder_embed_dim = args.encoder_embed_dim
|
|
|
|
decoder = TransformerDecoder(
|
|
roberta_args,
|
|
task.target_dictionary,
|
|
decoder_embed_tokens,
|
|
no_encoder_attn=False,
|
|
)
|
|
|
|
roberta_layers = roberta.model.encoder.sentence_encoder.layers
|
|
decoder_layers = decoder.layers
|
|
offset = len(roberta_layers) - len(decoder_layers)
|
|
assert offset >= 0
|
|
|
|
|
|
decoder_dict = roberta.state_dict()
|
|
new_decoder_dict = {}
|
|
for key, val in decoder_dict.items():
|
|
if key.startswith('model.encoder.sentence_encoder.layers.'):
|
|
layer_num = int(key[len('model.encoder.sentence_encoder.layers.'):].split('.')[0])
|
|
if layer_num - offset < 0:
|
|
continue
|
|
else:
|
|
new_key = 'model.encoder.sentence_encoder.layers.{}.'.format(
|
|
str(layer_num - offset)) + '.'.join(
|
|
key[len('model.encoder.sentence_encoder.layers.'):].split('.')[1:])
|
|
new_decoder_dict[new_key] = val
|
|
else:
|
|
new_decoder_dict[key] = val
|
|
decoder_dict = new_decoder_dict
|
|
|
|
for k, w in list(decoder_dict.items()):
|
|
if '.lm_head' in k:
|
|
k_proj = "output_projection." + k[len('model.encoder.lm_head.'):]
|
|
decoder_dict[k_proj] = w.detach().clone()
|
|
del decoder_dict[k]
|
|
|
|
del decoder_dict['_float_tensor']
|
|
del decoder_dict['output_projection.weight']
|
|
del decoder_dict['output_projection.bias']
|
|
del decoder_dict['output_projection.dense.weight']
|
|
del decoder_dict['output_projection.dense.bias']
|
|
del decoder_dict['output_projection.layer_norm.weight']
|
|
del decoder_dict['output_projection.layer_norm.bias']
|
|
|
|
new_decoder_dict = {}
|
|
for key, val in decoder_dict.items():
|
|
if "sentence_encoder" in key:
|
|
key = key[len('model.encoder.sentence_encoder.'):]
|
|
elif "encoder" in key:
|
|
key = key[len('model.encoder.'):]
|
|
new_decoder_dict[key] = val
|
|
|
|
if hasattr(args, 'only_keep_pretrained_decoder_structure') and args.only_keep_pretrained_decoder_structure:
|
|
logger.info('Only keep the pretrained decoder structure.')
|
|
pass
|
|
else:
|
|
missing_keys, unexpected_keys = decoder.load_state_dict(
|
|
new_decoder_dict, strict=False
|
|
)
|
|
|
|
elif getattr(args, "decoder_pretrained", None) == 'unilm':
|
|
logger.info('Decoder is pretrained using the unilm.')
|
|
|
|
prefix_of_parameter = 'bert'
|
|
|
|
decoder_embed_tokens = cls.build_embedding(
|
|
args, task.target_dictionary, args.decoder_embed_dim, args.decoder_embed_path
|
|
)
|
|
|
|
decoder = UniLMDecoder(
|
|
args,
|
|
task.target_dictionary,
|
|
decoder_embed_tokens,
|
|
no_encoder_attn=False,
|
|
)
|
|
|
|
if hasattr(args, 'decoder_pretrained_url') and args.decoder_pretrained_url != None and args.decoder_pretrained_url != '':
|
|
unilm_url = args.decoder_pretrained_url
|
|
logger.info('The unilm model url: {}.'.format(unilm_url[:unilm_url.find('?')]))
|
|
unilm_state_dict = torch.hub.load_state_dict_from_url(unilm_url)
|
|
|
|
unilm_layers = OrderedDict([(k, unilm_state_dict[k]) for k in unilm_state_dict.keys() if k.startswith(prefix_of_parameter + '.encoder.layer.')])
|
|
unilm_layers_num = []
|
|
for k in unilm_layers.keys():
|
|
t = k.replace(prefix_of_parameter + '.encoder.layer.', '')
|
|
t = t[:t.find('.')]
|
|
unilm_layers_num.append(int(t))
|
|
unilm_layers_num = max(unilm_layers_num) + 1
|
|
|
|
offset = unilm_layers_num - len(decoder.layers)
|
|
assert offset == 0
|
|
|
|
decoder_dict = decoder.state_dict()
|
|
# embedding
|
|
new_pos_weight = torch.zeros_like(decoder_dict['embed_positions.weight'])
|
|
# position padding will right offset padding idx + 1
|
|
new_pos_weight[task.target_dictionary.pad() + 1:, :] = unilm_state_dict[prefix_of_parameter + '.embeddings.position_embeddings.weight']
|
|
new_decoder_dict = {
|
|
'embed_tokens.weight': unilm_state_dict[prefix_of_parameter + '.embeddings.word_embeddings.weight'],
|
|
'embed_positions.weight': new_pos_weight,
|
|
'layernorm_embedding.weight': unilm_state_dict[prefix_of_parameter + '.embeddings.LayerNorm.weight'],
|
|
'layernorm_embedding.bias': unilm_state_dict[prefix_of_parameter + '.embeddings.LayerNorm.bias']
|
|
}
|
|
|
|
# layers
|
|
key_map = {
|
|
'self_attn.k_proj': 'attention.self.key',
|
|
'self_attn.v_proj': 'attention.self.value',
|
|
'self_attn.q_proj': 'attention.self.query',
|
|
'self_attn.out_proj': 'attention.output.dense',
|
|
'self_attn_layer_norm': 'attention.output.LayerNorm',
|
|
'fc1': 'intermediate.dense',
|
|
'fc2': 'output.dense',
|
|
'final_layer_norm': 'output.LayerNorm'
|
|
}
|
|
for layer_id in range(unilm_layers_num):
|
|
unilm_prefix = prefix_of_parameter + '.encoder.layer.{}.'.format(layer_id)
|
|
decoder_prefix = 'layers.{}.'.format(layer_id)
|
|
|
|
for key in key_map:
|
|
for suffix in ['.weight', '.bias']:
|
|
decoder_key = decoder_prefix + key + suffix
|
|
unilm_key = unilm_prefix + key_map[key] + suffix
|
|
if decoder_key in decoder_dict and unilm_key in unilm_state_dict:
|
|
new_decoder_dict[decoder_key] = unilm_state_dict[unilm_key]
|
|
|
|
if hasattr(args, "reset_dictionary") and args.reset_dictionary:
|
|
logger.info('Reset token embedding weights during decoder initialization.')
|
|
del new_decoder_dict['embed_tokens.weight']
|
|
elif hasattr(args, "adapt_dictionary") and args.adapt_dictionary:
|
|
unilm_embed_tokens_weight = new_decoder_dict['embed_tokens.weight']
|
|
logger.info('Adapt token embedding weights during decoder initialization from {} to {}'.format(unilm_embed_tokens_weight.shape[0], decoder_embed_tokens.weight.shape[0]))
|
|
new_decoder_dict['embed_tokens.weight'] = torch.zeros_like(decoder_dict['embed_tokens.weight'])
|
|
new_decoder_dict['embed_tokens.weight'][:min(unilm_embed_tokens_weight.shape[0], decoder_dict['embed_tokens.weight'].shape[0]), :] = unilm_embed_tokens_weight[:min(unilm_embed_tokens_weight.shape[0], decoder_dict['embed_tokens.weight'].shape[0]), :]
|
|
|
|
if hasattr(args, 'only_keep_pretrained_decoder_structure') and args.only_keep_pretrained_decoder_structure:
|
|
logger.info('Only keep the pretrained decoder structure.')
|
|
pass
|
|
else:
|
|
missing_keys, unexpected_keys = decoder.load_state_dict(
|
|
new_decoder_dict, strict=False
|
|
)
|
|
else:
|
|
logger.warning('You must specify the unilm model url or the decoder is randomly initialized.')
|
|
|
|
# freeze k_proj bias
|
|
for layer in decoder.layers:
|
|
layer.self_attn.k_proj.bias.requires_grad = False
|
|
|
|
elif getattr(args, "decoder_pretrained", None).startswith('roberta'):
|
|
logger.info('Using the old version loading roberta.')
|
|
decoder_embed_tokens = cls.build_embedding(
|
|
args, task.target_dictionary, args.decoder_embed_dim, args.decoder_embed_path
|
|
)
|
|
decoder = TransformerDecoder(
|
|
args = args,
|
|
dictionary=task.target_dictionary,
|
|
embed_tokens=decoder_embed_tokens,
|
|
no_encoder_attn=False
|
|
)
|
|
|
|
pretrained_model = getattr(args, "decoder_pretrained", None)
|
|
specified = pretrained_model.find('-')!=-1
|
|
|
|
if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] != '0':
|
|
torch.distributed.barrier()
|
|
|
|
if specified:
|
|
pretrained_model = pretrained_model.replace('-', '.')
|
|
logger.info('Load pre-trained decoder parameters from {}'.format(pretrained_model))
|
|
roberta = torch.hub.load('pytorch/fairseq:main', pretrained_model)
|
|
elif args.decoder_layers == 6:
|
|
logger.info('Load pre-trained decoder parameters from roberta.base')
|
|
roberta = torch.hub.load('pytorch/fairseq:main', 'roberta.base')
|
|
elif args.decoder_layers == 12:
|
|
logger.info('Load pre-trained decoder parameters from roberta.large')
|
|
roberta = torch.hub.load('pytorch/fairseq:main', 'roberta.large')
|
|
else:
|
|
raise AttributeError('Cannot determined the pre-trained model')
|
|
|
|
if 'LOCAL_RANK' in os.environ and os.environ['LOCAL_RANK'] == '0':
|
|
torch.distributed.barrier()
|
|
|
|
if hasattr(args, 'only_keep_pretrained_decoder_structure') and args.only_keep_pretrained_decoder_structure:
|
|
logger.info('Only keep the pretrained decoder structure.')
|
|
pass
|
|
else:
|
|
decoder.embed_tokens.load_state_dict(roberta.model.encoder.sentence_encoder.embed_tokens.state_dict())
|
|
roberta_layers = roberta.model.encoder.sentence_encoder.layers
|
|
decoder_layers = decoder.layers
|
|
offset = len(roberta_layers) - len(decoder_layers)
|
|
assert offset >= 0
|
|
|
|
for i in range(len(decoder_layers)):
|
|
roberta_i = i + offset
|
|
decoder_layers[i].self_attn.load_state_dict(roberta_layers[roberta_i].self_attn.state_dict())
|
|
decoder_layers[i].self_attn_layer_norm.load_state_dict(roberta_layers[roberta_i].self_attn_layer_norm.state_dict())
|
|
|
|
else:
|
|
raise Exception('Undefined decoder pretraining method.')
|
|
model = cls(encoder, decoder)
|
|
return model
|
|
|
|
@classmethod
|
|
def build_embedding(cls, args, dictionary, embed_dim, path=None):
|
|
num_embeddings = len(dictionary)
|
|
padding_idx = dictionary.pad()
|
|
|
|
emb = Embedding(num_embeddings, embed_dim, padding_idx)
|
|
# if provided, load from preloaded dictionaries
|
|
if path:
|
|
embed_dict = utils.parse_embedding(path)
|
|
utils.load_embedding(embed_dict, dictionary, emb)
|
|
return emb
|
|
|
|
def forward(self, imgs, prev_output_tokens, **kwargs):
|
|
encoder_out = self.encoder(imgs, **kwargs) # (seq_len, batch, embed_dim)
|
|
decoder_out = self.decoder(
|
|
prev_output_tokens, encoder_out=encoder_out, **kwargs
|
|
) # (batch, seq_len, vocab_size)
|
|
return decoder_out
|
|
|
|
|
|
@register_model_architecture('DeiT_TR', 'deit_base_decoder_base')
|
|
def deit_base_decoder_base(args):
|
|
# DeiT Encoder deit_base_distilled_patch16_384
|
|
args.deit_arch = getattr(args, "deit_arch", "deit_base_distilled_patch16_384")
|
|
# Transformer Decoder
|
|
# args.encoder_embed_dim = 768
|
|
base_transformer(args)
|
|
|
|
@register_model_architecture('DeiT_TR', 'deit_base_decoder_large')
|
|
def deit_base_decoder_large(args):
|
|
# DeiT Encoder deit_base_distilled_patch16_384
|
|
args.deit_arch = getattr(args, "deit_arch", "deit_base_distilled_patch16_384")
|
|
# Transformer Decoder
|
|
# args.encoder_embed_dim = 768
|
|
args.decoder_layers = getattr(args, "decoder_layers", 12)
|
|
args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1024)
|
|
args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 4096)
|
|
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
|
|
base_transformer(args)
|
|
|
|
@register_model_architecture('TrOCR', 'trocr_base')
|
|
@register_model_architecture('DeiT_TR', 'beit_base_decoder_large')
|
|
def beit_base_decoder_large(args):
|
|
# DeiT Encoder deit_base_distilled_patch16_384
|
|
args.deit_arch = getattr(args, "deit_arch", "beit_base_patch16_384")
|
|
# Transformer Decoder
|
|
# args.encoder_embed_dim = 768
|
|
args.decoder_layers = getattr(args, "decoder_layers", 12)
|
|
args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1024)
|
|
args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 4096)
|
|
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
|
|
base_transformer(args)
|
|
|
|
@register_model_architecture('TrOCR', 'trocr_large')
|
|
@register_model_architecture('DeiT_TR', 'beit_large_decoder_large')
|
|
def beit_large_decoder_large(args):
|
|
# DeiT Encoder deit_base_distilled_patch16_384
|
|
args.deit_arch = getattr(args, "deit_arch", "beit_large_patch16_384")
|
|
# Transformer Decoder
|
|
# args.encoder_embed_dim = 1024
|
|
args.decoder_layers = getattr(args, "decoder_layers", 12)
|
|
args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1024)
|
|
args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 4096)
|
|
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
|
|
base_transformer(args)
|
|
|
|
@register_model_architecture('DeiT_TR', 'deit_base_decoder_large_custom_size')
|
|
def deit_base_decoder_large_custom_size(args):
|
|
# DeiT Encoder deit_base_distilled_patch16_custom_size
|
|
args.deit_arch = getattr(args, "deit_arch", "deit_base_distilled_patch16_custom_size")
|
|
# Transformer Decoder
|
|
# args.encoder_embed_dim = 768
|
|
args.decoder_layers = getattr(args, "decoder_layers", 12)
|
|
args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 1024)
|
|
args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 4096)
|
|
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
|
|
base_transformer(args)
|
|
|
|
|
|
def nlrv4_compressed_tiny(args):
|
|
args.decoder_learned_pos = True
|
|
args.layernorm_embedding = True
|
|
args.decoder_attention_heads = 8
|
|
args.decoder_embed_dim = 256
|
|
args.decoder_output_dim = 256
|
|
args.decoder_ffn_embed_dim = 1024
|
|
args.dropout = 0.1
|
|
args.decoder_layers = 6
|
|
args.max_target_positions = 512
|
|
|
|
@register_model_architecture('TrOCR', 'trocr_small_224')
|
|
def trocr_small(args):
|
|
# DeiT Encoder deit_base_distilled_patch16_384
|
|
args.deit_arch = getattr(args, "deit_arch", "deit_small_distilled_patch16_224")
|
|
|
|
nlrv4_compressed_tiny(args)
|
|
# Transformer Decoder
|
|
base_transformer(args)
|
|
|
|
@register_model_architecture('TrOCR', 'trocr_small')
|
|
@register_model_architecture('TrOCR', 'trocr_small_384')
|
|
def trocr_small_384(args):
|
|
# DeiT Encoder deit_base_distilled_patch16_384
|
|
args.deit_arch = getattr(args, "deit_arch", "deit_small_distilled_patch16_384")
|
|
|
|
nlrv4_compressed_tiny(args)
|
|
# Transformer Decoder
|
|
base_transformer(args)
|
|
|
|
class TrOCREncoder(FairseqEncoder):
|
|
def __init__(self, args, dictionary):
|
|
super().__init__(dictionary)
|
|
|
|
if hasattr(args, 'only_keep_pretrained_encoder_structure') and args.only_keep_pretrained_encoder_structure:
|
|
pretrained = False
|
|
else:
|
|
pretrained = True
|
|
|
|
if 'custom_size' in args.deit_arch:
|
|
self.deit = create_model(args.deit_arch, pretrained=pretrained, img_size=args.input_size, ape=args.ape, mask_ratio=args.mask_ratio)
|
|
else:
|
|
self.deit = create_model(args.deit_arch, pretrained=pretrained, ape=args.ape, mask_ratio=args.mask_ratio)
|
|
|
|
self.fp16 = args.fp16
|
|
|
|
def forward(self, imgs):
|
|
if self.fp16:
|
|
imgs = imgs.half()
|
|
|
|
x, encoder_embedding = self.deit.forward_features(imgs) # bs, n + 2, dim
|
|
x = x.transpose(0, 1) # n + 2, bs, dim
|
|
|
|
encoder_padding_mask = torch.zeros(*x.shape[:2]).transpose(0, 1).to(imgs.device)
|
|
|
|
return {
|
|
"encoder_out": [x], # T x B x C
|
|
"encoder_padding_mask": [encoder_padding_mask], # B x T
|
|
"encoder_embedding": [encoder_embedding], # B x T x C
|
|
"encoder_states": [], # List[T x B x C]
|
|
"src_tokens": [],
|
|
"src_lengths": [],
|
|
}
|
|
|
|
def reorder_encoder_out(self, encoder_out, new_order):
|
|
"""
|
|
Reorder encoder output according to `new_order`.
|
|
|
|
Args:
|
|
encoder_out: output from the ``forward()`` method
|
|
new_order (LongTensor): desired order
|
|
|
|
Returns:
|
|
`encoder_out` rearranged according to `new_order`
|
|
"""
|
|
_encoder_out = encoder_out['encoder_out'][0]
|
|
_encoder_padding_mask = encoder_out['encoder_padding_mask'][0]
|
|
_encoder_embedding = encoder_out['encoder_embedding'][0]
|
|
return {
|
|
"encoder_out": [_encoder_out.index_select(1, new_order)],
|
|
"encoder_padding_mask": [_encoder_padding_mask.index_select(0, new_order)], # B x T
|
|
"encoder_embedding": [_encoder_padding_mask.index_select(0, new_order)], # B x T x C
|
|
"encoder_states": [],
|
|
"src_tokens": [],
|
|
"src_lengths": [],
|
|
}
|
|
|
|
if __name__ == '__main__':
|
|
pass
|