davanstrien/deepseek-ocr

1import torch.nn as nn2import torch3import torch.nn.functional as F4import copy5 6 7class MlpProjector(nn.Module):8 9    def __init__(self, cfg):10 11        super().__init__()12 13        self.cfg = cfg14 15        if cfg.projector_type == "identity":16            modules = nn.Identity()17 18        elif cfg.projector_type == "linear":19            modules = nn.Linear(cfg.input_dim, cfg.n_embed)20 21        elif cfg.projector_type == "mlp_gelu":22            mlp_depth = cfg.get("depth", 1)23            modules = [nn.Linear(cfg.input_dim, cfg.n_embed)]24            for _ in range(1, mlp_depth):25                modules.append(nn.GELU())26                modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))27            modules = nn.Sequential(*modules)28        29        elif cfg.projector_type == "normlayer_downsample_mlp_gelu":30            mlp_depth = cfg.get("depth", 1)31            mlp_ratio = cfg.get("mlp_ratio", 1)32            modules = [33                nn.LayerNorm(cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio),34                nn.Linear(cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio, cfg.n_embed * mlp_ratio)35            ]36            for _ in range(1, mlp_depth - 1):37                modules.append(nn.GELU())38                modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed * mlp_ratio))39            modules.append(nn.GELU())40            modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed))41            modules = nn.Sequential(*modules)42        43        elif cfg.projector_type == "downsample_mlp_gelu":44            mlp_depth = cfg.get("depth", 1)45            mlp_ratio = cfg.get("mlp_ratio", 1)46            modules = [nn.Linear(cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio, cfg.n_embed * mlp_ratio)]47            for _ in range(1, mlp_depth - 1):48                modules.append(nn.GELU())49                modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed * mlp_ratio))50            modules.append(nn.GELU())51            modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed))52            modules = nn.Sequential(*modules)53 54        elif cfg.projector_type == "low_high_hybrid_split_mlp_gelu":55            mlp_depth = cfg.get("depth", 1)56            self.high_up_proj = nn.Linear(cfg.input_dim, cfg.n_embed // 2)57            self.low_up_proj = nn.Linear(cfg.input_dim, cfg.n_embed // 2)58 59            modules = []60            for _ in range(1, mlp_depth):61                modules.append(nn.GELU())62                modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))63            modules = nn.Sequential(*modules)64 65        elif cfg.projector_type == "hybrid_split_feature_mlp_gelu":66            mlp_depth = cfg.get("depth", 1)67            channel_div = cfg.get("channel_div", 0.5)68            self.high_up_proj = nn.Linear(cfg.input_dim[0], int(cfg.n_embed * channel_div))69            self.low_up_proj = nn.Linear(cfg.input_dim[1], cfg.n_embed - int(cfg.n_embed * channel_div))70 71            modules = []72            for _ in range(1, mlp_depth):73                modules.append(nn.GELU())74                modules.append(nn.Linear(cfg.n_embed, cfg.n_embed))75            modules = nn.Sequential(*modules)76 77        elif cfg.projector_type == "low_high_split_mlp_gelu":78            mlp_depth = cfg.get("depth", 1)79            modules = []80            for _ in range(1, mlp_depth):81                modules.append(nn.GELU())82                modules.append(nn.Linear(cfg.n_embed // 2, cfg.n_embed // 2))83            modules = nn.Sequential(*modules)84            self.high_layers = nn.Sequential(*modules)85            self.low_layers = copy.deepcopy(modules)86 87        else:88            raise ValueError(f"Unknown projector type: {cfg.projector_type}")89 90        if cfg.get("token_pooling", False):91            self.token_pooling_layer = nn.Linear(cfg.input_dim * 4, cfg.input_dim)92 93        if cfg.get("conv_fusion_high_low_features", False):94            self.fusion_layer = nn.Linear(cfg.input_dim, cfg.input_dim)95        self.layers = modules96 97    def forward(self, x):98        if self.cfg.get("token_pooling", False):99            batch_size, wxh, channels = x.shape100            w = h = int(wxh**0.5)101            x = x.view(batch_size, w, h, channels)102            x = x.permute(0, 3, 1, 2)103            # import ipdb; ipdb.set_trace()104            patches = x.unfold(2, 2, 2).unfold(3, 2, 2)105            batch_size, channels, h_patches, w_patches, _, _ = patches.size()106            # 在通道维度上拼接107            patches = patches.contiguous().view(batch_size, channels, h_patches * w_patches, -1)108 109            # 通过线性层110            patches = patches.permute(0, 2, 1, 3).contiguous()111            patches = patches.view(batch_size, h_patches * w_patches, channels * 4)112 113            x = self.token_pooling_layer(patches)114        115        if self.cfg.get("conv_fusion_high_low_features", False):116            x = self.fusion_layer(x[:, 0]) + x[:, 1]117 118        if self.cfg.projector_type == 'low_high_hybrid_split_mlp_gelu':119            high_x, low_x = x[0], x[1]120            high_x = self.high_up_proj(high_x)121            low_x = self.low_up_proj(low_x)122            x = torch.concat([high_x, low_x], dim=-1)123        124        if self.cfg.projector_type == 'hybrid_split_feature_mlp_gelu':125            high_x = x[...,:self.cfg.input_dim[0]]126            low_x = x[...,self.cfg.input_dim[0]:]127            high_x = self.high_up_proj(high_x)128            low_x = self.low_up_proj(low_x)129            x = torch.concat([high_x, low_x], dim=-1)130        131        if self.cfg.projector_type == 'low_high_split_mlp_gelu':132            high_x, low_x = x[0], x[1]133            high_x = self.high_layers(high_x)134            low_x = self.low_layers(low_x)135            x = torch.concat([high_x, low_x], dim=-1)136            return x137        138        if self.cfg.projector_type == 'downsample_mlp_gelu' or self.cfg.projector_type == 'normlayer_downsample_mlp_gelu':139            bs, hw, input_dim = x.shape140            h = w = int((hw) ** 0.5)141 142            """compute padding"""143            if h % self.cfg.downsample_ratio:144                pad = self.cfg.downsample_ratio - h % self.cfg.downsample_ratio145            else:146                pad = 0147            x = x.reshape(bs, h, w, input_dim)148            if pad > 0:149                x = F.pad(x, (0, 0, 0, pad, 0, pad), "constant", 0)150 151            """4 to 1 concat"""152            x = x.permute(0, 3, 1, 2)  # B, C, H, W153            x = F.unfold(x, kernel_size=self.cfg.downsample_ratio, stride=self.cfg.downsample_ratio, padding=0) # B, C*4, HW // 4154            x = x.permute(0, 2, 1)155            156        return self.layers(x)157 158    @staticmethod159    def get_flops_per_sample(cfg):160        if cfg.projector_type == "linear":161            fwd = 2 * cfg.input_dim * cfg.n_embed162 163        elif "mlp_gelu" in cfg.projector_type :164            mlp_depth = cfg.get("depth", 1)165            downsample_ratio = cfg.get("downsample_ratio", 1)166            input_dim = sum(cfg.input_dim) if isinstance(cfg.input_dim, list) else cfg.input_dim167            input_dim = input_dim * downsample_ratio * downsample_ratio168            fwd = 2 * input_dim * cfg.n_embed + (mlp_depth - 1) * 2 * cfg.n_embed * cfg.n_embed169        else:170            fwd = 0171 172        return fwd * 3173 174 175