RuiTerrty/RemoteSensingChangeDetection-RSCD.HA2F

1# Copyright (c) Meta Platforms, Inc. and affiliates.2#3# This source code is licensed under the Apache License, Version 2.04# found in the LICENSE file in the root directory of this source tree.5 6# References:7#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py8#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py9 10from typing import Callable, Optional, Tuple, Union11 12from torch import Tensor13import torch.nn as nn14 15 16def make_2tuple(x):17    if isinstance(x, tuple):18        assert len(x) == 219        return x20 21    assert isinstance(x, int)22    return (x, x)23 24 25class PatchEmbed(nn.Module):26    """27    2D image to patch embedding: (B,C,H,W) -> (B,N,D)28 29    Args:30        img_size: Image size.31        patch_size: Patch token size.32        in_chans: Number of input image channels.33        embed_dim: Number of linear projection output channels.34        norm_layer: Normalization layer.35    """36 37    def __init__(38        self,39        img_size: Union[int, Tuple[int, int]] = 224,40        patch_size: Union[int, Tuple[int, int]] = 16,41        in_chans: int = 3,42        embed_dim: int = 768,43        norm_layer: Optional[Callable] = None,44        flatten_embedding: bool = True,45    ) -> None:46        super().__init__()47 48        image_HW = make_2tuple(img_size)49        patch_HW = make_2tuple(patch_size)50        patch_grid_size = (51            image_HW[0] // patch_HW[0],52            image_HW[1] // patch_HW[1],53        )54 55        self.img_size = image_HW56        self.patch_size = patch_HW57        self.patches_resolution = patch_grid_size58        self.num_patches = patch_grid_size[0] * patch_grid_size[1]59 60        self.in_chans = in_chans61        self.embed_dim = embed_dim62 63        self.flatten_embedding = flatten_embedding64 65        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)66        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()67 68    def forward(self, x: Tensor) -> Tensor:69        _, _, H, W = x.shape70        patch_H, patch_W = self.patch_size71 72        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"73        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"74 75        x = self.proj(x)  # B C H W76        H, W = x.size(2), x.size(3)77        x = x.flatten(2).transpose(1, 2)  # B HW C78        x = self.norm(x)79        if not self.flatten_embedding:80            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C81        return x82 83    def flops(self) -> float:84        Ho, Wo = self.patches_resolution85        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])86        if self.norm is not None:87            flops += Ho * Wo * self.embed_dim88        return flops89