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 6import os7from typing import Callable, Optional8import warnings9 10from torch import Tensor, nn11import torch.nn.functional as F12 13 14class SwiGLUFFN(nn.Module):15    def __init__(16        self,17        in_features: int,18        hidden_features: Optional[int] = None,19        out_features: Optional[int] = None,20        act_layer: Callable[..., nn.Module] = None,21        drop: float = 0.0,22        bias: bool = True,23    ) -> None:24        super().__init__()25        out_features = out_features or in_features26        hidden_features = hidden_features or in_features27        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)28        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)29 30    def forward(self, x: Tensor) -> Tensor:31        x12 = self.w12(x)32        x1, x2 = x12.chunk(2, dim=-1)33        hidden = F.silu(x1) * x234        return self.w3(hidden)35 36 37XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None38try:39    if XFORMERS_ENABLED:40        from xformers.ops import SwiGLU41 42        XFORMERS_AVAILABLE = True43        warnings.warn("xFormers is available (SwiGLU)")44    else:45        warnings.warn("xFormers is disabled (SwiGLU)")46        raise ImportError47except ImportError:48    SwiGLU = SwiGLUFFN49    XFORMERS_AVAILABLE = False50 51    warnings.warn("xFormers is not available (SwiGLU)")52 53 54class SwiGLUFFNFused(SwiGLU):55    def __init__(56        self,57        in_features: int,58        hidden_features: Optional[int] = None,59        out_features: Optional[int] = None,60        act_layer: Callable[..., nn.Module] = None,61        drop: float = 0.0,62        bias: bool = True,63    ) -> None:64        out_features = out_features or in_features65        hidden_features = hidden_features or in_features66        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 867        super().__init__(68            in_features=in_features,69            hidden_features=hidden_features,70            out_features=out_features,71            bias=bias,72        )73