Duplicate from InPeerReview/RemoteSensingChangeDetection-RSCD.HA2F

Co-authored-by: Agentic Motion Group (AMG) <InPeerReview@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,67 @@
+## 🛠️ Requirements
+
+### Environment
+- **Linux system**,
+- **Python** 3.8+, recommended 3.10
+- **PyTorch** 2.0 or higher, recommended 2.1.0
+- **CUDA** 11.7 or higher, recommended 12.1
+
+### Environment Installation
+
+It is recommended to use Miniconda for installation. The following commands will create a virtual environment named `stnr` and install PyTorch. In the following installation steps, the default installed CUDA version is 12.1. If your CUDA version is not 12.1, please modify it according to the actual situation.
+
+```bash
+# Create conda environment
+conda create -n stnr python=3.8 -y
+conda activate stnr
+
+# Install PyTorch
+pip install -r requirements.txt
+```
+
+## 📁 Dataset Preparation
+
+We evaluate our method on five remote sensing change detection datasets: **WHU-CD**, **LEVIR-CD**, **SYSU-CD**.
+
+| Dataset | Link |
+|---------|------|
+| WHU-CD | [Download](https://aistudio.baidu.com/datasetdetail/251669) |
+| LEVIR-CD | [Download](https://opendatalab.org.cn/OpenDataLab/LEVIR-CD) |
+| SYSU-CD | [Download](https://mail2sysueducn-my.sharepoint.com/personal/liumx23_mail2_sysu_edu_cn/_layouts/15/onedrive.aspx?id=%2Fpersonal%2Fliumx23%5Fmail2%5Fsysu%5Fedu%5Fcn%2FDocuments%2FSYSU%2DCD&ga=1) |
+
+
+
+### Example of Training on LEVIR-CD Dataset
+
+```bash
+python main.py --file_root LEVIR --max_steps 80000 --model_type small --batch_size 16 --lr 2e-4 --gpu_id 0
+```
+
+### Example of Training on LEVIR-CD Dataset
+
+```bash
+python eval.py --file_root LEVIR --max_steps 80000 --model_type small --batch_size 16 --lr 2e-4 --gpu_id 0
+```
+
+## 📂 DATA Structure
+
+```
+├─Train
+      ├─A          jpg/png
+      ├─B          jpg/png
+      └─label      jpg/png
+  ├─Val
+      ├─A 
+      ├─B
+      └─label
+  ├─Test
+      ├─A
+      ├─B
+      └─label
+```
+
+## 🙏 Acknowledgement
+
+We sincerely thank the following works for their contributions:
+
+- [ChangeViT](https://arxiv.org/pdf/2406.12847) – A state-of-the-art method for remote sensing change detection that inspired and influenced parts of this work.

dataset/Transforms.py

@@ -0,0 +1,217 @@
+import numpy
+import numpy as np
+import torch
+import random
+import cv2
+
+
+class Scale(object):
+    """
+    Resize the given image to a fixed scale
+    """
+
+    def __init__(self, wi, he):
+        '''
+        :param wi: width after resizing
+        :param he: height after reszing
+        '''
+        self.w = wi
+        self.h = he
+
+    # modified from torchvision to add support for max size
+
+    def __call__(self, img, label):
+        '''
+        :param img: RGB image
+        :param label: semantic label image
+        :return: resized images
+        '''
+        # bilinear interpolation for RGB image
+        img = cv2.resize(img, (self.w, self.h))
+        # nearest neighbour interpolation for label image
+        label = cv2.resize(label, (self.w, self.h), interpolation=cv2.INTER_NEAREST)
+        return [img, label]
+
+
+class Resize(object):
+    def __init__(self, min_size, max_size, strict=False):
+        if not isinstance(min_size, (list, tuple)):
+            min_size = (min_size,)
+        self.min_size = min_size
+        self.max_size = max_size
+        self.strict = strict
+
+    # modified from torchvision to add support for max size
+    def get_size(self, image_size):
+        w, h = image_size
+        if not self.strict:
+            size = random.choice(self.min_size)
+            max_size = self.max_size
+            if max_size is not None:
+                min_original_size = float(min((w, h)))
+                max_original_size = float(max((w, h)))
+                if max_original_size / min_original_size * size > max_size:
+                    size = int(round(max_size * min_original_size / max_original_size))
+
+            if (w <= h and w == size) or (h <= w and h == size):
+                return (h, w)
+
+            if w < h:
+                ow = size
+                oh = int(size * h / w)
+            else:
+                oh = size
+                ow = int(size * w / h)
+
+            return (oh, ow)
+        else:
+            if w < h:
+                return (self.max_size, self.min_size[0])
+            else:
+                return (self.min_size[0], self.max_size)
+
+    def __call__(self, image, label):
+        size = self.get_size(image.shape[:2])
+        image = cv2.resize(image, size)
+        # I confirm that the output size is right, not reversed
+        label = cv2.resize(label, size, interpolation=cv2.INTER_NEAREST)
+        return (image, label)
+
+
+class RandomCropResize(object):
+    """
+    Randomly crop and resize the given image with a probability of 0.5
+    """
+
+    def __init__(self, crop_area):
+        '''
+        :param crop_area: area to be cropped (this is the max value and we select between 0 and crop area
+        '''
+        self.cw = crop_area
+        self.ch = crop_area
+
+    def __call__(self, img, label):
+        if random.random() < 0.5:
+            h, w = img.shape[:2]
+            x1 = random.randint(0, self.ch)
+            y1 = random.randint(0, self.cw)
+
+            img_crop = img[y1:h - y1, x1:w - x1]
+            label_crop = label[y1:h - y1, x1:w - x1]
+
+            img_crop = cv2.resize(img_crop, (w, h))
+            label_crop = cv2.resize(label_crop, (w, h), interpolation=cv2.INTER_NEAREST)
+
+            return img_crop, label_crop
+        else:
+            return [img, label]
+
+
+class RandomFlip(object):
+    """
+    Randomly flip the given Image with a probability of 0.5
+    """
+
+    def __call__(self, image, label):
+        if random.random() < 0.5:
+                image = cv2.flip(image, 0)  # horizontal flip
+                label = cv2.flip(label, 0)  # horizontal flip
+        if random.random() < 0.5:
+                image = cv2.flip(image, 1)  # veritcal flip
+                label = cv2.flip(label, 1)  # veritcal flip
+        return [image, label]
+
+
+class RandomExchange(object):
+    """
+    Randomly flip the given Image with a probability of 0.5
+    """
+
+    def __call__(self, image, label):
+        if random.random() < 0.5:
+            pre_img = image[:, :, 0:3]
+            post_img = image[:, :, 3:6]
+            image = numpy.concatenate((post_img, pre_img), axis=2)
+        return [image, label]
+
+
+class Normalize(object):
+    """
+    Given mean: (B, G, R) and std: (B, G, R),
+    will normalize each channel of the torch.*Tensor, i.e.
+    channel = (channel - mean) / std
+    """
+
+    def __init__(self, mean, std):
+        '''
+        :param mean: global mean computed from dataset
+        :param std: global std computed from dataset
+        '''
+        self.mean = mean
+        self.std = std
+        self.depth_mean = [0.5]
+        self.depth_std = [0.5]
+
+    def __call__(self, image, label):
+        image = image.astype(np.float32)
+        image = image / 255
+        label = np.ceil(label / 255)
+        for i in range(6):
+            image[:, :, i] -= self.mean[i]
+        for i in range(6):
+            image[:, :, i] /= self.std[i]
+
+        return [image, label]
+
+
+class GaussianNoise(object):
+    def __init__(self, std=0.05):
+        '''
+        :param mean: global mean computed from dataset
+        :param std: global std computed from dataset
+        '''
+        self.std = std
+
+    def __call__(self, image, label):
+        noise = np.random.normal(loc=0, scale=self.std, size=image.shape)
+        image = image + noise.astype(np.float32)
+        return [image, label]
+
+
+class ToTensor(object):
+    '''
+    This class converts the data to tensor so that it can be processed by PyTorch
+    '''
+
+    def __init__(self, scale=1):
+        '''
+        :param scale: set this parameter according to the output scale
+        '''
+        self.scale = scale
+
+    def __call__(self, image, label):
+        if self.scale != 1:
+            h, w = label.shape[:2]
+            image = cv2.resize(image, (int(w), int(h)))
+            label = cv2.resize(label, (int(w / self.scale), int(h / self.scale)), \
+                               interpolation=cv2.INTER_NEAREST)
+        image = image[:, :, ::-1].copy()  # .copy() is to solve "torch does not support negative index"
+        image = image.transpose((2, 0, 1))
+        image_tensor = torch.from_numpy(image)
+        label_tensor = torch.LongTensor(np.array(label, dtype=np.int)).unsqueeze(dim=0)
+
+        return [image_tensor, label_tensor]
+
+
+class Compose(object):
+    """
+    Composes several transforms together.
+    """
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, *args):
+        for t in self.transforms:
+            args = t(*args)
+        return args

dataset/dataset.py

@@ -0,0 +1,39 @@
+import cv2
+import os
+from os.path import join as osp
+import numpy
+import torch.utils.data
+
+
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self, file_root='data/', mode='train', transform=None):
+        self.file_list = os.listdir(osp(file_root, mode, 'A'))
+
+        self.pre_images = [osp(file_root, mode, 'A', x) for x in self.file_list]
+        self.post_images = [osp(file_root, mode, 'B', x) for x in self.file_list]
+        self.gts = [osp(file_root, mode, 'label', x) for x in self.file_list]
+
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.pre_images)
+
+    def __getitem__(self, idx):
+        pre_image_name = self.pre_images[idx]
+        label_name = self.gts[idx]
+        post_image_name = self.post_images[idx]
+
+        pre_image = cv2.imread(pre_image_name)
+        label = cv2.imread(label_name, 0)
+        post_image = cv2.imread(post_image_name)
+
+        img = numpy.concatenate((pre_image, post_image), axis=2)
+
+        if self.transform:
+            [img, label] = self.transform(img, label)
+
+        return img, label
+
+    def get_img_info(self, idx):
+        img = cv2.imread(self.pre_images[idx])
+        return {"height": img.shape[0], "width": img.shape[1]}

eval.py

@@ -0,0 +1,222 @@
+import sys
+from model.trainer import Trainer
+
+sys.path.insert(0, '.')
+
+import torch
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.nn.parallel import gather
+import torch.optim.lr_scheduler
+
+import dataset.dataset as myDataLoader
+import dataset.Transforms as myTransforms
+from model.metric_tool import ConfuseMatrixMeter
+from model.utils import BCEDiceLoss, init_seed
+from PIL import Image
+import os
+import time
+import numpy as np
+from argparse import ArgumentParser
+from tqdm import tqdm
+
+
+@torch.no_grad()
+def validate(args, val_loader, model, save_masks=False):
+    model.eval()
+
+    # 确保所有BatchNorm层使用全局统计量
+    for m in model.modules():
+        if isinstance(m, (torch.nn.BatchNorm2d, torch.nn.BatchNorm1d)):
+            m.track_running_stats = True
+            m.eval()
+
+    salEvalVal = ConfuseMatrixMeter(n_class=2)
+    epoch_loss = []
+
+    if save_masks:
+        mask_dir = f"{args.savedir}/pred_masks"
+        os.makedirs(mask_dir, exist_ok=True)
+        print(f"Saving prediction masks to: {mask_dir}")
+
+    pbar = tqdm(enumerate(val_loader), total=len(val_loader), desc="Validating")
+
+    for batch_idx, batched_inputs in pbar:
+        img, target = batched_inputs
+        # 获取当前batch的所有文件名
+        batch_file_names = val_loader.sampler.data_source.file_list[
+            batch_idx * args.batch_size : (batch_idx + 1) * args.batch_size
+        ]
+        
+        pre_img = img[:, 0:3]
+        post_img = img[:, 3:6]
+
+        if args.onGPU:
+            pre_img = pre_img.cuda()
+            post_img = post_img.cuda()
+            target = target.cuda()
+
+        target = target.float()
+        output = model(pre_img, post_img)
+        loss = BCEDiceLoss(output, target)
+        pred = (output > 0.5).long()
+
+        if save_masks:
+            pred_np = pred.cpu().numpy().astype(np.uint8)
+            
+            print(f"\nDebug - Batch {batch_idx}: {len(batch_file_names)} files, Mask shape: {pred_np.shape}")
+            
+            try:
+                for i in range(pred_np.shape[0]):
+                    if i >= len(batch_file_names):  # 防止文件名不足
+                        print(f"Warning: Missing filename for mask {i}, using default")
+                        base_name = f"batch_{batch_idx}_mask_{i}"
+                    else:
+                        base_name = os.path.splitext(os.path.basename(batch_file_names[i]))[0]
+                    
+                    single_mask = pred_np[i, 0]  # 获取(1, 256, 256)中的(256, 256)
+                    
+                    if single_mask.ndim != 2:
+                        raise ValueError(f"Invalid mask shape: {single_mask.shape}")
+                    
+                    mask_path = f"{mask_dir}/{base_name}_pred.png"
+                    Image.fromarray(single_mask * 255).save(mask_path)
+                    print(f"Saved: {mask_path}")
+
+            except Exception as e:
+                print(f"\nError saving batch {batch_idx}: {str(e)}")
+                print(f"Current mask shape: {single_mask.shape if 'single_mask' in locals() else 'N/A'}")
+                print(f"Current file: {base_name if 'base_name' in locals() else 'N/A'}")
+
+        if args.onGPU and torch.cuda.device_count() > 1:
+            pred = gather(pred, 0, dim=0)
+
+        f1 = salEvalVal.update_cm(pr=pred.cpu().numpy(), gt=target.cpu().numpy())
+        epoch_loss.append(loss.item())
+
+        pbar.set_postfix({'Loss': f"{loss.item():.4f}", 'F1': f"{f1:.4f}"})
+
+    average_loss = sum(epoch_loss) / len(epoch_loss)
+    scores = salEvalVal.get_scores()
+    return average_loss, scores
+    
+def ValidateSegmentation(args):
+    """完整的验证流程主函数"""
+    # 初始化设置
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
+    torch.backends.cudnn.benchmark = True
+    init_seed(args.seed)  # 固定随机种子保证可重复性
+
+    # 模型路径设置
+    args.savedir = os.path.join(args.savedir,
+                                f"{args.file_root}_iter_{args.max_steps}_lr_{args.lr}")
+    os.makedirs(args.savedir, exist_ok=True)
+
+    # 数据集路径配置
+    dataset_mapping = {
+        'LEVIR': './levir_cd_256',
+        'WHU': './whu_cd_256',
+        'CLCD': './clcd_256',
+        'SYSU': './sysu_256',
+        'OSCD': './oscd_256'
+    }
+    args.file_root = dataset_mapping.get(args.file_root, args.file_root)
+
+    # 初始化模型
+    model = Trainer(args.model_type).float()
+    if args.onGPU:
+        model = model.cuda()
+
+    # 数据预处理 - 保持与训练时验证集相同的预处理
+    mean = [0.406, 0.456, 0.485, 0.406, 0.456, 0.485]
+    std = [0.225, 0.224, 0.229, 0.225, 0.224, 0.229]
+
+    valDataset = myTransforms.Compose([
+        myTransforms.Normalize(mean=mean, std=std),
+        myTransforms.Scale(args.inWidth, args.inHeight),
+        myTransforms.ToTensor()
+    ])
+
+    # 数据加载
+    test_data = myDataLoader.Dataset(file_root=args.file_root, mode="test", transform=valDataset)
+    testLoader = torch.utils.data.DataLoader(
+        test_data,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=True
+    )
+
+    # 日志设置
+    logFileLoc = os.path.join(args.savedir, args.logFile)
+    logger = open(logFileLoc, 'a' if os.path.exists(logFileLoc) else 'w')
+    if not os.path.exists(logFileLoc):
+        logger.write("\n%s\t%s\t%s\t%s\t%s\t%s\t%s" %
+                     ('Epoch', 'Kappa', 'IoU', 'F1', 'Recall', 'Precision', 'OA'))
+    logger.flush()
+
+    # 加载最佳模型
+    model_file_name = os.path.join(args.savedir, 'best_model.pth')
+    if not os.path.exists(model_file_name):
+        raise FileNotFoundError(f"Model file not found: {model_file_name}")
+
+    state_dict = torch.load(model_file_name)
+    model.load_state_dict(state_dict)
+    print(f"Loaded model from {model_file_name}")
+
+    # 执行验证
+    loss_test, score_test = validate(args, testLoader, model, save_masks=args.save_masks)
+
+    # 输出结果
+    print("\nTest Results:")
+    print(f"Loss: {loss_test:.4f}")
+    print(f"Kappa: {score_test['Kappa']:.4f}")
+    print(f"IoU: {score_test['IoU']:.4f}")
+    print(f"F1: {score_test['F1']:.4f}")
+    print(f"Recall: {score_test['recall']:.4f}")
+    print(f"Precision: {score_test['precision']:.4f}")
+    print(f"OA: {score_test['OA']:.4f}")
+
+    # 记录日志
+    logger.write("\n%s\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f" %
+                 ('Test', score_test['Kappa'], score_test['IoU'], score_test['F1'],
+                  score_test['recall'], score_test['precision'], score_test['OA']))
+    logger.close()
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument('--file_root', default="LEVIR",
+                        help='Data directory | LEVIR | WHU | CLCD | SYSU | OSCD')
+    parser.add_argument('--inWidth', type=int, default=256, help='Width of input image')
+    parser.add_argument('--inHeight', type=int, default=256, help='Height of input image')
+    parser.add_argument('--max_steps', type=int, default=80000,
+                        help='Max. number of iterations (for path naming)')
+    parser.add_argument('--num_workers', type=int, default=4,
+                        help='Number of data loading workers')
+    parser.add_argument('--model_type', type=str, default='small',
+                        help='Model type | tiny | small')
+    parser.add_argument('--batch_size', type=int, default=16,
+                        help='Batch size for validation')
+    parser.add_argument('--lr', type=float, default=2e-4,
+                        help='Learning rate (for path naming)')
+    parser.add_argument('--seed', type=int, default=16,
+                        help='Random seed for reproducibility')
+    parser.add_argument('--savedir', default='./results',
+                        help='Base directory to save results')
+    parser.add_argument('--logFile', default='testLog.txt',
+                        help='File to save validation logs')
+    parser.add_argument('--onGPU', default=True,
+                        type=lambda x: (str(x).lower() == 'true'),
+                        help='Run on GPU if True')
+    parser.add_argument('--gpu_id', type=int, default=0,
+                        help='GPU device id')
+    parser.add_argument('--save_masks', action='store_true',
+                        help='Save predicted masks to disk')
+
+    args = parser.parse_args()
+    print('Validation with args:')
+    print(args)
+
+    ValidateSegmentation(args)
+

main.py

@@ -0,0 +1,319 @@
+import sys
+
+from model.trainer import Trainer
+
+sys.path.insert(0, '.')
+
+import torch
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.nn.parallel import gather
+import torch.optim.lr_scheduler
+
+import dataset.dataset as myDataLoader
+import dataset.Transforms as myTransforms
+from model.metric_tool import ConfuseMatrixMeter
+from model.utils import BCEDiceLoss, init_seed, adjust_learning_rate
+
+import os, time
+import numpy as np
+from argparse import ArgumentParser
+
+
+
+@torch.no_grad()
+def val(args, val_loader, model):
+    model.eval()
+
+    salEvalVal = ConfuseMatrixMeter(n_class=2)
+
+    epoch_loss = []
+
+    total_batches = len(val_loader)
+    print(len(val_loader))
+    for iter, batched_inputs in enumerate(val_loader):
+
+        img, target = batched_inputs
+        pre_img = img[:, 0:3]
+        post_img = img[:, 3:6]
+
+        start_time = time.time()
+
+        if args.onGPU == True:
+            pre_img = pre_img.cuda()
+            target = target.cuda()
+            post_img = post_img.cuda()
+
+        pre_img_var = torch.autograd.Variable(pre_img).float()
+        post_img_var = torch.autograd.Variable(post_img).float()
+        target_var = torch.autograd.Variable(target).float()
+
+        # run the mdoel
+        output = model(pre_img_var, post_img_var)
+        loss = BCEDiceLoss(output, target_var)
+
+        pred = torch.where(output > 0.5, torch.ones_like(output), torch.zeros_like(output)).long()
+
+        # torch.cuda.synchronize()
+        time_taken = time.time() - start_time
+
+        epoch_loss.append(loss.data.item())
+
+        # compute the confusion matrix
+        if args.onGPU and torch.cuda.device_count() > 1:
+            output = gather(pred, 0, dim=0)
+        # salEvalVal.addBatch(pred, target_var)
+        f1 = salEvalVal.update_cm(pr=pred.cpu().numpy(), gt=target_var.cpu().numpy())
+        if iter % 5 == 0:
+            print('\r[%d/%d] F1: %3f loss: %.3f time: %.3f' % (iter, total_batches, f1, loss.data.item(), time_taken),
+                  end='')
+
+    average_epoch_loss_val = sum(epoch_loss) / len(epoch_loss)
+    scores = salEvalVal.get_scores()
+
+    return average_epoch_loss_val, scores
+
+
+def train(args, train_loader, model, optimizer, epoch, max_batches, cur_iter=0, lr_factor=1.):
+    # switch to train mode
+    model.train()
+
+    salEvalVal = ConfuseMatrixMeter(n_class=2)
+    epoch_loss = []
+
+    for iter, batched_inputs in enumerate(train_loader):
+
+        img, target = batched_inputs
+        pre_img = img[:, 0:3]
+        post_img = img[:, 3:6]
+
+        start_time = time.time()
+
+        # adjust the learning rate
+        lr = adjust_learning_rate(args, optimizer, epoch, iter + cur_iter, max_batches, lr_factor=lr_factor)
+
+        if args.onGPU == True:
+            pre_img = pre_img.cuda()
+            target = target.cuda()
+            post_img = post_img.cuda()
+
+        pre_img_var = torch.autograd.Variable(pre_img).float()
+        post_img_var = torch.autograd.Variable(post_img).float()
+        target_var = torch.autograd.Variable(target).float()
+
+        # run the model
+        output = model(pre_img_var, post_img_var)
+        loss = BCEDiceLoss(output, target_var)
+
+        pred = torch.where(output > 0.5, torch.ones_like(output), torch.zeros_like(output)).long()
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        epoch_loss.append(loss.data.item())
+        time_taken = time.time() - start_time
+        res_time = (max_batches * args.max_epochs - iter - cur_iter) * time_taken / 3600
+
+        if args.onGPU and torch.cuda.device_count() > 1:
+            output = gather(pred, 0, dim=0)
+
+        # Computing F-measure and IoU on GPU
+        with torch.no_grad():
+            f1 = salEvalVal.update_cm(pr=pred.cpu().numpy(), gt=target_var.cpu().numpy())
+
+        if iter % 5 == 0:
+            print('\riteration: [%d/%d] f1: %.3f lr: %.7f loss: %.3f time:%.3f h' % (
+                iter + cur_iter, max_batches * args.max_epochs, f1, lr, loss.data.item(),
+                res_time),
+                  end='')
+
+    average_epoch_loss_train = sum(epoch_loss) / len(epoch_loss)
+    scores = salEvalVal.get_scores()
+
+    return average_epoch_loss_train, scores, lr
+
+
+def trainValidateSegmentation(args):
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
+    
+    torch.backends.cudnn.benchmark = True
+
+    init_seed(args.seed)
+
+    args.savedir = args.savedir + '_' + args.file_root + '_iter_' + str(args.max_steps) + '_lr_' + str(args.lr) + '/'
+
+    if args.file_root == 'LEVIR':
+         args.file_root = './levir_cd_256'
+    elif args.file_root == 'WHU':
+        args.file_root = './whu_cd_256'
+    elif args.file_root == 'CLCD':
+        args.file_root = './clcd_256'
+    elif args.file_root == 'SYSU':
+        args.file_root = './sysu_256'
+    elif args.file_root == 'OSCD':
+        args.file_root = 'oscd_256'
+    else:
+        raise TypeError('%s has not defined' % args.file_root)
+
+    if not os.path.exists(args.savedir):
+        os.makedirs(args.savedir)
+
+
+    model = Trainer(args.model_type).float()
+    if args.onGPU:
+        model = model.cuda()
+
+    # mean = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
+    # std = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
+
+    mean = [0.406, 0.456, 0.485, 0.406, 0.456, 0.485]
+    std = [0.225, 0.224, 0.229, 0.225, 0.224, 0.229]
+
+    # compose the data with transforms
+    trainDataset_main = myTransforms.Compose([
+        myTransforms.Normalize(mean=mean, std=std),
+        myTransforms.Scale(args.inWidth, args.inHeight),
+        myTransforms.RandomCropResize(int(7. / 224. * args.inWidth)),
+        myTransforms.RandomFlip(),
+        myTransforms.RandomExchange(),
+        myTransforms.ToTensor()
+    ])
+
+    valDataset = myTransforms.Compose([
+        myTransforms.Normalize(mean=mean, std=std),
+        myTransforms.Scale(args.inWidth, args.inHeight),
+        myTransforms.ToTensor()
+    ])
+
+    train_data = myDataLoader.Dataset(file_root=args.file_root, mode="train", transform=trainDataset_main)
+
+    trainLoader = torch.utils.data.DataLoader(
+        train_data,
+        batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, pin_memory=True, drop_last=False
+    )
+
+    test_data = myDataLoader.Dataset(file_root=args.file_root, mode="test", transform=valDataset)
+    testLoader = torch.utils.data.DataLoader(
+        test_data, shuffle=False,
+        batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=True)
+
+
+    max_batches = len(trainLoader)
+    print('For each epoch, we have {} batches'.format(max_batches))
+
+    if args.onGPU:
+        cudnn.benchmark = True
+
+    args.max_epochs = int(np.ceil(args.max_steps / max_batches))
+    start_epoch = 0
+    cur_iter = 0
+    max_F1_val = 0
+
+    if args.resume is not None:
+        args.resume = args.savedir + 'checkpoint.pth.tar'
+        if os.path.isfile(args.resume):
+            print("=> loading checkpoint '{}'".format(args.resume))
+            checkpoint = torch.load(args.resume)
+            start_epoch = checkpoint['epoch']
+            cur_iter = start_epoch * len(trainLoader)
+            # args.lr = checkpoint['lr']
+            model.load_state_dict(checkpoint['state_dict'])
+            print("=> loaded checkpoint '{}' (epoch {})"
+                  .format(args.resume, checkpoint['epoch']))
+        else:
+            print("=> no checkpoint found at '{}'".format(args.resume))
+
+    logFileLoc = args.savedir + args.logFile
+    if os.path.isfile(logFileLoc):
+        logger = open(logFileLoc, 'a')
+    else:
+        logger = open(logFileLoc, 'w')
+        logger.write(
+            "\n%s\t%s\t%s\t%s\t%s\t%s\t%s" % ('Epoch', 'Kappa (val)', 'IoU (val)', 'F1 (val)', 'R (val)', 'P (val)', 'OA (val)'))
+    logger.flush()
+
+    optimizer = torch.optim.Adam(model.parameters(), args.lr, (0.9, 0.99), eps=1e-08, weight_decay=1e-4)
+
+    for epoch in range(start_epoch, args.max_epochs):
+        lossTr, score_tr, lr = \
+            train(args, trainLoader, model, optimizer, epoch, max_batches, cur_iter)
+        cur_iter += len(trainLoader)
+
+        torch.cuda.empty_cache()
+
+        # evaluate on validation set
+        if epoch == 0:
+            continue
+
+        lossVal, score_val = val(args, testLoader, model)
+        torch.cuda.empty_cache()
+        logger.write("\n%d\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f" % (epoch, score_val['Kappa'], score_val['IoU'],
+                                                                       score_val['F1'], score_val['recall'],
+                                                                       score_val['precision'], score_val['OA']))
+        logger.flush()
+
+        torch.save({
+            'epoch': epoch + 1,
+            'arch': str(model),
+            'state_dict': model.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'lossTr': lossTr,
+            'lossVal': lossVal,
+            'F_Tr': score_tr['F1'],
+            'F_val': score_val['F1'],
+            'lr': lr
+        }, args.savedir + 'checkpoint.pth.tar')
+
+        # save the model also
+        model_file_name = args.savedir + 'best_model.pth'
+        if epoch % 1 == 0 and max_F1_val <= score_val['F1']:
+            max_F1_val = score_val['F1']
+            torch.save(model.state_dict(), model_file_name)
+
+        print("Epoch " + str(epoch) + ': Details')
+        print("\nEpoch No. %d:\tTrain Loss = %.4f\tVal Loss = %.4f\t F1(tr) = %.4f\t F1(val) = %.4f" \
+              % (epoch, lossTr, lossVal, score_tr['F1'], score_val['F1']))
+        torch.cuda.empty_cache()
+
+    state_dict = torch.load(model_file_name)
+    model.load_state_dict(state_dict)
+
+    loss_test, score_test = val(args, testLoader, model)
+    print("\nTest :\t Kappa (te) = %.4f\t IoU (te) = %.4f\t F1 (te) = %.4f\t R (te) = %.4f\t P (te) = %.4f" \
+          % (score_test['Kappa'], score_test['IoU'], score_test['F1'], score_test['recall'], score_test['precision']))
+    logger.write("\n%s\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f" % ('Test', score_test['Kappa'], score_test['IoU'],
+                                                                   score_test['F1'], score_test['recall'],
+                                                                   score_test['precision'], score_test['OA']))
+    logger.flush()
+    logger.close()
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument('--file_root', default="LEVIR", help='Data directory | LEVIR | WHU | CLCD | SYSU | OSCD ')
+    parser.add_argument('--inWidth', type=int, default=256, help='Width of RGB image')
+    parser.add_argument('--inHeight', type=int, default=256, help='Height of RGB image')
+    parser.add_argument('--max_steps', type=int, default=80000, help='Max. number of iterations')
+    parser.add_argument('--num_workers', type=int, default=4, help='No. of parallel threads')
+    parser.add_argument('--model_type', type=str, default='small', help='select vit model type | tiny | small')
+    parser.add_argument('--batch_size', type=int, default=16, help='Batch size')
+    parser.add_argument('--step_loss', type=int, default=100, help='Decrease learning rate after how many epochs')
+    parser.add_argument('--lr', type=float, default=2e-4, help='Initial learning rate')
+    parser.add_argument('--lr_mode', default='poly', help='Learning rate policy, step or poly')
+    parser.add_argument('--seed', default=16, help='initialization seed number')
+    parser.add_argument('--savedir', default='./results', help='Directory to save the results')
+    parser.add_argument('--resume', default=None, help='Use this checkpoint to continue training | '
+                                                       './results_ep100/checkpoint.pth.tar')
+    parser.add_argument('--logFile', default='trainValLog.txt',
+                        help='File that stores the training and validation logs')
+    parser.add_argument('--onGPU', default=True, type=lambda x: (str(x).lower() == 'true'),
+                        help='Run on CPU or GPU. If TRUE, then GPU.')
+    parser.add_argument('--gpu_id', default=0, type=int, help='GPU id number')
+
+    args = parser.parse_args()
+    print('Called with args:')
+    print(args)
+
+    trainValidateSegmentation(args)

model/decoder.py

@@ -0,0 +1,301 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from model.utils import weight_init
+
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim1, dim2, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim1 // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.q = nn.Linear(dim1, dim1, bias=qkv_bias)
+        self.kv = nn.Linear(dim2, dim1 * 2, bias=qkv_bias)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim1, dim1)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, y):
+        B1, N1, C1 = x.shape
+        B2, N2, C2 = y.shape
+
+        q = self.q(x).reshape(B1, N1, self.num_heads, C1 // self.num_heads).permute(0, 2, 1, 3)
+        kv = self.kv(y).reshape(B2, N2, 2, self.num_heads, C1 // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        k, v = kv[0], kv[1]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B1, N1, C1)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+
+class Block(nn.Module):
+    def __init__(self, dim1, dim2, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim1)
+        self.norm2 = norm_layer(dim2)
+        self.attn = CrossAttention(dim1, dim2, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm3 = norm_layer(dim1)
+        mlp_hidden_dim = int(dim1 * mlp_ratio)
+        self.mlp = Mlp(in_features=dim1, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, y):
+        x = x + self.drop_path(self.attn(self.norm1(x), self.norm2(y)))
+        x = x + self.drop_path(self.mlp(self.norm3(x)))
+        return x
+
+
+
+class ContentAwareAggregation(nn.Module):
+    def __init__(self, low_dim, high_dim):
+        super().__init__()
+        self.project = nn.Sequential(
+            nn.Conv2d(high_dim, low_dim, kernel_size=1),
+            nn.BatchNorm2d(low_dim),
+            nn.ReLU(inplace=True)
+        )
+
+        self.attn_gen = nn.Sequential(
+            nn.Conv2d(low_dim, low_dim, kernel_size=3, padding=1, groups=low_dim),
+            nn.BatchNorm2d(low_dim),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(low_dim, low_dim, kernel_size=1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, low_feat, high_feat):
+        high_feat = F.interpolate(high_feat, size=low_feat.shape[2:], mode='bilinear', align_corners=False)
+        high_feat = self.project(high_feat)
+        attn = self.attn_gen(low_feat + high_feat)
+        out = attn * low_feat + high_feat
+        return out
+
+
+
+class FeatureInjector(nn.Module):
+    def __init__(self, dim1=384, dim2=[64, 128, 256], num_heads=8, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
+                    drop_path=0., act_layer=nn.ReLU, norm_layer=nn.LayerNorm):
+        super().__init__()
+
+        self.c2_c5 = Block(dim1, dim2[0], num_heads, mlp_ratio, qkv_bias, drop, attn_drop, drop_path, act_layer, norm_layer)
+        self.c3_c5 = Block(dim1, dim2[1], num_heads, mlp_ratio, qkv_bias, drop, attn_drop, drop_path, act_layer, norm_layer)
+        self.c4_c5 = Block(dim1, dim2[2], num_heads, mlp_ratio, qkv_bias, drop, attn_drop, drop_path, act_layer, norm_layer)
+
+        self.fuse = nn.Conv2d(dim1*3, dim1, 1, bias=False)
+        self.caa = ContentAwareAggregation(dim1, dim1)
+
+        weight_init(self)
+
+    def base_forward(self, c2, c3, c4, c5):
+        H, W = c5.shape[2:]
+
+        c2 = rearrange(c2, 'b c h w -> b (h w) c')
+        c3 = rearrange(c3, 'b c h w -> b (h w) c')
+        c4 = rearrange(c4, 'b c h w -> b (h w) c')
+        c5 = rearrange(c5, 'b c h w -> b (h w) c')
+
+        _c2 = self.c2_c5(c5, c2)
+        _c2 = rearrange(_c2, 'b (h w) c -> b c h w', h=H, w=W)
+
+        _c3 = self.c3_c5(c5, c3)
+        _c3 = rearrange(_c3, 'b (h w) c -> b c h w', h=H, w=W)
+
+        _c4 = self.c4_c5(c5, c4)
+        _c4 = rearrange(_c4, 'b (h w) c -> b c h w', h=H, w=W)
+
+        _c5 = self.fuse(torch.cat([_c2, _c3, _c4], dim=1))
+
+        return _c5
+
+    def forward(self, fx, fy):
+        _c5x = self.base_forward(fx[0], fx[1], fx[2], fx[3])
+        _c5y = self.base_forward(fy[0], fy[1], fy[2], fy[3])
+
+
+        _c5x = self.caa(_c5x, _c5y)
+        _c5y = self.caa(_c5y, _c5x)
+
+        return _c5x, _c5y
+
+
+class DualAttentionGate(nn.Module):
+    def __init__(self, channels, ratio=8):
+        super().__init__()
+        self.channel_att = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),  # [B,C,1,1]
+            nn.Conv2d(channels, channels // ratio, 1, bias=False),  # [B,C/8,1,1]
+            nn.ReLU(),
+            nn.Conv2d(channels // ratio, channels, 1, bias=False),  # [B,C,1,1]
+            nn.Sigmoid()
+        )
+
+        self.spatial_att = nn.Sequential(
+            nn.Conv2d(2, 1, 7, padding=3, bias=False),  # 输入2通道(mean+std)
+            nn.Sigmoid()  # 输出[B,1,H,W]
+        )
+
+    def forward(self, x):
+
+        c_att = self.channel_att(x)
+        mean = torch.mean(x, dim=1, keepdim=True)
+        std = torch.std(x, dim=1, keepdim=True)
+        s_att = self.spatial_att(torch.cat([mean, std], dim=1))
+
+
+        return x * c_att * s_att
+
+
+class SimplifiedFGFM(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.down = nn.Conv2d(in_channels, out_channels, 1, bias=False)
+        self.flow_make = nn.Conv2d(out_channels * 2, 4, 3, padding=1, bias=False)
+        self.dual_att = DualAttentionGate(out_channels)
+
+    def flow_warp(self, input, flow, size):
+
+        out_h, out_w = size
+        n, c, h, w = input.size()
+
+        norm = torch.tensor([[[[out_w, out_h]]]]).type_as(input).to(input.device)
+        grid = torch.meshgrid(
+            torch.linspace(-1.0, 1.0, out_h),
+            torch.linspace(-1.0, 1.0, out_w),
+            indexing='ij'
+        )
+        grid = torch.stack((grid[1], grid[0]), 2).repeat(n, 1, 1, 1).type_as(input)
+        grid = grid + flow.permute(0, 2, 3, 1) / norm
+
+        return F.grid_sample(input, grid, align_corners=True)
+
+    def forward(self, lowres_feature, highres_feature):
+
+        l_feature = self.down(lowres_feature)
+        l_feature_up = F.interpolate(l_feature, size=highres_feature.shape[2:], mode='bilinear', align_corners=True)
+
+        flow = self.flow_make(torch.cat([l_feature_up, highres_feature], dim=1))
+        flow_l, flow_h = flow[:, :2, :, :], flow[:, 2:, :, :]
+
+        l_warp = self.flow_warp(l_feature, flow_l, highres_feature.shape[2:])
+        h_warp = self.flow_warp(highres_feature, flow_h, highres_feature.shape[2:])
+
+
+        fused = self.dual_att(l_warp + h_warp)
+        return fused
+
+
+
+class Decoder(nn.Module):
+    def __init__(self, in_dim=[64, 128, 256, 384], decay=4, num_class=1):
+        super().__init__()
+        c2_channel, c3_channel, c4_channel, c5_channel = in_dim
+
+        self.structure_enhance = FeatureInjector(dim1=c5_channel)
+
+
+        self.fgfm_c4 = SimplifiedFGFM(in_channels=c5_channel, out_channels=c4_channel)
+        self.fgfm_c3 = SimplifiedFGFM(in_channels=c4_channel, out_channels=c3_channel)
+        self.fgfm_c2 = SimplifiedFGFM(in_channels=c3_channel, out_channels=c2_channel)
+
+
+        self.classfier = nn.Sequential(
+            nn.ConvTranspose2d(c2_channel, c2_channel, kernel_size=4, stride=2, padding=1),
+            nn.Conv2d(c2_channel, num_class, 3, 1, padding=1, bias=False)
+        )
+
+
+        self.mlp = nn.ModuleList([
+            nn.Sequential(
+                nn.Conv2d(dim * 3, dim // decay, 1, bias=False),
+                nn.BatchNorm2d(dim // decay),
+                nn.ReLU(),
+                nn.Conv2d(dim // decay, dim // decay, 3, 1, padding=1, bias=False),
+                nn.ReLU(),
+                nn.Conv2d(dim // decay, dim // decay, 3, 1, padding=1, bias=False),
+                nn.ReLU(),
+                nn.Conv2d(dim // decay, dim, 3, 1, padding=1, bias=False)
+            ) for dim in in_dim
+        ])
+
+    def difference_modeling(self, x, y, block):
+        f = torch.cat([x, y, torch.abs(x - y)], dim=1)
+        return block(f)
+
+    def forward(self, fx, fy):
+        c2x, c3x, c4x = fx[:-1]
+        c2y, c3y, c4y = fy[:-1]
+
+
+        c5x, c5y = self.structure_enhance(fx, fy)
+
+
+        c2 = self.difference_modeling(c2x, c2y, self.mlp[0])
+        c3 = self.difference_modeling(c3x, c3y, self.mlp[1])
+        c4 = self.difference_modeling(c4x, c4y, self.mlp[2])
+        c5 = self.difference_modeling(c5x, c5y, self.mlp[3])
+
+
+        c4f = self.fgfm_c4(c5, c4)
+        c3f = self.fgfm_c3(c4f, c3)
+        c2f = self.fgfm_c2(c3f, c2)
+
+
+        pred = self.classfier(c2f)
+        pred_mask = torch.sigmoid(pred)
+
+        return pred_mask

model/encoder.py

@@ -0,0 +1,391 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from torch.nn.init import trunc_normal_
+from einops import rearrange
+
+from model.layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
+from model.resnet import resnet18
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+    if not depth_first and include_root:
+        fn(module=module, name=name)
+    for child_name, child_module in module.named_children():
+        child_name = ".".join((name, child_name)) if name else child_name
+        named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+    if depth_first and include_root:
+        fn(module=module, name=name)
+    return module
+
+
+class BlockChunk(nn.ModuleList):
+    def forward(self, x):
+        for b in self:
+            x = b(x)
+        return x
+
+
+class DinoVisionTransformer(nn.Module):
+    def __init__(
+            self,
+            img_size=224,
+            patch_size=16,
+            in_chans=3,
+            embed_dim=768,
+            depth=12,
+            num_heads=12,
+            mlp_ratio=4.0,
+            qkv_bias=True,
+            ffn_bias=True,
+            proj_bias=True,
+            drop_path_rate=0.0,
+            drop_path_uniform=False,
+            init_values=None,  # for layerscale: None or 0 => no layerscale
+            embed_layer=PatchEmbed,
+            act_layer=nn.GELU,
+            block_fn=Block,
+            ffn_layer="mlp",
+            block_chunks=0,
+            num_register_tokens=0,
+            interpolate_antialias=False,
+            interpolate_offset=0.1,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            proj_bias (bool): enable bias for proj in attn if True
+            ffn_bias (bool): enable bias for ffn if True
+            drop_path_rate (float): stochastic depth rate
+            drop_path_uniform (bool): apply uniform drop rate across blocks
+            weight_init (str): weight init scheme
+            init_values (float): layer-scale init values
+            embed_layer (nn.Module): patch embedding layer
+            act_layer (nn.Module): MLP activation layer
+            block_fn (nn.Module): transformer block class
+            ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+            block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+            num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+            interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+            interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+        """
+        super().__init__()
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.interpolate_antialias = interpolate_antialias
+        self.interpolate_offset = interpolate_offset
+
+        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+        assert num_register_tokens >= 0
+        self.register_tokens = (
+            nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+        )
+
+        if drop_path_uniform is True:
+            dpr = [drop_path_rate] * depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        if ffn_layer == "mlp":
+            print("using MLP layer as FFN")
+            ffn_layer = Mlp
+        elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+            print("using SwiGLU layer as FFN")
+            ffn_layer = SwiGLUFFNFused
+        elif ffn_layer == "identity":
+            print("using Identity layer as FFN")
+
+            def f(*args, **kwargs):
+                return nn.Identity()
+
+            ffn_layer = f
+        else:
+            raise NotImplementedError
+
+        blocks_list = [
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                ffn_layer=ffn_layer,
+                init_values=init_values,
+            )
+            for i in range(depth)
+        ]
+        if block_chunks > 0:
+            self.chunked_blocks = True
+            chunked_blocks = []
+            chunksize = depth // block_chunks
+            for i in range(0, depth, chunksize):
+                # this is to keep the block index consistent if we chunk the block list
+                chunked_blocks.append([nn.Identity()] * i + blocks_list[i: i + chunksize])
+            self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+        else:
+            self.chunked_blocks = False
+            self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer(embed_dim)
+        self.head = nn.Identity()
+
+        self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+        self.init_weights()
+
+    def init_weights(self):
+        trunc_normal_(self.pos_embed, std=0.02)
+        if self.register_tokens is not None:
+            nn.init.normal_(self.register_tokens, std=1e-6)
+        named_apply(init_weights_vit_timm, self)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        previous_dtype = x.dtype
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1]
+        if npatch == N and w == h:
+            return self.pos_embed
+        patch_pos_embed = self.pos_embed.float()
+        dim = x.shape[-1]
+        w0 = w // self.patch_size
+        h0 = h // self.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset
+
+        sqrt_N = math.sqrt(N)
+        sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(sqrt_N), int(sqrt_N), dim).permute(0, 3, 1, 2),
+            scale_factor=(sx, sy),
+            mode="bicubic",
+            antialias=self.interpolate_antialias,
+        )
+
+        assert int(w0) == patch_pos_embed.shape[-2]
+        assert int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return patch_pos_embed.to(previous_dtype)
+
+    def prepare_tokens_with_masks(self, x, masks=None):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        if masks is not None:
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        if self.register_tokens is not None:
+            x = torch.cat(
+                (
+                    x[:, :1],
+                    self.register_tokens.expand(x.shape[0], -1, -1),
+                    x[:, 1:],
+                ),
+                dim=1,
+            )
+
+        return x
+
+    def forward_features_list(self, x_list, masks_list):
+        x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+        for blk in self.blocks:
+            x = blk(x)
+
+        all_x = x
+        output = []
+        for x, masks in zip(all_x, masks_list):
+            x_norm = self.norm(x)
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm[:, 0],
+                    "x_norm_regtokens": x_norm[:, 1: self.num_register_tokens + 1],
+                    "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1:],
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward(self, x, masks=None):
+        if isinstance(x, list):
+            return self.forward_features_list(x, masks)
+
+        x = self.prepare_tokens_with_masks(x, masks)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x_norm = self.norm(x)
+        return x_norm
+
+    def _get_intermediate_layers_not_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        # If n is an int, take the n last blocks. If it's a list, take them
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if i in blocks_to_take:
+                output.append(x)
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def _get_intermediate_layers_chunked(self, x, n=1):
+        x = self.prepare_tokens_with_masks(x)
+        output, i, total_block_len = [], 0, len(self.blocks[-1])
+        # If n is an int, take the n last blocks. If it's a list, take them
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+        for block_chunk in self.blocks:
+            for blk in block_chunk[i:]:  # Passing the nn.Identity()
+                x = blk(x)
+                if i in blocks_to_take:
+                    output.append(x)
+                i += 1
+        assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        return output
+
+    def get_intermediate_layers(
+            self,
+            x: torch.Tensor,
+            n: Union[int, Sequence] = 1,  # Layers or n last layers to take
+            reshape: bool = False,
+            return_class_token: bool = False,
+            norm=True,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        if self.chunked_blocks:
+            outputs = self._get_intermediate_layers_chunked(x, n)
+        else:
+            outputs = self._get_intermediate_layers_not_chunked(x, n)
+        if norm:
+            outputs = [self.norm(out) for out in outputs]
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
+        if reshape:
+            B, _, w, h = x.shape
+            outputs = [
+                out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+                for out in outputs
+            ]
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+    """ViT weight initialization, original timm impl (for reproducibility)"""
+    if isinstance(module, nn.Linear):
+        trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+class Encoder(nn.Module):
+    def __init__(self, model_type='small'):
+        super().__init__()
+        if model_type == 'tiny':
+            self.vit = DinoVisionTransformer(
+                img_size=256,
+                patch_size=16,
+                embed_dim=192,
+                depth=12,
+                num_heads=6,
+                mlp_ratio=4,
+                block_fn=partial(Block, attn_class=MemEffAttention),
+                num_register_tokens=0
+            )
+            path = "checkpoint/deit_tiny_patch16_224-a1311bcf.pth"
+
+        elif model_type == 'small':
+            self.vit = DinoVisionTransformer(
+                img_size=256,
+                patch_size=16,
+                embed_dim=384,
+                depth=12,
+                num_heads=6,
+                mlp_ratio=4,
+                block_fn=partial(Block, attn_class=MemEffAttention),
+                num_register_tokens=0
+            )
+            path = "checkpoint/dinov2_vits14_pretrain.pth"
+
+        else:
+            assert False, r'Encoder: check the vit model type'
+
+        state_dict = torch.load(path, map_location='cpu')['model'] \
+            if model_type == 'tiny' else torch.load(path, map_location='cpu')
+
+        for k in ['pos_embed', 'patch_embed.proj.weight']:
+            del state_dict[k]
+        msg = self.vit.load_state_dict(state_dict, strict=False)
+        print(' missing_keys:{},\n unexpected_keys:{}'.format(msg.missing_keys, msg.unexpected_keys))
+        print('model_type: {},\n checkpoint_path: {}'.format(model_type, path))
+
+        self.resnet = resnet18(pretrained=True)
+        self.drop = nn.Dropout(p=0.01)
+
+        # 新增特征融合模块
+        self.fusion_conv = nn.Sequential(
+            nn.Conv2d(512 + 384, 384, kernel_size=1),  # 假设ViT embed_dim=384
+            nn.BatchNorm2d(384),
+            nn.ReLU(inplace=True)
+        )
+
+    def detail_capture(self, x):
+        x = self.resnet.conv1(x)
+        x = self.resnet.bn1(x)
+        x = self.resnet.relu(x)
+
+        x2 = self.drop(self.resnet.layer1(x))
+        x3 = self.resnet.layer2(x2)
+        x4 = self.resnet.layer3(x3)
+        x5 = self.resnet.layer4(x4)
+        return [x2, x3, x4, x5]
+
+    def forward(self, x, y):
+
+        v_x = self.vit(x)
+        v_y = self.vit(y)
+
+        v_x = rearrange(v_x, 'b (h w) c -> b c h w', h=16, w=16)
+        v_y = rearrange(v_y, 'b (h w) c -> b c h w', h=16, w=16)
+
+        c_x = self.detail_capture(x)
+        c_y = self.detail_capture(y)
+
+        fused_v_x = self.fusion_conv(torch.cat([c_x[-1], v_x], dim=1))
+        fused_v_y = self.fusion_conv(torch.cat([c_y[-1], v_y], dim=1))
+        return c_x[:-1] + [fused_v_x], c_y[:-1] + [fused_v_y]

model/layers/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+from .dino_head import DINOHead
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
+from .block import NestedTensorBlock
+from .attention import MemEffAttention

model/layers/attention.py

@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+import os
+import warnings
+
+from torch import Tensor
+from torch import nn
+
+
+logger = logging.getLogger("dinov2")
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import memory_efficient_attention, unbind
+
+        XFORMERS_AVAILABLE = True
+        warnings.warn("xFormers is available (Attention)")
+    else:
+        warnings.warn("xFormers is disabled (Attention)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+    warnings.warn("xFormers is not available (Attention)")
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+        attn = q @ k.transpose(-2, -1)
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MemEffAttention(Attention):
+    def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+        if not XFORMERS_AVAILABLE:
+            if attn_bias is not None:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return super().forward(x)
+
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+        q, k, v = unbind(qkv, 2)
+
+        x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+        x = x.reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

model/layers/block.py

@@ -0,0 +1,260 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+import os
+from typing import Callable, List, Any, Tuple, Dict
+import warnings
+
+import torch
+from torch import nn, Tensor
+
+from .attention import Attention, MemEffAttention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+logger = logging.getLogger("dinov2")
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import fmha, scaled_index_add, index_select_cat
+
+        XFORMERS_AVAILABLE = True
+        warnings.warn("xFormers is available (Block)")
+    else:
+        warnings.warn("xFormers is disabled (Block)")
+        raise ImportError
+except ImportError:
+    XFORMERS_AVAILABLE = False
+
+    warnings.warn("xFormers is not available (Block)")
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values=None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = Attention,
+        ffn_layer: Callable[..., nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.sample_drop_ratio = drop_path
+
+    def forward(self, x: Tensor) -> Tensor:
+        def attn_residual_func(x: Tensor) -> Tensor:
+            return self.ls1(self.attn(self.norm1(x)))
+
+        def ffn_residual_func(x: Tensor) -> Tensor:
+            return self.ls2(self.mlp(self.norm2(x)))
+
+        if self.training and self.sample_drop_ratio > 0.1:
+            # the overhead is compensated only for a drop path rate larger than 0.1
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+            x = drop_add_residual_stochastic_depth(
+                x,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+            )
+        elif self.training and self.sample_drop_ratio > 0.0:
+            x = x + self.drop_path1(attn_residual_func(x))
+            x = x + self.drop_path1(ffn_residual_func(x))  # FIXME: drop_path2
+        else:
+            x = x + attn_residual_func(x)
+            x = x + ffn_residual_func(x)
+        return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: Tensor,
+    residual_func: Callable[[Tensor], Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> Tensor:
+    # 1) extract subset using permutation
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    x_subset = x[brange]
+
+    # 2) apply residual_func to get residual
+    residual = residual_func(x_subset)
+
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+
+    residual_scale_factor = b / sample_subset_size
+
+    # 3) add the residual
+    x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+    b, n, d = x.shape
+    sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+    brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+    residual_scale_factor = b / sample_subset_size
+    return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+    if scaling_vector is None:
+        x_flat = x.flatten(1)
+        residual = residual.flatten(1)
+        x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+    else:
+        x_plus_residual = scaled_index_add(
+            x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+        )
+    return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+    """
+    this will perform the index select, cat the tensors, and provide the attn_bias from cache
+    """
+    batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+    all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+    if all_shapes not in attn_bias_cache.keys():
+        seqlens = []
+        for b, x in zip(batch_sizes, x_list):
+            for _ in range(b):
+                seqlens.append(x.shape[1])
+        attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+        attn_bias._batch_sizes = batch_sizes
+        attn_bias_cache[all_shapes] = attn_bias
+
+    if branges is not None:
+        cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+    else:
+        tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+        cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+    return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[Tensor],
+    residual_func: Callable[[Tensor, Any], Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> Tensor:
+    # 1) generate random set of indices for dropping samples in the batch
+    branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+    branges = [s[0] for s in branges_scales]
+    residual_scale_factors = [s[1] for s in branges_scales]
+
+    # 2) get attention bias and index+concat the tensors
+    attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+    # 3) apply residual_func to get residual, and split the result
+    residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+    outputs = []
+    for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+        outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+    return outputs
+
+
+class NestedTensorBlock(Block):
+    def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+        """
+        x_list contains a list of tensors to nest together and run
+        """
+        assert isinstance(self.attn, MemEffAttention)
+
+        if self.training and self.sample_drop_ratio > 0.0:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.mlp(self.norm2(x))
+
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=attn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            x_list = drop_add_residual_stochastic_depth_list(
+                x_list,
+                residual_func=ffn_residual_func,
+                sample_drop_ratio=self.sample_drop_ratio,
+                scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+            )
+            return x_list
+        else:
+
+            def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+            def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+                return self.ls2(self.mlp(self.norm2(x)))
+
+            attn_bias, x = get_attn_bias_and_cat(x_list)
+            x = x + attn_residual_func(x, attn_bias=attn_bias)
+            x = x + ffn_residual_func(x)
+            return attn_bias.split(x)
+
+    def forward(self, x_or_x_list):
+        if isinstance(x_or_x_list, Tensor):
+            return super().forward(x_or_x_list)
+        elif isinstance(x_or_x_list, list):
+            if not XFORMERS_AVAILABLE:
+                raise AssertionError("xFormers is required for using nested tensors")
+            return self.forward_nested(x_or_x_list)
+        else:
+            raise AssertionError

model/layers/dino_head.py

@@ -0,0 +1,58 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+from torch.nn.init import trunc_normal_
+from torch.nn.utils import weight_norm
+
+
+class DINOHead(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        use_bn=False,
+        nlayers=3,
+        hidden_dim=2048,
+        bottleneck_dim=256,
+        mlp_bias=True,
+    ):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        self.mlp = _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=hidden_dim, use_bn=use_bn, bias=mlp_bias)
+        self.apply(self._init_weights)
+        self.last_layer = weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+        self.last_layer.weight_g.data.fill_(1)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        eps = 1e-6 if x.dtype == torch.float16 else 1e-12
+        x = nn.functional.normalize(x, dim=-1, p=2, eps=eps)
+        x = self.last_layer(x)
+        return x
+
+
+def _build_mlp(nlayers, in_dim, bottleneck_dim, hidden_dim=None, use_bn=False, bias=True):
+    if nlayers == 1:
+        return nn.Linear(in_dim, bottleneck_dim, bias=bias)
+    else:
+        layers = [nn.Linear(in_dim, hidden_dim, bias=bias)]
+        if use_bn:
+            layers.append(nn.BatchNorm1d(hidden_dim))
+        layers.append(nn.GELU())
+        for _ in range(nlayers - 2):
+            layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias))
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+        layers.append(nn.Linear(hidden_dim, bottleneck_dim, bias=bias))
+        return nn.Sequential(*layers)

model/layers/drop_path.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    output = x * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

model/layers/layer_scale.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch import nn
+
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

model/layers/mlp.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x

model/layers/patch_embed.py

@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+#   https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+from typing import Callable, Optional, Tuple, Union
+
+from torch import Tensor
+import torch.nn as nn
+
+
+def make_2tuple(x):
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """
+    2D image to patch embedding: (B,C,H,W) -> (B,N,D)
+
+    Args:
+        img_size: Image size.
+        patch_size: Patch token size.
+        in_chans: Number of input image channels.
+        embed_dim: Number of linear projection output channels.
+        norm_layer: Normalization layer.
+    """
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int]] = 224,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+
+        image_HW = make_2tuple(img_size)
+        patch_HW = make_2tuple(patch_size)
+        patch_grid_size = (
+            image_HW[0] // patch_HW[0],
+            image_HW[1] // patch_HW[1],
+        )
+
+        self.img_size = image_HW
+        self.patch_size = patch_HW
+        self.patches_resolution = patch_grid_size
+        self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.flatten_embedding = flatten_embedding
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)  # B C H W
+        H, W = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)  # B HW C
+        x = self.norm(x)
+        if not self.flatten_embedding:
+            x = x.reshape(-1, H, W, self.embed_dim)  # B H W C
+        return x
+
+    def flops(self) -> float:
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops

model/layers/swiglu_ffn.py

@@ -0,0 +1,72 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import os
+from typing import Callable, Optional
+import warnings
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+    if XFORMERS_ENABLED:
+        from xformers.ops import SwiGLU
+
+        XFORMERS_AVAILABLE = True
+        warnings.warn("xFormers is available (SwiGLU)")
+    else:
+        warnings.warn("xFormers is disabled (SwiGLU)")
+        raise ImportError
+except ImportError:
+    SwiGLU = SwiGLUFFN
+    XFORMERS_AVAILABLE = False
+
+    warnings.warn("xFormers is not available (SwiGLU)")
+
+
+class SwiGLUFFNFused(SwiGLU):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[..., nn.Module] = None,
+        drop: float = 0.0,
+        bias: bool = True,
+    ) -> None:
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+        super().__init__(
+            in_features=in_features,
+            hidden_features=hidden_features,
+            out_features=out_features,
+            bias=bias,
+        )

model/metric_tool.py

@@ -0,0 +1,131 @@
+import numpy as np
+
+
+###################       metrics      ###################
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.initialized = False
+        self.val = None
+        self.avg = None
+        self.sum = None
+        self.count = None
+
+    def initialize(self, val, weight):
+        self.val = val
+        self.avg = val
+        self.sum = val * weight
+        self.count = weight
+        self.initialized = True
+
+    def update(self, val, weight=1):
+        if not self.initialized:
+            self.initialize(val, weight)
+        else:
+            self.add(val, weight)
+
+    def add(self, val, weight):
+        self.val = val
+        self.sum += val * weight
+        self.count += weight
+        self.avg = self.sum / self.count
+
+    def value(self):
+        return self.val
+
+    def average(self):
+        return self.avg
+
+    def get_scores(self):
+        scores_dict = cm2score(self.sum)
+        return scores_dict
+
+    def clear(self):
+        self.initialized = False
+
+
+###################      cm metrics      ###################
+class ConfuseMatrixMeter(AverageMeter):
+    """Computes and stores the average and current value"""
+
+    def __init__(self, n_class):
+        super(ConfuseMatrixMeter, self).__init__()
+        self.n_class = n_class
+
+    def update_cm(self, pr, gt, weight=1):
+        """获得当前混淆矩阵，并计算当前F1得分，并更新混淆矩阵"""
+        val = get_confuse_matrix(num_classes=self.n_class, label_gts=gt, label_preds=pr)
+        self.update(val, weight)
+        current_score = cm2F1(val)
+        return current_score
+
+    def get_scores(self):
+        scores_dict = cm2score(self.sum)
+        return scores_dict
+
+
+def harmonic_mean(xs):
+    harmonic_mean = len(xs) / sum((x + 1e-6) ** -1 for x in xs)
+    return harmonic_mean
+
+
+def cm2F1(confusion_matrix):
+    hist = confusion_matrix
+    tp = hist[1, 1]
+    fn = hist[1, 0]
+    fp = hist[0, 1]
+    tn = hist[0, 0]
+    # recall
+    recall = tp / (tp + fn + np.finfo(np.float32).eps)
+    # precision
+    precision = tp / (tp + fp + np.finfo(np.float32).eps)
+    # F1 score
+    f1 = 2 * recall * precision / (recall + precision + np.finfo(np.float32).eps)
+    return f1
+
+
+def cm2score(confusion_matrix):
+    hist = confusion_matrix
+    tp = hist[1, 1]
+    fn = hist[1, 0]
+    fp = hist[0, 1]
+    tn = hist[0, 0]
+    # acc
+    oa = (tp + tn) / (tp + fn + fp + tn + np.finfo(np.float32).eps)
+    # recall
+    recall = tp / (tp + fn + np.finfo(np.float32).eps)
+    # precision
+    precision = tp / (tp + fp + np.finfo(np.float32).eps)
+    # F1 score
+    f1 = 2 * recall * precision / (recall + precision + np.finfo(np.float32).eps)
+    # IoU
+    iou = tp / (tp + fp + fn + np.finfo(np.float32).eps)
+    # pre
+    pre = ((tp + fn) * (tp + fp) + (tn + fp) * (tn + fn)) / (tp + fp + tn + fn) ** 2
+    # kappa
+    kappa = (oa - pre) / (1 - pre)
+    score_dict = {'Kappa': kappa, 'IoU': iou, 'F1': f1, 'OA': oa, 'recall': recall, 'precision': precision, 'Pre': pre}
+    return score_dict
+
+
+def get_confuse_matrix(num_classes, label_gts, label_preds):
+    """计算一组预测的混淆矩阵"""
+
+    def __fast_hist(label_gt, label_pred):
+        """
+        Collect values for Confusion Matrix
+        For reference, please see: https://en.wikipedia.org/wiki/Confusion_matrix
+        :param label_gt: <np.array> ground-truth
+        :param label_pred: <np.array> prediction
+        :return: <np.ndarray> values for confusion matrix
+        """
+        mask = (label_gt >= 0) & (label_gt < num_classes)
+        hist = np.bincount(num_classes * label_gt[mask].astype(int) + label_pred[mask],
+                           minlength=num_classes ** 2).reshape(num_classes, num_classes)
+        return hist
+
+    confusion_matrix = np.zeros((num_classes, num_classes))
+    for lt, lp in zip(label_gts, label_preds):
+        confusion_matrix += __fast_hist(lt.flatten(), lp.flatten())
+    return confusion_matrix

model/resnet.py

@@ -0,0 +1,213 @@
+import torch.nn as nn
+import math
+import torch
+import torch.utils.model_zoo as model_zoo
+import torch.nn.functional as F
+from einops import rearrange
+
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+def resnet18(pretrained=False, **kwargs):
+    """Constructs a ResNet-18 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
+    return model
+
+
+def resnet34(pretrained=False, **kwargs):
+    """Constructs a ResNet-34 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+    return model
+
+
+def resnet50(pretrained=False, **kwargs):
+    """Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model
+
+
+def resnet101(pretrained=False, **kwargs):
+    """Constructs a ResNet-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+    return model
+
+
+def resnet152(pretrained=False, **kwargs):
+    """Constructs a ResNet-152 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+    return model
+
+
+if __name__ == '__main__':
+    m = resnet18(pretrained=True, vit_dim=768)
+    x = torch.rand(1, 3, 256, 256)
+    vit = [torch.rand(1, 256, 768), torch.rand(1, 256, 768), torch.rand(1, 256, 768)]
+    x2, x3, x4 = m(x, vit)
+    print(x2.shape, x3.shape, x4.shape)

model/trainer.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+
+from model.encoder import Encoder
+from model.decoder import Decoder
+
+from model.utils import weight_init
+
+
+class Trainer(nn.Module):
+    def __init__(self, model_type='small'):
+        super().__init__()
+        if model_type == 'tiny':
+            embed_dim = 192
+        elif model_type == 'small':
+            embed_dim = 384
+        else:
+            assert False, r'Trainer: check the vit model type'
+
+        self.encoder = Encoder(model_type)
+
+        self.decoder = Decoder(in_dim=[64, 128, 256, embed_dim])
+        weight_init(self.decoder)
+        
+    def forward(self, x, y):
+        fx, fy = self.encoder(x, y)
+        pred = self.decoder(fx, fy)
+
+        return pred
+        

model/utils.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+import torch.nn as nn
+import random
+
+
+def weight_init(module):
+    for n, m in module.named_children():
+        print('initialize: '+n)
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+            nn.init.ones_(m.weight)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Linear):
+            nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Sequential):
+            for f, g in m.named_children():
+                print('initialize: ' + f)
+                if isinstance(g, nn.Conv2d):
+                    nn.init.kaiming_normal_(g.weight, mode='fan_in', nonlinearity='relu')
+                    if g.bias is not None:
+                        nn.init.zeros_(g.bias)
+                elif isinstance(g, (nn.BatchNorm2d, nn.GroupNorm)):
+                    nn.init.ones_(g.weight)
+                    if g.bias is not None:
+                        nn.init.zeros_(g.bias)
+                elif isinstance(g, nn.Linear):
+                    nn.init.kaiming_normal_(g.weight, mode='fan_in', nonlinearity='relu')
+                    if g.bias is not None:
+                        nn.init.zeros_(g.bias)
+        elif isinstance(m, nn.AdaptiveAvgPool2d) or isinstance(m, nn.AdaptiveMaxPool2d) or isinstance(m, nn.ModuleList) or isinstance(m, nn.BCELoss):
+            a=1
+        else:
+            pass
+
+
+def init_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+
+def BCEDiceLoss(inputs, targets):
+    # print(inputs.shape, targets.shape)
+    bce = F.binary_cross_entropy(inputs, targets)
+    inter = (inputs * targets).sum()
+    eps = 1e-5
+    dice = (2 * inter + eps) / (inputs.sum() + targets.sum() + eps)
+    # print(bce.item(), inter.item(), inputs.sum().item(), dice.item())
+    return bce + 1 - dice
+
+
+def BCE(inputs, targets):
+    # print(inputs.shape, targets.shape)
+    bce = F.binary_cross_entropy(inputs, targets)
+    return bce
+
+
+def adjust_learning_rate(args, optimizer, epoch, iter, max_batches, lr_factor=1):
+    if args.lr_mode == 'step':
+        lr = args.lr * (0.1 ** (epoch // args.step_loss))
+    elif args.lr_mode == 'poly':
+        cur_iter = iter
+        max_iter = max_batches * args.max_epochs
+        lr = args.lr * (1 - cur_iter * 1.0 / max_iter) ** 0.9
+    else:
+        raise ValueError('Unknown lr mode {}'.format(args.lr_mode))
+    if epoch == 0 and iter < 200:
+        lr = args.lr * 0.9 * (iter + 1) / 200 + 0.1 * args.lr  # warm_up
+    lr *= lr_factor
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+    return lr

requirements.txt

@@ -0,0 +1,10 @@
+torch==1.13.1
+torchaudio==0.13.1
+torchcam==0.3.2
+torchgeo==0.4.1
+torchmetrics==0.11.4
+torchvision==0.14.1
+numpy==1.21.6
+Pillow==9.2.0
+einops==0.6.0
+opencv-python==4.6.0.66