Logics-Thinking-8B / convnext_encoder.py

安觅

update

deb7abb about 1 year ago

19 kB

from transformers import CLIPImageProcessor
from torch.utils.checkpoint import checkpoint

from functools import partial
from typing import Callable, Optional, Tuple, Union

import torch
import torch.nn as nn

from timm.layers import trunc_normal_,AvgPool2dSame, DropPath, Mlp, GlobalResponseNormMlp, \
    LayerNorm2d, LayerNorm, create_conv2d, get_act_layer, make_divisible, to_ntuple
from timm.layers import NormMlpClassifierHead, ClassifierHead
from timm.models._manipulate import named_apply, checkpoint_seq


__all__ = ['ConvNeXt']  


class Downsample(nn.Module):

    def __init__(self, in_chs, out_chs, stride=1, dilation=1):
        super().__init__()
        avg_stride = stride if dilation == 1 else 1
        if stride > 1 or dilation > 1:
            avg_pool_fn = AvgPool2dSame if avg_stride == 1 and dilation > 1 else nn.AvgPool2d
            self.pool = avg_pool_fn(2, avg_stride, ceil_mode=True, count_include_pad=False)
        else:
            self.pool = nn.Identity()

        if in_chs != out_chs:
            self.conv = create_conv2d(in_chs, out_chs, 1, stride=1)
        else:
            self.conv = nn.Identity()

    def forward(self, x):
        x = self.pool(x)
        x = self.conv(x)
        return x




class ConvNeXtBlock(nn.Module):
    """ ConvNeXt Block
    There are two equivalent implementations:
      (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
      (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
    Unlike the official impl, this one allows choice of 1 or 2, 1x1 conv can be faster with appropriate
    choice of LayerNorm impl, however as model size increases the tradeoffs appear to change and nn.Linear
    is a better choice. This was observed with PyTorch 1.10 on 3090 GPU, it could change over time & w/ different HW.
    """

    def __init__(
            self,
            in_chs: int,
            out_chs: Optional[int] = None,
            kernel_size: int = 7,
            stride: int = 1,
            dilation: Union[int, Tuple[int, int]] = (1, 1),
            mlp_ratio: float = 4,
            conv_mlp: bool = False,
            conv_bias: bool = True,
            use_grn: bool = False,
            ls_init_value: Optional[float] = 1e-6,
            act_layer: Union[str, Callable] = 'gelu',
            norm_layer: Optional[Callable] = None,
            drop_path: float = 0.,
    ):
        """
        Args:
            in_chs: Block input channels.
            out_chs: Block output channels (same as in_chs if None).
            kernel_size: Depthwise convolution kernel size.
            stride: Stride of depthwise convolution.
            dilation: Tuple specifying input and output dilation of block.
            mlp_ratio: MLP expansion ratio.
            conv_mlp: Use 1x1 convolutions for MLP and a NCHW compatible norm layer if True.
            conv_bias: Apply bias for all convolution (linear) layers.
            use_grn: Use GlobalResponseNorm in MLP (from ConvNeXt-V2)
            ls_init_value: Layer-scale init values, layer-scale applied if not None.
            act_layer: Activation layer.
            norm_layer: Normalization layer (defaults to LN if not specified).
            drop_path: Stochastic depth probability.
        """
        super().__init__()
        out_chs = out_chs or in_chs
        dilation = to_ntuple(2)(dilation)
        act_layer = get_act_layer(act_layer)
        if not norm_layer:
            norm_layer = LayerNorm2d if conv_mlp else LayerNorm
        mlp_layer = partial(GlobalResponseNormMlp if use_grn else Mlp, use_conv=conv_mlp)
        self.use_conv_mlp = conv_mlp
        self.conv_dw = create_conv2d(
            in_chs,
            out_chs,
            kernel_size=kernel_size,
            stride=stride,
            dilation=dilation[0],
            depthwise=True,
            bias=conv_bias,
        )
        self.norm = norm_layer(out_chs)
        self.mlp = mlp_layer(out_chs, int(mlp_ratio * out_chs), act_layer=act_layer)
        self.weight = nn.Parameter(ls_init_value * torch.ones(out_chs)) if ls_init_value is not None else None
        if in_chs != out_chs or stride != 1 or dilation[0] != dilation[1]:
            self.shortcut = Downsample(in_chs, out_chs, stride=stride, dilation=dilation[0])
        else:
            self.shortcut = nn.Identity()
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        shortcut = x
        x = self.conv_dw(x)
        if self.use_conv_mlp:
            x = self.norm(x)
            x = self.mlp(x)
        else:
            x = x.permute(0, 2, 3, 1)
            x = self.norm(x)
            x = self.mlp(x)
            x = x.permute(0, 3, 1, 2)
        if self.weight is not None:
            x = x.mul(self.weight.reshape(1, -1, 1, 1))

        x = self.drop_path(x) + self.shortcut(shortcut)
        return x


class ConvNeXtStage(nn.Module):

    def __init__(
            self,
            in_chs,
            out_chs,
            kernel_size=7,
            stride=2,
            depth=2,
            dilation=(1, 1),
            drop_path_rates=None,
            ls_init_value=1.0,
            conv_mlp=False,
            conv_bias=True,
            use_grn=False,
            act_layer='gelu',
            norm_layer=None,
            norm_layer_cl=None
    ):
        super().__init__()
        self.grad_checkpointing = True

        if in_chs != out_chs or stride > 1 or dilation[0] != dilation[1]:
            ds_ks = 2 if stride > 1 or dilation[0] != dilation[1] else 1
            pad = 'same' if dilation[1] > 1 else 0  # same padding needed if dilation used
            self.downsample = nn.Sequential(
                norm_layer(in_chs),
                create_conv2d(
                    in_chs,
                    out_chs,
                    kernel_size=ds_ks,
                    stride=stride,
                    dilation=dilation[0],
                    padding=pad,
                    bias=conv_bias,
                ),
            )
            in_chs = out_chs
        else:
            self.downsample = nn.Identity()

        drop_path_rates = drop_path_rates or [0.] * depth
        stage_blocks = []
        for i in range(depth):
            stage_blocks.append(ConvNeXtBlock(
                in_chs=in_chs,
                out_chs=out_chs,
                kernel_size=kernel_size,
                dilation=dilation[1],
                drop_path=drop_path_rates[i],
                ls_init_value=ls_init_value,
                conv_mlp=conv_mlp,
                conv_bias=conv_bias,
                use_grn=use_grn,
                act_layer=act_layer,
                norm_layer=norm_layer if conv_mlp else norm_layer_cl,
            ))
            in_chs = out_chs
        self.blocks = nn.Sequential(*stage_blocks)

    def forward(self, x):
        x = self.downsample(x)
        if self.grad_checkpointing and not torch.jit.is_scripting():
            x = checkpoint_seq(self.blocks, x)
        else:
            x = self.blocks(x)
        return x


class ConvNeXt(nn.Module):
    r""" ConvNeXt
        A PyTorch impl of : `A ConvNet for the 2020s`  - https://arxiv.org/pdf/2201.03545.pdf
    """

    def __init__(
            self,
            in_chans: int = 3,
            num_classes: int = 1024,
            global_pool: str = 'avg',
            output_stride: int = 32,
            depths: Tuple[int, ...] = (3, 3, 9, 3),
            dims: Tuple[int, ...] = (96, 192, 384, 768),
            kernel_sizes: Union[int, Tuple[int, ...]] = 7,
            ls_init_value: Optional[float] = 1e-6,
            stem_type: str = 'patch',
            patch_size: int = 4,
            head_init_scale: float = 1.,
            head_norm_first: bool = False,
            head_hidden_size: Optional[int] = None,
            conv_mlp: bool = False,
            conv_bias: bool = True,
            use_grn: bool = False,
            act_layer: Union[str, Callable] = 'gelu',
            norm_layer: Optional[Union[str, Callable]] = None,
            norm_eps: Optional[float] = None,
            drop_rate: float = 0.,
            drop_path_rate: float = 0.,
    ):
        """
        Args:
            in_chans: Number of input image channels.
            num_classes: Number of classes for classification head.
            global_pool: Global pooling type.
            output_stride: Output stride of network, one of (8, 16, 32).
            depths: Number of blocks at each stage.
            dims: Feature dimension at each stage.
            kernel_sizes: Depthwise convolution kernel-sizes for each stage.
            ls_init_value: Init value for Layer Scale, disabled if None.
            stem_type: Type of stem.
            patch_size: Stem patch size for patch stem.
            head_init_scale: Init scaling value for classifier weights and biases.
            head_norm_first: Apply normalization before global pool + head.
            head_hidden_size: Size of MLP hidden layer in head if not None and head_norm_first == False.
            conv_mlp: Use 1x1 conv in MLP, improves speed for small networks w/ chan last.
            conv_bias: Use bias layers w/ all convolutions.
            use_grn: Use Global Response Norm (ConvNeXt-V2) in MLP.
            act_layer: Activation layer type.
            norm_layer: Normalization layer type.
            drop_rate: Head pre-classifier dropout rate.
            drop_path_rate: Stochastic depth drop rate.
        """
        super().__init__()
        assert output_stride in (8, 16, 32)
        kernel_sizes = to_ntuple(4)(kernel_sizes)
        if norm_layer is None:
            norm_layer = LayerNorm2d
            norm_layer_cl = norm_layer if conv_mlp else LayerNorm
            if norm_eps is not None:
                norm_layer = partial(norm_layer, eps=norm_eps)
                norm_layer_cl = partial(norm_layer_cl, eps=norm_eps)
        else:
            assert conv_mlp,\
                'If a norm_layer is specified, conv MLP must be used so all norm expect rank-4, channels-first input'
            norm_layer_cl = norm_layer
            if norm_eps is not None:
                norm_layer_cl = partial(norm_layer_cl, eps=norm_eps)

        self.num_classes = num_classes
        self.drop_rate = drop_rate
        self.feature_info = []

        assert stem_type in ('patch', 'overlap', 'overlap_tiered')
        if stem_type == 'patch':
            # NOTE: this stem is a minimal form of ViT PatchEmbed, as used in SwinTransformer w/ patch_size = 4
            self.stem = nn.Sequential(
                nn.Conv2d(in_chans, dims[0], kernel_size=patch_size, stride=patch_size, bias=conv_bias),
                norm_layer(dims[0]),
            )
            stem_stride = patch_size
        else:
            mid_chs = make_divisible(dims[0] // 2) if 'tiered' in stem_type else dims[0]
            self.stem = nn.Sequential(
                nn.Conv2d(in_chans, mid_chs, kernel_size=3, stride=2, padding=1, bias=conv_bias),
                nn.Conv2d(mid_chs, dims[0], kernel_size=3, stride=2, padding=1, bias=conv_bias),
                norm_layer(dims[0]),
            )
            stem_stride = 4

        self.stages = nn.Sequential()
        dp_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depths)).split(depths)]
        stages = []
        prev_chs = dims[0]
        curr_stride = stem_stride
        dilation = 1
        # 4 feature resolution stages, each consisting of multiple residual blocks
        for i in range(4):
            stride = 2 if curr_stride == 2 or i > 0 else 1
            if curr_stride >= output_stride and stride > 1:
                dilation *= stride
                stride = 1
            curr_stride *= stride
            first_dilation = 1 if dilation in (1, 2) else 2
            out_chs = dims[i]
            stages.append(ConvNeXtStage(
                prev_chs,
                out_chs,
                kernel_size=kernel_sizes[i],
                stride=stride,
                dilation=(first_dilation, dilation),
                depth=depths[i],
                drop_path_rates=dp_rates[i],
                ls_init_value=ls_init_value,
                conv_mlp=conv_mlp,
                conv_bias=conv_bias,
                use_grn=use_grn,
                act_layer=act_layer,
                norm_layer=norm_layer,
                norm_layer_cl=norm_layer_cl,
            ))
            prev_chs = out_chs
            # NOTE feature_info use currently assumes stage 0 == stride 1, rest are stride 2
            self.feature_info += [dict(num_chs=prev_chs, reduction=curr_stride, module=f'stages.{i}')]
        self.stages = nn.Sequential(*stages)
        self.num_features = prev_chs

        # if head_norm_first == true, norm -> global pool -> fc ordering, like most other nets
        # otherwise pool -> norm -> fc, the default ConvNeXt ordering (pretrained FB weights)
        if head_norm_first:
            assert not head_hidden_size
            self.norm_pre = norm_layer(self.num_features)
            self.head = ClassifierHead(
                self.num_features,
                num_classes,
                pool_type=global_pool,
                drop_rate=self.drop_rate,
            )
        else:
            self.norm_pre = nn.Identity()
            self.head = NormMlpClassifierHead(
                self.num_features,
                num_classes,
                hidden_size=head_hidden_size,
                pool_type=global_pool,
                drop_rate=self.drop_rate,
                norm_layer=norm_layer,
                act_layer='gelu',
            )
        named_apply(partial(_init_weights, head_init_scale=head_init_scale), self)

    @torch.jit.ignore
    def group_matcher(self, coarse=False):
        return dict(
            stem=r'^stem',
            blocks=r'^stages\.(\d+)' if coarse else [
                (r'^stages\.(\d+)\.downsample', (0,)),  # blocks
                (r'^stages\.(\d+)\.blocks\.(\d+)', None),
                (r'^norm_pre', (99999,))
            ]
        )

    @torch.jit.ignore
    def set_grad_checkpointing(self, enable=True):
        for s in self.stages:
            s.grad_checkpointing = enable

    @torch.jit.ignore
    def get_classifier(self):
        return self.head.fc

    def reset_classifier(self, num_classes=0, global_pool=None):
        self.head.reset(num_classes, global_pool)

    def forward_features(self, x):
        x = self.stem(x)
        x = self.stages(x)
        x = self.norm_pre(x)
        return x

    def forward_head(self, x, pre_logits: bool = False):
        return self.head(x, pre_logits=True) if pre_logits else self.head(x)

    def forward(self, x):
        x = self.forward_features(x)
        x = self.forward_head(x)
        return x


def _init_weights(module, name=None, head_init_scale=1.0):
    if isinstance(module, nn.Conv2d):
        trunc_normal_(module.weight, std=.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Linear):
        trunc_normal_(module.weight, std=.02)
        nn.init.zeros_(module.bias)
        if name and 'head.' in name:
            module.weight.data.mul_(head_init_scale)
            module.bias.data.mul_(head_init_scale)


cfg={
    "crop_size": 256,
    "do_center_crop": True,
    "do_normalize": True,
    "do_resize": True,
    "feature_extractor_type": "CLIPFeatureExtractor",
    "image_mean": [
        0.48145466,
        0.4578275,
        0.40821073
    ],
    "image_std": [
        0.26862954,
        0.26130258,
        0.27577711
    ],
    "resample": 3,
    "size": 256
}

class ConvNextVisionTower(nn.Module):
    def __init__(self, vision_tower, args, delay_load=False):
        super().__init__()

        self.is_loaded = False
        self.freeze_vision=False
        self.input_image_size=args.input_image_size
        self.vision_tower_name = vision_tower
        self.select_layer = -1 # hardcode
        self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')
        self.xpfs = args.xpfs if hasattr(args, 'xpfs') else None
        print(f"self.xpfs:{self.xpfs}")

    def load_model(self, gradient_checkpointing=True):


        self.image_processor = CLIPImageProcessor(**cfg)
        if 'xxlarge' in self.vision_tower_name:
            model_args = dict(
                depths=[3, 4, 30, 3], 
                dims=[384, 768, 1536, 3072], 
                norm_eps=1e-5, 
                num_classes=1024
            )
            self.vision_tower = ConvNeXt(**model_args) 
            setattr(self.vision_tower, 'hidden_size', 3072)
        else:
            raise NotImplementedError
        
        if self.freeze_vision:
            self.vision_tower.requires_grad_(False)

        for s in self.vision_tower.stages:
            s.grad_checkpointing = gradient_checkpointing

        if self.input_image_size is not None:
            self.image_processor.size=self.input_image_size
            self.image_processor.crop_size={
                'height':self.input_image_size,
                'width': self.input_image_size
            }

        self.is_loaded = True

    def feature_select(self, image_forward_outs):
        image_features = image_forward_outs[self.select_layer]
        return image_features

    def forward_features(self, x):
        x = self.vision_tower.stem(x)
        image_forward_out=[]
        for blk in self.vision_tower.stages:
            x = torch.utils.checkpoint.checkpoint(blk, x)
            b,c,h,w=x.shape
            image_forward_out.append(x.view(b,c,-1).transpose(1,2))
        return image_forward_out

    def forward(self, images):
        if self.freeze_vision:
            with torch.no_grad():
                image_features = self._forward_images(images)
        else:
            image_features = self._forward_images(images)

        return image_features

    def _forward_images(self, images):

        image_forward_outs = self.forward_features(images.to(device=self.device, dtype=self.dtype))
        image_features = self.feature_select(image_forward_outs)

        return image_features

    @property
    def dummy_feature(self):
        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)

    @property
    def dtype(self):
        return next(self.vision_tower.parameters()).dtype

    @property
    def device(self):
        return next(self.vision_tower.parameters()).device

    @property
    def config(self):
        assert  NotImplementedError
        pass

    @property
    def num_attention_heads(self):
        # as constant
        return 16
    @property
    def num_layers(self):
        # as constant
        return 4
    @property
    def hidden_size(self):
        return self.vision_tower.hidden_size

    @property
    def num_patches(self):
        return (cfg['image_size'] // self.patch_embed.patch_size[0]) ** 2

	from transformers import CLIPImageProcessor
	from torch.utils.checkpoint import checkpoint

	from functools import partial
	from typing import Callable, Optional, Tuple, Union

	import torch
	import torch.nn as nn

	from timm.layers import trunc_normal_,AvgPool2dSame, DropPath, Mlp, GlobalResponseNormMlp, \
	LayerNorm2d, LayerNorm, create_conv2d, get_act_layer, make_divisible, to_ntuple
	from timm.layers import NormMlpClassifierHead, ClassifierHead
	from timm.models._manipulate import named_apply, checkpoint_seq


	__all__ = ['ConvNeXt']


	class Downsample(nn.Module):

	def __init__(self, in_chs, out_chs, stride=1, dilation=1):
	super().__init__()
	avg_stride = stride if dilation == 1 else 1
	if stride > 1 or dilation > 1:
	avg_pool_fn = AvgPool2dSame if avg_stride == 1 and dilation > 1 else nn.AvgPool2d
	self.pool = avg_pool_fn(2, avg_stride, ceil_mode=True, count_include_pad=False)
	else:
	self.pool = nn.Identity()

	if in_chs != out_chs:
	self.conv = create_conv2d(in_chs, out_chs, 1, stride=1)
	else:
	self.conv = nn.Identity()

	def forward(self, x):
	x = self.pool(x)
	x = self.conv(x)
	return x




	class ConvNeXtBlock(nn.Module):
	""" ConvNeXt Block
	There are two equivalent implementations:
	(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
	(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back

	Unlike the official impl, this one allows choice of 1 or 2, 1x1 conv can be faster with appropriate
	choice of LayerNorm impl, however as model size increases the tradeoffs appear to change and nn.Linear
	is a better choice. This was observed with PyTorch 1.10 on 3090 GPU, it could change over time & w/ different HW.
	"""

	def __init__(
	self,
	in_chs: int,
	out_chs: Optional[int] = None,
	kernel_size: int = 7,
	stride: int = 1,
	dilation: Union[int, Tuple[int, int]] = (1, 1),
	mlp_ratio: float = 4,
	conv_mlp: bool = False,
	conv_bias: bool = True,
	use_grn: bool = False,
	ls_init_value: Optional[float] = 1e-6,
	act_layer: Union[str, Callable] = 'gelu',
	norm_layer: Optional[Callable] = None,
	drop_path: float = 0.,
	):
	"""

	Args:
	in_chs: Block input channels.
	out_chs: Block output channels (same as in_chs if None).
	kernel_size: Depthwise convolution kernel size.
	stride: Stride of depthwise convolution.
	dilation: Tuple specifying input and output dilation of block.
	mlp_ratio: MLP expansion ratio.
	conv_mlp: Use 1x1 convolutions for MLP and a NCHW compatible norm layer if True.
	conv_bias: Apply bias for all convolution (linear) layers.
	use_grn: Use GlobalResponseNorm in MLP (from ConvNeXt-V2)
	ls_init_value: Layer-scale init values, layer-scale applied if not None.
	act_layer: Activation layer.
	norm_layer: Normalization layer (defaults to LN if not specified).
	drop_path: Stochastic depth probability.
	"""
	super().__init__()
	out_chs = out_chs or in_chs
	dilation = to_ntuple(2)(dilation)
	act_layer = get_act_layer(act_layer)
	if not norm_layer:
	norm_layer = LayerNorm2d if conv_mlp else LayerNorm
	mlp_layer = partial(GlobalResponseNormMlp if use_grn else Mlp, use_conv=conv_mlp)
	self.use_conv_mlp = conv_mlp
	self.conv_dw = create_conv2d(
	in_chs,
	out_chs,
	kernel_size=kernel_size,
	stride=stride,
	dilation=dilation[0],
	depthwise=True,
	bias=conv_bias,
	)
	self.norm = norm_layer(out_chs)
	self.mlp = mlp_layer(out_chs, int(mlp_ratio * out_chs), act_layer=act_layer)
	self.weight = nn.Parameter(ls_init_value * torch.ones(out_chs)) if ls_init_value is not None else None
	if in_chs != out_chs or stride != 1 or dilation[0] != dilation[1]:
	self.shortcut = Downsample(in_chs, out_chs, stride=stride, dilation=dilation[0])
	else:
	self.shortcut = nn.Identity()
	self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

	def forward(self, x):
	shortcut = x
	x = self.conv_dw(x)
	if self.use_conv_mlp:
	x = self.norm(x)
	x = self.mlp(x)
	else:
	x = x.permute(0, 2, 3, 1)
	x = self.norm(x)
	x = self.mlp(x)
	x = x.permute(0, 3, 1, 2)
	if self.weight is not None:
	x = x.mul(self.weight.reshape(1, -1, 1, 1))

	x = self.drop_path(x) + self.shortcut(shortcut)
	return x


	class ConvNeXtStage(nn.Module):

	def __init__(
	self,
	in_chs,
	out_chs,
	kernel_size=7,
	stride=2,
	depth=2,
	dilation=(1, 1),
	drop_path_rates=None,
	ls_init_value=1.0,
	conv_mlp=False,
	conv_bias=True,
	use_grn=False,
	act_layer='gelu',
	norm_layer=None,
	norm_layer_cl=None
	):
	super().__init__()
	self.grad_checkpointing = True

	if in_chs != out_chs or stride > 1 or dilation[0] != dilation[1]:
	ds_ks = 2 if stride > 1 or dilation[0] != dilation[1] else 1
	pad = 'same' if dilation[1] > 1 else 0 # same padding needed if dilation used
	self.downsample = nn.Sequential(
	norm_layer(in_chs),
	create_conv2d(
	in_chs,
	out_chs,
	kernel_size=ds_ks,
	stride=stride,
	dilation=dilation[0],
	padding=pad,
	bias=conv_bias,
	),
	)
	in_chs = out_chs
	else:
	self.downsample = nn.Identity()

	drop_path_rates = drop_path_rates or [0.] * depth
	stage_blocks = []
	for i in range(depth):
	stage_blocks.append(ConvNeXtBlock(
	in_chs=in_chs,
	out_chs=out_chs,
	kernel_size=kernel_size,
	dilation=dilation[1],
	drop_path=drop_path_rates[i],
	ls_init_value=ls_init_value,
	conv_mlp=conv_mlp,
	conv_bias=conv_bias,
	use_grn=use_grn,
	act_layer=act_layer,
	norm_layer=norm_layer if conv_mlp else norm_layer_cl,
	))
	in_chs = out_chs
	self.blocks = nn.Sequential(*stage_blocks)

	def forward(self, x):
	x = self.downsample(x)
	if self.grad_checkpointing and not torch.jit.is_scripting():
	x = checkpoint_seq(self.blocks, x)
	else:
	x = self.blocks(x)
	return x


	class ConvNeXt(nn.Module):
	r""" ConvNeXt
	A PyTorch impl of : `A ConvNet for the 2020s` - https://arxiv.org/pdf/2201.03545.pdf
	"""

	def __init__(
	self,
	in_chans: int = 3,
	num_classes: int = 1024,
	global_pool: str = 'avg',
	output_stride: int = 32,
	depths: Tuple[int, ...] = (3, 3, 9, 3),
	dims: Tuple[int, ...] = (96, 192, 384, 768),
	kernel_sizes: Union[int, Tuple[int, ...]] = 7,
	ls_init_value: Optional[float] = 1e-6,
	stem_type: str = 'patch',
	patch_size: int = 4,
	head_init_scale: float = 1.,
	head_norm_first: bool = False,
	head_hidden_size: Optional[int] = None,
	conv_mlp: bool = False,
	conv_bias: bool = True,
	use_grn: bool = False,
	act_layer: Union[str, Callable] = 'gelu',
	norm_layer: Optional[Union[str, Callable]] = None,
	norm_eps: Optional[float] = None,
	drop_rate: float = 0.,
	drop_path_rate: float = 0.,
	):
	"""
	Args:
	in_chans: Number of input image channels.
	num_classes: Number of classes for classification head.
	global_pool: Global pooling type.
	output_stride: Output stride of network, one of (8, 16, 32).
	depths: Number of blocks at each stage.
	dims: Feature dimension at each stage.
	kernel_sizes: Depthwise convolution kernel-sizes for each stage.
	ls_init_value: Init value for Layer Scale, disabled if None.
	stem_type: Type of stem.
	patch_size: Stem patch size for patch stem.
	head_init_scale: Init scaling value for classifier weights and biases.
	head_norm_first: Apply normalization before global pool + head.
	head_hidden_size: Size of MLP hidden layer in head if not None and head_norm_first == False.
	conv_mlp: Use 1x1 conv in MLP, improves speed for small networks w/ chan last.
	conv_bias: Use bias layers w/ all convolutions.
	use_grn: Use Global Response Norm (ConvNeXt-V2) in MLP.
	act_layer: Activation layer type.
	norm_layer: Normalization layer type.
	drop_rate: Head pre-classifier dropout rate.
	drop_path_rate: Stochastic depth drop rate.
	"""
	super().__init__()
	assert output_stride in (8, 16, 32)
	kernel_sizes = to_ntuple(4)(kernel_sizes)
	if norm_layer is None:
	norm_layer = LayerNorm2d
	norm_layer_cl = norm_layer if conv_mlp else LayerNorm
	if norm_eps is not None:
	norm_layer = partial(norm_layer, eps=norm_eps)
	norm_layer_cl = partial(norm_layer_cl, eps=norm_eps)
	else:
	assert conv_mlp,\
	'If a norm_layer is specified, conv MLP must be used so all norm expect rank-4, channels-first input'
	norm_layer_cl = norm_layer
	if norm_eps is not None:
	norm_layer_cl = partial(norm_layer_cl, eps=norm_eps)

	self.num_classes = num_classes
	self.drop_rate = drop_rate
	self.feature_info = []

	assert stem_type in ('patch', 'overlap', 'overlap_tiered')
	if stem_type == 'patch':
	# NOTE: this stem is a minimal form of ViT PatchEmbed, as used in SwinTransformer w/ patch_size = 4
	self.stem = nn.Sequential(
	nn.Conv2d(in_chans, dims[0], kernel_size=patch_size, stride=patch_size, bias=conv_bias),
	norm_layer(dims[0]),
	)
	stem_stride = patch_size
	else:
	mid_chs = make_divisible(dims[0] // 2) if 'tiered' in stem_type else dims[0]
	self.stem = nn.Sequential(
	nn.Conv2d(in_chans, mid_chs, kernel_size=3, stride=2, padding=1, bias=conv_bias),
	nn.Conv2d(mid_chs, dims[0], kernel_size=3, stride=2, padding=1, bias=conv_bias),
	norm_layer(dims[0]),
	)
	stem_stride = 4

	self.stages = nn.Sequential()
	dp_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(depths)).split(depths)]
	stages = []
	prev_chs = dims[0]
	curr_stride = stem_stride
	dilation = 1
	# 4 feature resolution stages, each consisting of multiple residual blocks
	for i in range(4):
	stride = 2 if curr_stride == 2 or i > 0 else 1
	if curr_stride >= output_stride and stride > 1:
	dilation *= stride
	stride = 1
	curr_stride *= stride
	first_dilation = 1 if dilation in (1, 2) else 2
	out_chs = dims[i]
	stages.append(ConvNeXtStage(
	prev_chs,
	out_chs,
	kernel_size=kernel_sizes[i],
	stride=stride,
	dilation=(first_dilation, dilation),
	depth=depths[i],
	drop_path_rates=dp_rates[i],
	ls_init_value=ls_init_value,
	conv_mlp=conv_mlp,
	conv_bias=conv_bias,
	use_grn=use_grn,
	act_layer=act_layer,
	norm_layer=norm_layer,
	norm_layer_cl=norm_layer_cl,
	))
	prev_chs = out_chs
	# NOTE feature_info use currently assumes stage 0 == stride 1, rest are stride 2
	self.feature_info += [dict(num_chs=prev_chs, reduction=curr_stride, module=f'stages.{i}')]
	self.stages = nn.Sequential(*stages)
	self.num_features = prev_chs

	# if head_norm_first == true, norm -> global pool -> fc ordering, like most other nets
	# otherwise pool -> norm -> fc, the default ConvNeXt ordering (pretrained FB weights)
	if head_norm_first:
	assert not head_hidden_size
	self.norm_pre = norm_layer(self.num_features)
	self.head = ClassifierHead(
	self.num_features,
	num_classes,
	pool_type=global_pool,
	drop_rate=self.drop_rate,
	)
	else:
	self.norm_pre = nn.Identity()
	self.head = NormMlpClassifierHead(
	self.num_features,
	num_classes,
	hidden_size=head_hidden_size,
	pool_type=global_pool,
	drop_rate=self.drop_rate,
	norm_layer=norm_layer,
	act_layer='gelu',
	)
	named_apply(partial(_init_weights, head_init_scale=head_init_scale), self)

	@torch.jit.ignore
	def group_matcher(self, coarse=False):
	return dict(
	stem=r'^stem',
	blocks=r'^stages\.(\d+)' if coarse else [
	(r'^stages\.(\d+)\.downsample', (0,)), # blocks
	(r'^stages\.(\d+)\.blocks\.(\d+)', None),
	(r'^norm_pre', (99999,))
	]
	)

	@torch.jit.ignore
	def set_grad_checkpointing(self, enable=True):
	for s in self.stages:
	s.grad_checkpointing = enable

	@torch.jit.ignore
	def get_classifier(self):
	return self.head.fc

	def reset_classifier(self, num_classes=0, global_pool=None):
	self.head.reset(num_classes, global_pool)

	def forward_features(self, x):
	x = self.stem(x)
	x = self.stages(x)
	x = self.norm_pre(x)
	return x

	def forward_head(self, x, pre_logits: bool = False):
	return self.head(x, pre_logits=True) if pre_logits else self.head(x)

	def forward(self, x):
	x = self.forward_features(x)
	x = self.forward_head(x)
	return x


	def _init_weights(module, name=None, head_init_scale=1.0):
	if isinstance(module, nn.Conv2d):
	trunc_normal_(module.weight, std=.02)
	if module.bias is not None:
	nn.init.zeros_(module.bias)
	elif isinstance(module, nn.Linear):
	trunc_normal_(module.weight, std=.02)
	nn.init.zeros_(module.bias)
	if name and 'head.' in name:
	module.weight.data.mul_(head_init_scale)
	module.bias.data.mul_(head_init_scale)


	cfg={
	"crop_size": 256,
	"do_center_crop": True,
	"do_normalize": True,
	"do_resize": True,
	"feature_extractor_type": "CLIPFeatureExtractor",
	"image_mean": [
	0.48145466,
	0.4578275,
	0.40821073
	],
	"image_std": [
	0.26862954,
	0.26130258,
	0.27577711
	],
	"resample": 3,
	"size": 256
	}

	class ConvNextVisionTower(nn.Module):
	def __init__(self, vision_tower, args, delay_load=False):
	super().__init__()

	self.is_loaded = False
	self.freeze_vision=False
	self.input_image_size=args.input_image_size
	self.vision_tower_name = vision_tower
	self.select_layer = -1 # hardcode
	self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')
	self.xpfs = args.xpfs if hasattr(args, 'xpfs') else None
	print(f"self.xpfs:{self.xpfs}")

	def load_model(self, gradient_checkpointing=True):


	self.image_processor = CLIPImageProcessor(**cfg)
	if 'xxlarge' in self.vision_tower_name:
	model_args = dict(
	depths=[3, 4, 30, 3],
	dims=[384, 768, 1536, 3072],
	norm_eps=1e-5,
	num_classes=1024
	)
	self.vision_tower = ConvNeXt(**model_args)
	setattr(self.vision_tower, 'hidden_size', 3072)
	else:
	raise NotImplementedError

	if self.freeze_vision:
	self.vision_tower.requires_grad_(False)

	for s in self.vision_tower.stages:
	s.grad_checkpointing = gradient_checkpointing

	if self.input_image_size is not None:
	self.image_processor.size=self.input_image_size
	self.image_processor.crop_size={
	'height':self.input_image_size,
	'width': self.input_image_size
	}

	self.is_loaded = True

	def feature_select(self, image_forward_outs):
	image_features = image_forward_outs[self.select_layer]
	return image_features

	def forward_features(self, x):
	x = self.vision_tower.stem(x)
	image_forward_out=[]
	for blk in self.vision_tower.stages:
	x = torch.utils.checkpoint.checkpoint(blk, x)
	b,c,h,w=x.shape
	image_forward_out.append(x.view(b,c,-1).transpose(1,2))
	return image_forward_out

	def forward(self, images):
	if self.freeze_vision:
	with torch.no_grad():
	image_features = self._forward_images(images)
	else:
	image_features = self._forward_images(images)

	return image_features

	def _forward_images(self, images):

	image_forward_outs = self.forward_features(images.to(device=self.device, dtype=self.dtype))
	image_features = self.feature_select(image_forward_outs)

	return image_features

	@property
	def dummy_feature(self):
	return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)

	@property
	def dtype(self):
	return next(self.vision_tower.parameters()).dtype

	@property
	def device(self):
	return next(self.vision_tower.parameters()).device

	@property
	def config(self):
	assert NotImplementedError
	pass

	@property
	def num_attention_heads(self):
	# as constant
	return 16
	@property
	def num_layers(self):
	# as constant
	return 4
	@property
	def hidden_size(self):
	return self.vision_tower.hidden_size

	@property
	def num_patches(self):
	return (cfg['image_size'] // self.patch_embed.patch_size[0]) ** 2