Shortcuts

Source code for torchgeo.models.vit

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

"""Pre-trained Vision Transformer models."""

from typing import Any

import kornia.augmentation as K
import timm
import torch
from timm.models.vision_transformer import VisionTransformer
from torchvision.models._api import Weights, WeightsEnum

__all__ = ["ViTSmall16_Weights"]

# https://github.com/zhu-xlab/SSL4EO-S12/blob/d2868adfada65e40910bfcedfc49bc3b20df2248/src/benchmark/transfer_classification/linear_BE_moco.py#L167 # noqa: E501
# https://github.com/zhu-xlab/SSL4EO-S12/blob/d2868adfada65e40910bfcedfc49bc3b20df2248/src/benchmark/transfer_classification/datasets/EuroSat/eurosat_dataset.py#L97 # noqa: E501
# Normalization either by 10K or channel-wise with band statistics
_zhu_xlab_transforms = K.AugmentationSequential(
    K.Resize(256),
    K.CenterCrop(224),
    K.Normalize(mean=torch.tensor(0), std=torch.tensor(10000)),
    data_keys=None,
)

# https://github.com/microsoft/torchgeo/blob/8b53304d42c269f9001cb4e861a126dc4b462606/torchgeo/datamodules/ssl4eo_benchmark.py#L43 # noqa: E501
_ssl4eo_l_transforms = K.AugmentationSequential(
    K.Normalize(mean=torch.tensor(0), std=torch.tensor(255)),
    K.CenterCrop((224, 224)),
    data_keys=None,
)

# https://github.com/pytorch/vision/pull/6883
# https://github.com/pytorch/vision/pull/7107
# Can be removed once torchvision>=0.15 is required
Weights.__deepcopy__ = lambda *args, **kwargs: args[0]


[docs]class ViTSmall16_Weights(WeightsEnum): # type: ignore[misc] """Vision Transformer Small Patch Size 16 weights. For `timm <https://github.com/rwightman/pytorch-image-models>`_ *vit_small_patch16_224* implementation. .. versionadded:: 0.4 """ LANDSAT_TM_TOA_MOCO = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_tm_toa_moco-a1c967d8.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 7, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "moco", }, ) LANDSAT_TM_TOA_SIMCLR = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_tm_toa_simclr-7c2d9799.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 7, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "simclr", }, ) LANDSAT_ETM_TOA_MOCO = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_etm_toa_moco-26d19bcf.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 9, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "moco", }, ) LANDSAT_ETM_TOA_SIMCLR = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_etm_toa_simclr-34fb12cb.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 9, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "simclr", }, ) LANDSAT_ETM_SR_MOCO = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_etm_sr_moco-eaa4674e.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 6, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "moco", }, ) LANDSAT_ETM_SR_SIMCLR = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_etm_sr_simclr-a14c466a.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 6, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "simclr", }, ) LANDSAT_OLI_TIRS_TOA_MOCO = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_oli_tirs_toa_moco-c7c2cceb.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 11, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "moco", }, ) LANDSAT_OLI_TIRS_TOA_SIMCLR = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_oli_tirs_toa_simclr-ad43e9a4.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 11, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "simclr", }, ) LANDSAT_OLI_SR_MOCO = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_oli_sr_moco-c9b8898d.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 7, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "moco", }, ) LANDSAT_OLI_SR_SIMCLR = Weights( url="https://hf.co/torchgeo/ssl4eo_landsat/resolve/1c88bb51b6e17a21dde5230738fa38b74bd74f76/vits16_landsat_oli_sr_simclr-4e8f6102.pth", # noqa: E501 transforms=_ssl4eo_l_transforms, meta={ "dataset": "SSL4EO-L", "in_chans": 7, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2306.09424", "repo": "https://github.com/microsoft/torchgeo", "ssl_method": "simclr", }, ) SENTINEL2_ALL_DINO = Weights( url="https://hf.co/torchgeo/vit_small_patch16_224_sentinel2_all_dino/resolve/5b41dd418a79de47ac9f5be3e035405a83818a62/vit_small_patch16_224_sentinel2_all_dino-36bcc127.pth", # noqa: E501 transforms=_zhu_xlab_transforms, meta={ "dataset": "SSL4EO-S12", "in_chans": 13, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2211.07044", "repo": "https://github.com/zhu-xlab/SSL4EO-S12", "ssl_method": "dino", }, ) SENTINEL2_ALL_MOCO = Weights( url="https://hf.co/torchgeo/vit_small_patch16_224_sentinel2_all_moco/resolve/1cb683f6c14739634cdfaaceb076529adf898c74/vit_small_patch16_224_sentinel2_all_moco-67c9032d.pth", # noqa: E501 transforms=_zhu_xlab_transforms, meta={ "dataset": "SSL4EO-S12", "in_chans": 13, "model": "vit_small_patch16_224", "publication": "https://arxiv.org/abs/2211.07044", "repo": "https://github.com/zhu-xlab/SSL4EO-S12", "ssl_method": "moco", }, )
[docs]def vit_small_patch16_224( weights: ViTSmall16_Weights | None = None, *args: Any, **kwargs: Any ) -> VisionTransformer: """Vision Transform (ViT) small patch size 16 model. If you use this model in your research, please cite the following paper: * https://arxiv.org/abs/2010.11929 .. versionadded:: 0.4 Args: weights: Pre-trained model weights to use. *args: Additional arguments to pass to :func:`timm.create_model`. **kwargs: Additional keywork arguments to pass to :func:`timm.create_model`. Returns: A ViT small 16 model. """ if weights: kwargs["in_chans"] = weights.meta["in_chans"] model: VisionTransformer = timm.create_model( "vit_small_patch16_224", *args, **kwargs ) if weights: missing_keys, unexpected_keys = model.load_state_dict( weights.get_state_dict(progress=True), strict=False ) assert set(missing_keys) <= {"head.weight", "head.bias"} assert not unexpected_keys return model

© Copyright 2021, Microsoft Corporation. Revision acad7d47.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
v0.5.2
v0.5.1
v0.5.0
v0.4.1
v0.4.0
v0.3.1
v0.3.0
v0.2.1
v0.2.0
v0.1.1
v0.1.0
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources