TTXian/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 torch7import torch.nn as nn8from torch.nn.init import trunc_normal_9from torch.nn.utils import weight_norm10 11 12class DINOHead(nn.Module):13    def __init__(14        self,15        in_dim,16        out_dim,17        use_bn=False,18        nlayers=3,19        hidden_dim=2048,20        bottleneck_dim=256,21        mlp_bias=True,22    ):23        super().__init__()24        nlayers = max(nlayers, 1)25        self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias)26        self.apply(self._init_weights)27        self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))28        self.last_layer.weight_g.data.fill_(1)29 30    def _init_weights(self, m):31        if isinstance(m, nn.Linear):32            trunc_normal_(m.weight, std=0.02)33            if isinstance(m, nn.Linear) and m.bias is not None:34                nn.init.constant_(m.bias, 0)35 36    def forward(self, x):37        x = self.mlp(x)38        eps = 1e-6 if x.dtype == torch.float16 else 1e-1239        x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)40        x = self.last_layer(x)41        return x42 43 44def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True):45    if nlayers == 1:46        return nn.Linear(in_dim, bottleneck_dim, bias=bias)47    else:48        layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]49        if use_bn:50            layers.append(nn.BatchNorm1d(hidden_dim))51        layers.append(nn.GELU())52        for _ in range(nlayers - 2):53            layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))54            if use_bn:55                layers.append(nn.BatchNorm1d(hidden_dim))56            layers.append(nn.GELU())57        layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))58        return nn.Sequential(*layers)59