LESSViT-Base / modeling_lessvit.py

Release LESSViT checkpoint and directly loadable encoder

586325c verified 1 day ago

1.47 kB

"""Hugging Face interface for the original LESSViT encoder."""
# Explicit dependencies let Transformers copy the full local/remote module tree.
from .pos_chan_embed import RopePositionChannelEmbedding
from .hyperspectral_patch_embed import HyperspectralPatchEmbed
from .low_rank_attention import LowRankAttention
from .attention_ops import LowRankBlock
from .spatial_spectral_low_rank_vit import (
    SpatialSpectralLowRankViTConfig,
    SpatialSpectralLowRankViTEncoder,
)


class LESSViTConfig(SpatialSpectralLowRankViTConfig):
    model_type = "lessvit"

    def __init__(self, channel_embed_dims_per_head=2, **kwargs):
        # The original config constructor recomputes these derived dimensions.
        kwargs.pop("channel_dim", None)
        kwargs.pop("spatial_dim", None)
        kwargs.pop("pretrain", None)
        kwargs.pop("mask_ratio", None)
        kwargs.pop("channel_mask_ratio", None)
        kwargs.pop("channel_dropout", None)
        super().__init__(channel_embed_dims_per_head=channel_embed_dims_per_head,
                         mask_ratio=0.0, channel_mask_ratio=0.0,
                         channel_dropout=None, **kwargs)
        self.channel_embed_dims_per_head = channel_embed_dims_per_head
        self.pretrain = False


class LESSViTModel(SpatialSpectralLowRankViTEncoder):
    config_class = LESSViTConfig
    main_input_name = "optical"

    def __init__(self, config):
        super().__init__(config)
        self.post_init()

	"""Hugging Face interface for the original LESSViT encoder."""
	# Explicit dependencies let Transformers copy the full local/remote module tree.
	from .pos_chan_embed import RopePositionChannelEmbedding
	from .hyperspectral_patch_embed import HyperspectralPatchEmbed
	from .low_rank_attention import LowRankAttention
	from .attention_ops import LowRankBlock
	from .spatial_spectral_low_rank_vit import (
	SpatialSpectralLowRankViTConfig,
	SpatialSpectralLowRankViTEncoder,
	)


	class LESSViTConfig(SpatialSpectralLowRankViTConfig):
	model_type = "lessvit"

	def __init__(self, channel_embed_dims_per_head=2, **kwargs):
	# The original config constructor recomputes these derived dimensions.
	kwargs.pop("channel_dim", None)
	kwargs.pop("spatial_dim", None)
	kwargs.pop("pretrain", None)
	kwargs.pop("mask_ratio", None)
	kwargs.pop("channel_mask_ratio", None)
	kwargs.pop("channel_dropout", None)
	super().__init__(channel_embed_dims_per_head=channel_embed_dims_per_head,
	mask_ratio=0.0, channel_mask_ratio=0.0,
	channel_dropout=None, **kwargs)
	self.channel_embed_dims_per_head = channel_embed_dims_per_head
	self.pretrain = False


	class LESSViTModel(SpatialSpectralLowRankViTEncoder):
	config_class = LESSViTConfig
	main_input_name = "optical"

	def __init__(self, config):
	super().__init__(config)
	self.post_init()