init

app.py

@@ -0,0 +1,398 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# pip uninstall nvidia_cublas_cu11
+
+import sys
+sys.path.append('..')
+import os
+os.system(f'pip install dlib')
+import torch
+import numpy as np
+from PIL import Image
+from torch.nn import functional as F
+
+import gradio as gr
+
+import models_vit
+from util.datasets import build_dataset
+import argparse
+from engine_finetune import test_all
+import dlib
+from huggingface_hub import hf_hub_download
+
+
+P = os.path.abspath(__file__)
+FRAME_SAVE_PATH = os.path.join(P[:-6], 'frame')
+CKPT_SAVE_PATH = os.path.join(P[:-6], 'checkpoints')
+CKPT_LIST = ['DfD Checkpoint_Fine-tuned on FF++',
+             'FAS Checkpoint_Fine-tuned on MCIO']
+CKPT_NAME = {'DfD Checkpoint_Fine-tuned on FF++': 'finetuned_models/FF++_c23_32frames/checkpoint-min_val_loss.pth',
+             'FAS Checkpoint_Fine-tuned on MCIO': 'finetuned_models/MCIO_protocol/Both_MCIO/checkpoint-min_val_loss.pth' }
+os.makedirs(FRAME_SAVE_PATH, exist_ok=True)
+os.makedirs(CKPT_SAVE_PATH, exist_ok=True)
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('MAE fine-tuning for image classification', add_help=False)
+    parser.add_argument('--batch_size', default=64, type=int,
+                        help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
+    parser.add_argument('--epochs', default=50, type=int)
+    parser.add_argument('--accum_iter', default=1, type=int,
+                        help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+
+    # Model parameters
+    parser.add_argument('--model', default='vit_large_patch16', type=str, metavar='MODEL',
+                        help='Name of model to train')
+
+    parser.add_argument('--input_size', default=224, type=int,
+                        help='images input size')
+    parser.add_argument('--normalize_from_IMN', action='store_true',
+                        help='cal mean and std from imagenet, else from pretrain datasets')
+    parser.set_defaults(normalize_from_IMN=True)
+    parser.add_argument('--apply_simple_augment', action='store_true',
+                        help='apply simple data augment')
+
+    parser.add_argument('--drop_path', type=float, default=0.1, metavar='PCT',
+                        help='Drop path rate (default: 0.1)')
+
+    # Optimizer parameters
+    parser.add_argument('--clip_grad', type=float, default=None, metavar='NORM',
+                        help='Clip gradient norm (default: None, no clipping)')
+    parser.add_argument('--weight_decay', type=float, default=0.05,
+                        help='weight decay (default: 0.05)')
+
+    parser.add_argument('--lr', type=float, default=None, metavar='LR',
+                        help='learning rate (absolute lr)')
+    parser.add_argument('--blr', type=float, default=1e-3, metavar='LR',
+                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
+    parser.add_argument('--layer_decay', type=float, default=0.75,
+                        help='layer-wise lr decay from ELECTRA/BEiT')
+
+    parser.add_argument('--min_lr', type=float, default=1e-6, metavar='LR',
+                        help='lower lr bound for cyclic schedulers that hit 0')
+
+    parser.add_argument('--warmup_epochs', type=int, default=5, metavar='N',
+                        help='epochs to warmup LR')
+
+    # Augmentation parameters
+    parser.add_argument('--color_jitter', type=float, default=None, metavar='PCT',
+                        help='Color jitter factor (enabled only when not using Auto/RandAug)')
+    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
+                        help='Use AutoAugment policy. "v0" or "original". " + "(default: rand-m9-mstd0.5-inc1)'),
+    parser.add_argument('--smoothing', type=float, default=0.1,
+                        help='Label smoothing (default: 0.1)')
+
+    # * Random Erase params
+    parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
+                        help='Random erase prob (default: 0.25)')
+    parser.add_argument('--remode', type=str, default='pixel',
+                        help='Random erase mode (default: "pixel")')
+    parser.add_argument('--recount', type=int, default=1,
+                        help='Random erase count (default: 1)')
+    parser.add_argument('--resplit', action='store_true', default=False,
+                        help='Do not random erase first (clean) augmentation split')
+
+    # * Mixup params
+    parser.add_argument('--mixup', type=float, default=0,
+                        help='mixup alpha, mixup enabled if > 0.')
+    parser.add_argument('--cutmix', type=float, default=0,
+                        help='cutmix alpha, cutmix enabled if > 0.')
+    parser.add_argument('--cutmix_minmax', type=float, nargs='+', default=None,
+                        help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
+    parser.add_argument('--mixup_prob', type=float, default=1.0,
+                        help='Probability of performing mixup or cutmix when either/both is enabled')
+    parser.add_argument('--mixup_switch_prob', type=float, default=0.5,
+                        help='Probability of switching to cutmix when both mixup and cutmix enabled')
+    parser.add_argument('--mixup_mode', type=str, default='batch',
+                        help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')
+
+    # * Finetuning params
+    parser.add_argument('--finetune', default='',
+                        help='finetune from checkpoint')
+    parser.add_argument('--global_pool', action='store_true')
+    parser.set_defaults(global_pool=True)
+    parser.add_argument('--cls_token', action='store_false', dest='global_pool',
+                        help='Use class token instead of global pool for classification')
+
+    # Dataset parameters
+    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
+                        help='dataset path')
+    parser.add_argument('--nb_classes', default=1000, type=int,
+                        help='number of the classification types')
+
+    parser.add_argument('--output_dir', default='',
+                        help='path where to save, empty for no saving')
+    parser.add_argument('--log_dir', default='',
+                        help='path where to tensorboard log')
+    parser.add_argument('--device', default='cuda',
+                        help='device to use for training / testing')
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--resume', default='',
+                        help='resume from checkpoint')
+
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='start epoch')
+    parser.add_argument('--eval', action='store_true',
+                        help='Perform evaluation only')
+    parser.set_defaults(eval=True)
+    parser.add_argument('--dist_eval', action='store_true', default=False,
+                        help='Enabling distributed evaluation (recommended during training for faster monitor')
+    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--pin_mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
+    parser.set_defaults(pin_mem=True)
+
+    # distributed training parameters
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='number of distributed processes')
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', action='store_true')
+    parser.add_argument('--dist_url', default='env://',
+                        help='url used to set up distributed training')
+
+    return parser
+
+
+args = get_args_parser()
+args = args.parse_args()
+args.nb_classes = 2
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = models_vit.__dict__['vit_base_patch16'](
+    num_classes=args.nb_classes,
+    drop_path_rate=args.drop_path,
+    global_pool=args.global_pool,
+)
+
+def load_model(ckpt):
+    if ckpt=='hoose from here':
+        return gr.update()
+    args.resume = os.path.join(CKPT_SAVE_PATH, ckpt)
+    if os.path.isfile(args.resume) == False:
+        hf_hub_download(local_dir=CKPT_SAVE_PATH,
+                        repo_id='Wolowolo/fsfm-3c',
+                        filename=ckpt)
+    checkpoint = torch.load(args.resume, map_location='cpu')
+    model.load_state_dict(checkpoint['model'])
+    return gr.update()
+
+
+def get_boundingbox(face, width, height, minsize=None):
+    """
+    From FF++:
+    https://github.com/ondyari/FaceForensics/blob/master/classification/detect_from_video.py
+    Expects a dlib face to generate a quadratic bounding box.
+    :param face: dlib face class
+    :param width: frame width
+    :param height: frame height
+    :param cfg.face_scale: bounding box size multiplier to get a bigger face region
+    :param minsize: set minimum bounding box size
+    :return: x, y, bounding_box_size in opencv form
+    """
+    x1 = face.left()
+    y1 = face.top()
+    x2 = face.right()
+    y2 = face.bottom()
+    size_bb = int(max(x2 - x1, y2 - y1) * 1.3)
+    if minsize:
+        if size_bb < minsize:
+            size_bb = minsize
+    center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
+
+    # Check for out of bounds, x-y top left corner
+    x1 = max(int(center_x - size_bb // 2), 0)
+    y1 = max(int(center_y - size_bb // 2), 0)
+    # Check for too big bb size for given x, y
+    size_bb = min(width - x1, size_bb)
+    size_bb = min(height - y1, size_bb)
+
+    return x1, y1, size_bb
+
+
+def extract_face(frame):
+    face_detector = dlib.get_frontal_face_detector()
+    image = np.array(frame.convert('RGB'))
+    faces = face_detector(image, 1)
+    if len(faces) > 0:
+        # For now only take the biggest face
+        face = faces[0]
+        # Face crop and rescale(follow FF++)
+        x, y, size = get_boundingbox(face, image.shape[1], image.shape[0])
+        # Get the landmarks/parts for the face in box d only with the five key points
+        cropped_face = image[y:y + size, x:x + size]
+        # cropped_face = cv2.resize(cropped_face, (224, 224), interpolation=cv2.INTER_CUBIC)
+        return Image.fromarray(cropped_face)
+    else:
+        return None
+    
+    
+def get_frame_index_uniform_sample(total_frame_num, extract_frame_num):
+    interval = np.linspace(0, total_frame_num - 1, num=extract_frame_num, dtype=int)
+    return interval.tolist()
+
+
+import cv2
+def extract_face_from_fixed_num_frames(src_video, dst_path, num_frames=None, device='cpu'):
+    """
+    1) extract specific num of frames from videos in [1st(index 0) frame, last frame] with uniform sample interval
+    2) extract face from frame with specific enlarge size
+    """
+    video_capture = cv2.VideoCapture(src_video)
+    total_frames = video_capture.get(7)
+
+    # extract from the 1st(index 0) frame
+    if num_frames is not None:
+        frame_indices = get_frame_index_uniform_sample(total_frames, num_frames)
+    else:
+        frame_indices = range(int(total_frames))
+
+    for frame_index in frame_indices:
+        video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
+        ret, frame = video_capture.read()
+        image = Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)) 
+        img = extract_face(image)
+        if img == None:
+            continue
+        img = img.resize((224, 224), Image.BICUBIC)
+        if not ret:
+            continue
+        save_img_name = f"frame_{frame_index}.png"
+        
+        img.save(os.path.join(dst_path, '0', save_img_name))
+        # cv2.imwrite(os.path.join(dst_path, '0', save_img_name), frame)
+      
+    video_capture.release()
+    # cv2.destroyAllWindows()
+
+
+def FSFM3C_video_detection(video):
+    model.to(device)
+    
+    # extract frames
+    num_frames = 32
+    
+    files = os.listdir(FRAME_SAVE_PATH)
+    num_files = len(files)
+    frame_path = os.path.join(FRAME_SAVE_PATH, str(num_files))
+    os.makedirs(frame_path, exist_ok=True)
+    os.makedirs(os.path.join(frame_path, '0'), exist_ok=True)
+    extract_face_from_fixed_num_frames(video, frame_path, num_frames=num_frames, device=device)
+    
+    args.data_path = frame_path
+    args.batch_size = 32
+    dataset_val = build_dataset(is_train=False, args=args)
+    sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val, sampler=sampler_val,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=False
+    )
+    
+    frame_preds_list, video_y_pred_list = test_all(data_loader_val, model, device)
+
+    return video_y_pred_list
+
+
+def FSFM3C_image_detection(image):
+    model.to(device)
+    
+    files = os.listdir(FRAME_SAVE_PATH)
+    num_files = len(files)
+    frame_path = os.path.join(FRAME_SAVE_PATH, str(num_files)) 
+    os.makedirs(frame_path, exist_ok=True)
+    os.makedirs(os.path.join(frame_path, '0'), exist_ok=True)
+    
+    save_img_name = f"frame_0.png"
+    img = extract_face(image)
+    if img is None:
+        return ['Invalid Input']
+    img = img.resize((224, 224), Image.BICUBIC)
+    img.save(os.path.join(frame_path, '0', save_img_name))
+    
+    args.data_path = frame_path
+    args.batch_size = 1
+    dataset_val = build_dataset(is_train=False, args=args)
+    sampler_val = torch.utils.data.SequentialSampler(dataset_val)
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val, sampler=sampler_val,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=False
+    )
+    
+    frame_preds_list, video_y_pred_list = test_all(data_loader_val, model, device)
+
+    return video_y_pred_list
+
+
+# WebUI
+with gr.Blocks() as demo:
+    gr.HTML("<h1 style='text-align: center;'>🦱 Real Facial Image&Video Detection <br> Against Face Forgery and Spoofing (Deepfake/Diffusion/Presentation-attacks)</h1>")
+    gr.Markdown("# ---Based on the fine-tuned model that is pre-trained from [FSFM-3C](https://fsfm-3c.github.io/)")
+
+    gr.Markdown("## Release  <br>"
+                "V1.0 [2024-12] (Current): <br>"
+                "[1] Create this page with basic detectors (simply fine-tuned models that follow the paper implementation): <br> "
+                "   - DfD Checkpoint_Fine-tuned on FF++: FSFM VIT-B fine-tuned on the FF++ (c23, train&val sets, 32 frames per video, 4 manipulations) dataset <br>"
+                "   - FAS Checkpoint_Fine-tuned on MCIO: FSFM VIT-B fine-tuned on the MCIO datasets (2 frames per video) <br> "
+                "   Performance is limited because no any optimization of data, models, hyperparameters, etc. is done for downstream tasks")
+
+    gr.Markdown("### TODO: We will soon update practical models, and optimized interfaces, and provide more functions such as visualizations, a unified detector, and multi-modal diagnosis.")
+
+    gr.Markdown("> Please provide an <b>image</b> or a <b>video (<100s </b>, default to uniform sampling 32 frames)</b> for detection:")
+
+
+    with gr.Column():
+        ckpt_select_dropdown = gr.Dropdown(
+            label = "Select the Model Checkpoint for Detection (🖱️ below)",
+            choices = ['choose from here'] + CKPT_LIST + ['Continuously updating...'],
+            multiselect = False,
+            value = 'choose from here',
+            interactive = True,
+            )
+        with gr.Row(elem_classes="center-align"):
+            with gr.Column(scale=5):
+                gr.Markdown(
+                    "## Image Detection"
+                )
+                image = gr.Image(label="Upload/Capture/Paste your image", type="pil")
+                image_submit_btn = gr.Button("Submit")  
+                output_results_image = gr.Textbox(label="Detection Result")
+            with gr.Column(scale=5):
+                gr.Markdown(
+                    "## Video Detection"
+                )
+                video = gr.Video(label="Upload/Capture your video")
+                video_submit_btn = gr.Button("Submit")
+                output_results_video = gr.Textbox(label="Detection Result")
+
+    image_submit_btn.click(
+        fn=FSFM3C_image_detection,
+        inputs=[image],
+        outputs=[output_results_image],
+    )
+    video_submit_btn.click(
+        fn=FSFM3C_video_detection,
+        inputs=[video],
+        outputs=[output_results_video],
+    )
+    ckpt_select_dropdown.change(
+        fn=load_model,
+        inputs=[ckpt_select_dropdown],
+        outputs=[ckpt_select_dropdown],
+    )
+
+    
+if __name__ == "__main__":
+    gr.close_all()
+    demo.queue()
+    demo.launch()

engine_finetune.py

@@ -0,0 +1,323 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import math
+import sys
+from typing import Iterable, Optional
+
+import numpy as np
+import torch
+
+from timm.data import Mixup
+from timm.utils import accuracy
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+from util.metrics import *
+
+import torch.nn.functional as F
+
+
+def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+                    mixup_fn: Optional[Mixup] = None, log_writer=None,
+                    args=None):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        with torch.cuda.amp.autocast():
+            # outputs = model(samples)
+            outputs = model(samples).to(device, non_blocking=True)  # modified
+            loss = criterion(outputs, targets)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, clip_grad=max_norm,
+                    parameters=model.parameters(), create_graph=False,
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        metric_logger.update(loss=loss_value)
+        min_lr = 10.
+        max_lr = 0.
+        for group in optimizer.param_groups:
+            min_lr = min(min_lr, group["lr"])
+            max_lr = max(max_lr, group["lr"])
+
+        metric_logger.update(lr=max_lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
+            """ We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar('loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', max_lr, epoch_1000x)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def evaluate(data_loader, model, device):
+    criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+
+    for batch in metric_logger.log_every(data_loader, 10, header):
+        images = batch[0]
+        target = batch[-1]
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        # compute output
+        with torch.cuda.amp.autocast():
+            # output = model(images)
+            output = model(images).to(device, non_blocking=True)  # modified
+            loss = criterion(output, target)
+
+        # acc1, acc5 = accuracy(output, target, topk=(1, 5))
+        acc = float(accuracy(output, target, topk=(1,))[0])
+        preds = (F.softmax(output, dim=1)[:, 1].detach().cpu().numpy())
+        trues = (target.detach().cpu().numpy())
+        auc_score = roc_auc_score(trues, preds) * 100.
+
+        batch_size = images.shape[0]
+        metric_logger.update(loss=loss.item())
+        # metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+        # metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+        metric_logger.meters['acc'].update(acc, n=batch_size)
+        metric_logger.meters['auc'].update(auc_score, n=batch_size)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    # print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
+    #       .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
+    print('* Acc {acc.global_avg:.3f} Auc {auc.global_avg:.3f}  loss {losses.global_avg:.3f}'
+          .format(acc=metric_logger.acc, auc=metric_logger.auc, losses=metric_logger.loss))
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def test_ori(data_loader, model, device):
+    criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+
+    labels = np.array([])
+    preds = np.array([])
+
+    for batch in metric_logger.log_every(data_loader, 10, header):
+        images = batch[0]
+        target = batch[-1]
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        # compute output
+        with torch.cuda.amp.autocast():
+            # output = model(images)
+            output = model(images).to(device, non_blocking=True)  # modified
+            loss = criterion(output, target)
+
+        # acc1, acc5 = accuracy(output, target, topk=(1, 5))
+        acc = float(accuracy(output, target, topk=(1,))[0])
+        pred = (F.softmax(output, dim=1)[:, 1].detach().cpu().numpy())
+        preds = np.append(preds, pred)
+        label = (target.detach().cpu().numpy())
+        labels = np.append(labels, label)
+
+        batch_size = images.shape[0]
+        metric_logger.update(loss=loss.item())
+        # metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+        # metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+        metric_logger.meters['acc'].update(acc, n=batch_size)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    # print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
+    #       .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
+    auc_score = roc_auc_score(labels, preds) * 100.
+    metric_logger.meters['auc'].update(auc_score)
+    print('* Acc {acc.global_avg:.3f} Auc {auc.global_avg:.3f}  loss {losses.global_avg:.3f}'
+          .format(acc=metric_logger.acc, auc=metric_logger.auc, losses=metric_logger.loss))
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def test(data_loader, model, device):
+    criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+
+    frame_labels = np.array([])  # int label
+    frame_preds = np.array([])  # pred logit
+    frame_y_preds = np.array([])  # pred int
+
+    # for batch in metric_logger.log_every(data_loader, print_freq=len(data_loader), header=header):
+    for batch in data_loader:
+        images = batch[0]  # torch.Size([BS, C, H, W])
+        target = batch[1]  # torch.Size([BS])
+
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        # compute output
+        with torch.cuda.amp.autocast():
+            # output = model(images)
+            output = model(images).to(device, non_blocking=True)  # modified
+            loss = criterion(output, target)
+
+        frame_pred = (F.softmax(output, dim=1)[:, 1].detach().cpu().numpy())
+        frame_preds = np.append(frame_preds, frame_pred)
+
+        frame_y_pred = np.argmax(output.detach().cpu().numpy(), axis=1)
+        frame_y_preds = np.append(frame_y_preds, frame_y_pred)
+
+        frame_label = (target.detach().cpu().numpy())
+        frame_labels = np.append(frame_labels, frame_label)
+
+        metric_logger.update(loss=loss.item())
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    metric_logger.meters['frame_acc'].update(frame_level_acc(frame_labels, frame_y_preds))
+    metric_logger.meters['frame_balanced_acc'].update(frame_level_balanced_acc(frame_labels, frame_y_preds))
+    metric_logger.meters['frame_auc'].update(frame_level_auc(frame_labels, frame_preds))
+    metric_logger.meters['frame_eer'].update(frame_level_eer(frame_labels, frame_preds))
+
+    print('*[------FRAME-LEVEL------] \n'
+          'Acc {frame_acc.global_avg:.3f} Balanced_Acc {frame_balanced_acc.global_avg:.3f} '
+          'Auc {frame_auc.global_avg:.3f} EER {frame_eer.global_avg:.3f} loss {losses.global_avg:.3f}'
+          .format(frame_acc=metric_logger.frame_acc, frame_balanced_acc=metric_logger.frame_balanced_acc,
+                  frame_auc=metric_logger.frame_auc, frame_eer=metric_logger.frame_eer, losses=metric_logger.loss))
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def test_all(data_loader, model, device):
+    criterion = torch.nn.CrossEntropyLoss()
+
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+
+    frame_labels = np.array([])  # int label
+    frame_preds = np.array([])  # pred logit
+    frame_y_preds = np.array([])  # pred int
+    video_names_list = list()
+
+    # for batch in metric_logger.log_every(data_loader, print_freq=len(data_loader), header=header):
+    for batch in data_loader:
+        images = batch[0]  # torch.Size([BS, C, H, W])
+        target = batch[1]  # torch.Size([BS])
+        video_name = batch[-1]  # list[BS]
+
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        # compute output
+        # with torch.cuda.amp.autocast():
+            # output = model(images)
+        output = model(images).to(device, non_blocking=True)  # modified
+        loss = criterion(output, target)
+
+        frame_pred = (F.softmax(output, dim=1)[:, 1].detach().cpu().numpy())
+        frame_preds = np.append(frame_preds, frame_pred)
+
+        frame_y_pred = np.argmax(output.detach().cpu().numpy(), axis=1)
+        frame_y_preds = np.append(frame_y_preds, frame_y_pred)
+
+        frame_label = (target.detach().cpu().numpy())
+        frame_labels = np.append(frame_labels, frame_label)
+
+        video_names_list.extend(list(video_name))
+
+        metric_logger.update(loss=loss.item())
+
+    # gather the stats from all processes
+    # metric_logger.synchronize_between_processes()
+    # metric_logger.meters['frame_acc'].update(frame_level_acc(frame_labels, frame_y_preds))
+    # metric_logger.meters['frame_balanced_acc'].update(frame_level_balanced_acc(frame_labels, frame_y_preds))
+    # metric_logger.meters['frame_auc'].update(frame_level_auc(frame_labels, frame_preds))
+    # metric_logger.meters['frame_eer'].update(frame_level_eer(frame_labels, frame_preds))
+
+    # print('*[------FRAME-LEVEL------] \n'
+    #       'Acc {frame_acc.global_avg:.3f} Balanced_Acc {frame_balanced_acc.global_avg:.3f} '
+    #       'Auc {frame_auc.global_avg:.3f} EER {frame_eer.global_avg:.3f} loss {losses.global_avg:.3f}'
+    #       .format(frame_acc=metric_logger.frame_acc, frame_balanced_acc=metric_logger.frame_balanced_acc,
+    #               frame_auc=metric_logger.frame_auc, frame_eer=metric_logger.frame_eer, losses=metric_logger.loss))
+
+    # video-level metrics:
+    frame_labels_list = frame_labels.tolist()
+    frame_preds_list = frame_preds.tolist()
+
+    video_label_list, video_pred_list, video_y_pred_list = get_video_level_label_pred(frame_labels_list, video_names_list, frame_preds_list)
+    # print(len(video_label_list), len(video_pred_list), len(video_y_pred_list))
+    # metric_logger.meters['video_acc'].update(video_level_acc(video_label_list, video_y_pred_list))
+    # metric_logger.meters['video_balanced_acc'].update(video_level_balanced_acc(video_label_list, video_y_pred_list))
+    # metric_logger.meters['video_auc'].update(video_level_auc(video_label_list, video_pred_list))
+    # metric_logger.meters['video_eer'].update(frame_level_eer(video_label_list, video_pred_list))
+
+    # print('*[------VIDEO-LEVEL------] \n'
+    #       'Acc {video_acc.global_avg:.3f} Balanced_Acc {video_balanced_acc.global_avg:.3f} '
+    #       'Auc {video_auc.global_avg:.3f} EER {video_eer.global_avg:.3f}'
+    #       .format(video_acc=metric_logger.video_acc, video_balanced_acc=metric_logger.video_balanced_acc,
+    #               video_auc=metric_logger.video_auc, video_eer=metric_logger.video_eer))
+
+    # return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+    return frame_preds_list, video_y_pred_list

models_vit.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+import timm.models.vision_transformer
+
+
+class VisionTransformer(timm.models.vision_transformer.VisionTransformer):
+    """ Vision Transformer with support for global average pooling
+    """
+    def __init__(self, global_pool=False, **kwargs):
+        super(VisionTransformer, self).__init__(**kwargs)
+
+        self.global_pool = global_pool
+        if self.global_pool:
+            norm_layer = kwargs['norm_layer']
+            embed_dim = kwargs['embed_dim']
+            self.fc_norm = norm_layer(embed_dim)
+
+            del self.norm  # remove the original norm
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        if self.global_pool:
+            x = x[:, 1:, :].mean(dim=1)  # global pool without cls token
+            outcome = self.fc_norm(x)
+        else:
+            x = self.norm(x)
+            outcome = x[:, 0]
+
+        return outcome
+
+
+def vit_small_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=384, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,  # ViT-small config in MOCO_V3
+        # patch_size=16, embed_dim=768, depth=8, num_heads=8, mlp_ratio=3, qkv_bias=True,  # ViT-small config in timm
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_base_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_large_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_huge_patch14(**kwargs):
+    model = VisionTransformer(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+# class VisionTransformerWithProjector(VisionTransformer):
+#     def __init__(self, vit_model, model_encoder, feat_cl_dim=128):
+#         super(VisionTransformerWithProjector, self).__init__()
+#         self.encoder = vit_model
+#         embed_dim = {'vit_base_patch16': 768, 'vit_large_patch16': 1024, 'vit_huge_patch14': 1280}
+#         self.projection_head = nn.Sequential(
+#             nn.Linear(embed_dim[model_encoder], embed_dim[model_encoder]),
+#             nn.ReLU(inplace=True),
+#             nn.Linear(embed_dim[model_encoder], feat_cl_dim)
+#         )
+#
+#     def forward(self, x):
+#         x = self.encoder(x)
+#         latent_cl = self.projection_head(x)  # [N, feat_cl_dim]
+#         features = nn.functional.normalize(latent_cl, dim=-1)  # [N, feat_cl_dim]
+#         return features

requirements.txt

@@ -0,0 +1,77 @@
+aiofiles==23.2.1
+annotated-types==0.7.0
+anyio==4.6.2.post1
+certifi==2024.8.30
+charset-normalizer==3.4.0
+click==8.1.7
+contourpy==1.3.0
+cycler==0.12.1
+# dlib==19.24.6
+exceptiongroup==1.2.2
+fastapi==0.115.5
+ffmpy==0.4.0
+filelock==3.16.1
+fonttools==4.55.0
+fsspec==2024.10.0
+gradio==4.44.1
+gradio_client==1.3.0
+h11==0.14.0
+httpcore==1.0.7
+httpx==0.27.2
+huggingface-hub==0.26.2
+idna==3.10
+imageio==2.36.0
+importlib_resources==6.4.5
+Jinja2==3.1.4
+joblib==1.4.2
+kiwisolver==1.4.7
+lazy_loader==0.4
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.2
+mdurl==0.1.2
+networkx==3.2.1
+numpy==1.24.4
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+opencv-python==4.10.0.84
+orjson==3.10.11
+packaging==24.2
+pandas==2.2.3
+pillow==10.4.0
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.2.0
+python-dateutil==2.9.0.post0
+python-multipart==0.0.17
+pytz==2024.2
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.9.4
+ruff==0.7.4
+scikit-image==0.24.0
+scikit-learn==1.5.2
+scipy==1.13.1
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+starlette==0.41.2
+threadpoolctl==3.5.0
+tifffile==2024.8.30
+timm==0.4.5
+tomlkit==0.12.0
+torch==1.13.1
+torchvision==0.14.1
+tqdm==4.67.0
+typer==0.13.0
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+uvicorn==0.32.0
+validators==0.34.0
+websockets==12.0
+zipp==3.21.0

util/crop.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import math
+
+import torch
+
+from torchvision import transforms
+from torchvision.transforms import functional as F
+
+
+class RandomResizedCrop(transforms.RandomResizedCrop):
+    """
+    RandomResizedCrop for matching TF/TPU implementation: no for-loop is used.
+    This may lead to results different with torchvision's version.
+    Following BYOL's TF code:
+    https://github.com/deepmind/deepmind-research/blob/master/byol/utils/dataset.py#L206
+    """
+    @staticmethod
+    def get_params(img, scale, ratio):
+        width, height = F._get_image_size(img)
+        area = height * width
+
+        target_area = area * torch.empty(1).uniform_(scale[0], scale[1]).item()
+        log_ratio = torch.log(torch.tensor(ratio))
+        aspect_ratio = torch.exp(
+            torch.empty(1).uniform_(log_ratio[0], log_ratio[1])
+        ).item()
+
+        w = int(round(math.sqrt(target_area * aspect_ratio)))
+        h = int(round(math.sqrt(target_area / aspect_ratio)))
+
+        w = min(w, width)
+        h = min(h, height)
+
+        i = torch.randint(0, height - h + 1, size=(1,)).item()
+        j = torch.randint(0, width - w + 1, size=(1,)).item()
+
+        return i, j, h, w

util/datasets.py

@@ -0,0 +1,350 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import os
+import json
+import shutil
+
+from torchvision import datasets, transforms
+
+from timm.data import create_transform
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+import numpy as np
+from PIL import Image
+import random
+import torch
+from torch.utils.data import DataLoader, Dataset, ConcatDataset
+from torchvision import transforms
+from torch.nn import functional as F
+
+
+class collate_fn_crfrp:
+    def __init__(self, input_size=224, patch_size=16, mask_ratio=0.75):
+        self.img_size = input_size
+        self.patch_size = patch_size
+        self.num_patches_axis = input_size // patch_size
+        self.num_patches = (input_size // patch_size) ** 2
+        self.mask_ratio = mask_ratio
+
+        # --------------------------------------------------------------------------
+        # self.facial_region_group = [
+        #     [2],  # right eyebrow
+        #     [3],  # left eyebrow
+        #     [4],  # right eye
+        #     [5],  # left eye
+        #     [6],  # nose
+        #     [7, 8],  # upper mouth
+        #     [8, 9],  # lower mouth
+        #     [10, 1, 0],  # facial boundaries
+        #     [10],  # hair
+        #     [1],  # facial skin
+        #     [0]  # background
+        # ]
+        self.facial_region_group = [
+            [2, 3],  # eyebrows
+            [4, 5],  # eyes
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10, 1, 0],  # face boundaries
+            [10],  # hair
+            [1],  # facial skin
+            [0]  # background
+        ]  # ['background', 'face', 'rb', 'lb', 're', 'le', 'nose', 'ulip', 'imouth', 'llip', 'hair']
+
+    def __call__(self, samples):
+        image, img_mask, facial_region_mask, random_specific_facial_region \
+            = self.CRFR_P_masking(samples, specified_facial_region=None)
+
+        return {'image': image, 'img_mask': img_mask, 'specific_facial_region_mask': facial_region_mask}
+
+        # # using following code if using different data augmentation for target view
+        # image, img_mask, facial_region_mask, random_specific_facial_region \
+        #     = self.CRFR_P_masking(samples, specified_facial_region=None)
+        # image_cl, img_mask_cl, facial_region_mask_cl, random_specific_facial_region_cl \
+        #     = self.CRFR_P_masking(samples, specified_facial_region=random_specific_facial_region)
+        #
+        # return {'image': image, 'img_mask': img_mask, 'specific_facial_region_mask': facial_region_mask,
+        #         'image_cl': image_cl, 'img_mask_cl': img_mask_cl, 'specific_facial_region_mask_cl': facial_region_mask_cl}
+
+    def CRFR_P_masking(self, samples, specified_facial_region=None):
+        image = torch.stack([sample['image'] for sample in samples])  # torch.Size([bs, 3, 224, 224])
+        parsing_map = torch.stack([sample['parsing_map'] for sample in samples])  # torch.Size([bs, 1, 224, 224])
+        parsing_map = parsing_map.squeeze(1)  # torch.Size([BS, 1, 224, 224]) → torch.Size([BS, 224, 224])
+
+        # covering a randomly select facial_region_group and get fr_mask(masking all patches include this region)
+        facial_region_mask = torch.zeros(parsing_map.size(0), self.num_patches_axis, self.num_patches_axis,
+                                         dtype=torch.float32)  # torch.Size([BS, H/P, W/P])
+        facial_region_mask, random_specific_facial_region \
+            = self.masking_all_patches_in_random_specific_facial_region(parsing_map, facial_region_mask)
+        # torch.Size([num_patches,]), list
+
+        img_mask, facial_region_mask \
+            = self.variable_proportional_masking(parsing_map, facial_region_mask, random_specific_facial_region)
+        # torch.Size([num_patches,]), torch.Size([num_patches,])
+
+        del parsing_map
+        return image, img_mask, facial_region_mask, random_specific_facial_region
+
+    def masking_all_patches_in_random_specific_facial_region(self, parsing_map, facial_region_mask,
+                                                             # specified_facial_region=None
+                                                             ):
+        # while True:
+        #     random_specific_facial_region = random.choice(self.facial_region_group[:-2])
+        #     if random_specific_facial_region != specified_facial_region:
+        #         break
+        random_specific_facial_region = random.choice(self.facial_region_group[:-2])
+        if random_specific_facial_region == [10, 1, 0]:  # facial boundaries, 10-hair 1-skin 0-background
+            # True for hair(10) or bg(0) patches:
+            patch_hair_bg = F.max_pool2d(((parsing_map == 10) + (parsing_map == 0)).float(),
+                                         kernel_size=self.patch_size)
+            # True for skin(1) patches:
+            patch_skin = F.max_pool2d((parsing_map == 1).float(), kernel_size=self.patch_size)
+            # skin&hair or skin&bg is defined as facial boundaries：
+            facial_region_mask = (patch_hair_bg.bool() & patch_skin.bool()).float()
+        else:
+            for facial_region_index in random_specific_facial_region:
+                facial_region_mask = torch.maximum(facial_region_mask,
+                                                   F.max_pool2d((parsing_map == facial_region_index).float(),
+                                                                kernel_size=self.patch_size))
+
+        return facial_region_mask.view(parsing_map.size(0), -1), random_specific_facial_region
+
+    def variable_proportional_masking(self, parsing_map, facial_region_mask, random_specific_facial_region):
+        img_mask = facial_region_mask.clone()
+
+        # proportional masking patches in other regions
+        other_facial_region_group = [region for region in self.facial_region_group if
+                                     region != random_specific_facial_region]
+        # print(other_facial_region_group)
+        for i in range(facial_region_mask.size(0)):  # iterate each map in BS
+            num_mask_to_change = (self.mask_ratio * self.num_patches - facial_region_mask[i].sum(dim=-1)).int()
+            # mask_change_to = 1 if num_mask_to_change >= 0 else 0
+            mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+
+            if mask_change_to == 1:
+                # proportional masking patches in other facial regions according to the corresponding ratio
+                mask_ratio_other_fr = (
+                        num_mask_to_change / (self.num_patches - facial_region_mask[i].sum(dim=-1)))
+
+                masked_patches = facial_region_mask[i].clone()
+                for other_fr in other_facial_region_group:
+                    to_mask_patches = torch.zeros(1, self.num_patches_axis, self.num_patches_axis,
+                                                  dtype=torch.float32)
+                    if other_fr == [10, 1, 0]:
+                        patch_hair_bg = F.max_pool2d(
+                            ((parsing_map[i].unsqueeze(0) == 10) + (parsing_map[i].unsqueeze(0) == 0)).float(),
+                            kernel_size=self.patch_size)
+                        patch_skin = F.max_pool2d((parsing_map[i].unsqueeze(0) == 1).float(),
+                                                  kernel_size=self.patch_size)
+                        # skin&hair or skin&bg defined as facial boundaries：
+                        to_mask_patches = (patch_hair_bg.bool() & patch_skin.bool()).float()
+                    else:
+                        for facial_region_index in other_fr:
+                            to_mask_patches = torch.maximum(to_mask_patches,
+                                                            F.max_pool2d((parsing_map[i].unsqueeze(
+                                                                0) == facial_region_index).float(),
+                                                                         kernel_size=self.patch_size))
+
+                    # ignore already masked patches:
+                    to_mask_patches = (to_mask_patches.view(-1) - masked_patches) > 0
+                    select_indices = to_mask_patches.nonzero(as_tuple=False).view(-1)
+                    change_indices = torch.randperm(len(select_indices))[
+                                     :torch.round(to_mask_patches.sum() * mask_ratio_other_fr).int()]
+                    img_mask[i, select_indices[change_indices]] = mask_change_to
+                    # prevent overlap
+                    masked_patches = masked_patches + to_mask_patches.float()
+
+                # mask/unmask patch from other facial regions to get img_mask with fixed size
+                num_mask_to_change = (self.mask_ratio * self.num_patches - img_mask[i].sum(dim=-1)).int()
+                # mask_change_to = 1 if num_mask_to_change >= 0 else 0
+                mask_change_to = torch.clamp(num_mask_to_change, 0, 1).item()
+                # prevent unmasking facial_region_mask
+                select_indices = ((img_mask[i] + facial_region_mask[i]) == (1 - mask_change_to)).nonzero(
+                    as_tuple=False).view(-1)
+                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+                img_mask[i, select_indices[change_indices]] = mask_change_to
+
+            else:
+                # Extreme situations:
+                # if fr_mask is already over(>=) num_patches*mask_ratio, unmask it to get img_mask with fixed ratio
+                select_indices = (facial_region_mask[i] == (1 - mask_change_to)).nonzero(as_tuple=False).view(-1)
+                change_indices = torch.randperm(len(select_indices))[:torch.abs(num_mask_to_change)]
+                img_mask[i, select_indices[change_indices]] = mask_change_to
+                facial_region_mask[i] = img_mask[i]
+
+        return img_mask, facial_region_mask
+
+
+class FaceParsingDataset(Dataset):
+    def __init__(self, root, transform=None):
+        self.root_dir = root
+        self.transform = transform
+        self.image_folder = os.path.join(root, 'images')
+        self.parsing_map_folder = os.path.join(root, 'parsing_maps')
+        self.image_names = os.listdir(self.image_folder)
+
+    def __len__(self):
+        return len(self.image_names)
+
+    def __getitem__(self, idx):
+        img_name = os.path.join(self.image_folder, self.image_names[idx])
+        parsing_map_name = os.path.join(self.parsing_map_folder, self.image_names[idx].replace('.png', '.npy'))
+
+        image = Image.open(img_name).convert("RGB")
+        parsing_map_np = np.load(parsing_map_name)
+
+        if self.transform:
+            image = self.transform(image)
+
+        # Convert mask to tensor
+        parsing_map = torch.from_numpy(parsing_map_np)
+        del parsing_map_np  # may save mem
+
+        return {'image': image, 'parsing_map': parsing_map}
+
+
+class TestImageFolder(datasets.ImageFolder):
+    def __init__(self, root, transform=None, target_transform=None):
+        super(TestImageFolder, self).__init__(root, transform, target_transform)
+
+    def __getitem__(self, index):
+        # Call the parent class method to load image and label
+        original_tuple = super(TestImageFolder, self).__getitem__(index)
+
+        # Get the video name
+        video_name = self.imgs[index][0].split('/')[-1].split('_frame_')[0]  # the separator of video name
+
+        # Extend the tuple to include video name
+        extended_tuple = (original_tuple + (video_name,))
+
+        return extended_tuple
+
+
+def get_mean_std(args):
+    print('dataset_paths:', args.data_path)
+    transform = transforms.Compose([transforms.ToTensor(),
+                                    transforms.Resize((args.input_size, args.input_size),
+                                                      interpolation=transforms.InterpolationMode.BICUBIC)])
+
+    if len(args.data_path) > 1:
+        pretrain_datasets = [FaceParsingDataset(root=path, transform=transform) for path in args.data_path]
+        dataset_pretrain = ConcatDataset(pretrain_datasets)
+    else:
+        pretrain_datasets = args.data_path[0]
+        dataset_pretrain = FaceParsingDataset(root=pretrain_datasets, transform=transform)
+
+    print('Compute mean and variance for pretraining data.')
+    print('len(dataset_train): ', len(dataset_pretrain))
+
+    loader = DataLoader(
+        dataset_pretrain,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+    )
+
+    channels_sum, channels_squared_sum, num_batches = 0, 0, 0
+    for sample in loader:
+        data = sample['image']
+        channels_sum += torch.mean(data, dim=[0, 2, 3])
+        channels_squared_sum += torch.mean(data ** 2, dim=[0, 2, 3])
+        num_batches += 1
+
+    mean = channels_sum / num_batches
+    std = (channels_squared_sum / num_batches - mean ** 2) ** 0.5
+
+    print(f'train dataset mean%: {mean.numpy()} std: %{std.numpy()} ')
+    del pretrain_datasets, dataset_pretrain, loader
+    return mean.numpy(), std.numpy()
+
+
+def build_dataset(is_train, args):
+    transform = build_transform(is_train, args)
+    dataset = datasets.ImageFolder(args.data_path, transform=transform)
+    # if args.eval:
+    #     # no loading training set
+    #     root = os.path.join(args.data_path, 'test' if is_train else 'test')
+    #     dataset = TestImageFolder(root, transform=transform)
+    # else:
+    #     root = os.path.join(args.data_path, 'train' if is_train else 'val')
+    #     dataset = datasets.ImageFolder(root, transform=transform)
+    #     print(dataset)
+
+    return dataset
+
+
+def build_transform(is_train, args):
+    if args.normalize_from_IMN:
+        mean = IMAGENET_DEFAULT_MEAN
+        std = IMAGENET_DEFAULT_STD
+        # print(f'mean:{mean}, std:{std}')
+    else:
+        if not os.path.exists(os.path.join(args.output_dir, "/pretrain_ds_mean_std.txt")) and not args.eval:
+            shutil.copyfile(os.path.dirname(args.finetune) + '/pretrain_ds_mean_std.txt',
+                            os.path.join(args.output_dir) + '/pretrain_ds_mean_std.txt')
+        with open(os.path.join(os.path.dirname(args.resume)) + '/pretrain_ds_mean_std.txt' if args.eval
+                  else os.path.join(args.output_dir) + '/pretrain_ds_mean_std.txt', 'r') as file:
+            ds_stat = json.loads(file.readline())
+            mean = ds_stat['mean']
+            std = ds_stat['std']
+            # print(f'mean:{mean}, std:{std}')
+
+    if args.apply_simple_augment:
+        if is_train:
+            # this should always dispatch to transforms_imagenet_train
+            transform = create_transform(
+                input_size=args.input_size,
+                is_training=True,
+                color_jitter=args.color_jitter,
+                auto_augment=args.aa,
+                interpolation=transforms.InterpolationMode.BICUBIC,
+                re_prob=args.reprob,
+                re_mode=args.remode,
+                re_count=args.recount,
+                mean=mean,
+                std=std,
+            )
+            return transform
+
+        # no augment / eval transform
+        t = []
+        if args.input_size <= 224:
+            crop_pct = 224 / 256
+        else:
+            crop_pct = 1.0
+        size = int(args.input_size / crop_pct)  # 256
+        t.append(
+            transforms.Resize(size, interpolation=transforms.InterpolationMode.BICUBIC),
+            # to maintain same ratio w.r.t. 224 images
+        )
+        t.append(transforms.CenterCrop(args.input_size))  # 224
+
+        t.append(transforms.ToTensor())
+        t.append(transforms.Normalize(mean, std))
+        return transforms.Compose(t)
+
+    else:
+        t = []
+        if args.input_size < 224:
+            crop_pct = input_size / 224
+        else:
+            crop_pct = 1.0
+        size = int(args.input_size / crop_pct)  # size = 224
+        t.append(
+            transforms.Resize(size, interpolation=transforms.InterpolationMode.BICUBIC),
+            # to maintain same ratio w.r.t. 224 images
+        )
+        # t.append(
+        #     transforms.Resize((224, 224), interpolation=transforms.InterpolationMode.BICUBIC),
+        #     # to maintain same ratio w.r.t. 224 images
+        # )
+
+        t.append(transforms.ToTensor())
+        t.append(transforms.Normalize(mean, std))
+        return transforms.Compose(t)

util/lars.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import torch
+
+
+class LARS(torch.optim.Optimizer):
+    """
+    LARS optimizer, no rate scaling or weight decay for parameters <= 1D.
+    """
+    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, trust_coefficient=0.001):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, trust_coefficient=trust_coefficient)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for g in self.param_groups:
+            for p in g['params']:
+                dp = p.grad
+
+                if dp is None:
+                    continue
+
+                if p.ndim > 1: # if not normalization gamma/beta or bias
+                    dp = dp.add(p, alpha=g['weight_decay'])
+                    param_norm = torch.norm(p)
+                    update_norm = torch.norm(dp)
+                    one = torch.ones_like(param_norm)
+                    q = torch.where(param_norm > 0.,
+                                    torch.where(update_norm > 0,
+                                    (g['trust_coefficient'] * param_norm / update_norm), one),
+                                    one)
+                    dp = dp.mul(q)
+
+                param_state = self.state[p]
+                if 'mu' not in param_state:
+                    param_state['mu'] = torch.zeros_like(p)
+                mu = param_state['mu']
+                mu.mul_(g['momentum']).add_(dp)
+                p.add_(mu, alpha=-g['lr'])

util/loss_contrastive.py

@@ -0,0 +1,360 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+from __future__ import print_function
+
+import torch
+import torch.nn as nn
+import math
+
+
+class SimSiamLoss(nn.Module):
+    def __init__(self):
+        super(SimSiamLoss, self).__init__()
+        self.criterion = nn.CosineSimilarity(dim=1)
+
+    def forward(self, cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+        loss = -(self.criterion(cl_features_1, cl_features_2).mean()) * 0.5
+
+        # if not math.isfinite(loss):
+        #     print(cl_features_1, '\n', cl_features_2)
+        #     print(self.criterion(cl_features_1, cl_features_2))
+
+        return loss
+
+
+class BYOLLoss(nn.Module):
+    def __init__(self):
+        super(BYOLLoss, self).__init__()
+
+    @staticmethod
+    def forward(cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+        loss = 2 - 2 * (cl_features_1 * cl_features_2).sum(dim=-1)
+        # loss = 1 - (cl_features_1 * cl_features_2).sum(dim=-1)
+        loss = loss.mean()
+
+        if not math.isfinite(loss):
+            print(cl_features_1, '\n', cl_features_2)
+            print(2 - 2 * (cl_features_1 * cl_features_2).sum(dim=-1))
+
+        return loss
+
+
+# different implementation of InfoNCELoss, including MOCOV3Loss; SupConLoss
+class InfoNCELoss(nn.Module):
+    def __init__(self, temperature=0.1, contrast_sample='all'):
+        """
+        from CMAE: https://github.com/ZhichengHuang/CMAE/issues/5
+        :param temperature: 0.1 0.5 1.0, 1.5 2.0
+        """
+        super(InfoNCELoss, self).__init__()
+        self.temperature = temperature
+        self.criterion = nn.CrossEntropyLoss()
+        self.contrast_sample = contrast_sample
+
+    def forward(self, cl_features):
+        """
+        Args:
+            :param cl_features: : hidden vector of shape [bsz, n_views, ...]
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if cl_features.is_cuda
+                  else torch.device('cpu'))
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+        score_all = torch.matmul(cl_features_1, cl_features_2.transpose(1, 0))  # [BS, BS]
+        score_all = score_all / self.temperature
+        bs = score_all.size(0)
+
+        if self.contrast_sample == 'all':
+            score = score_all
+        elif self.contrast_sample == 'positive':
+            mask = torch.eye(bs, dtype=torch.float).to(device)  # torch.Size([BS, BS])
+            score = score_all * mask
+        else:
+            raise ValueError('Contrastive sample: all{pos&neg} or positive(positive)')
+
+        # label = (torch.arange(bs, dtype=torch.long) +
+        #          bs * torch.distributed.get_rank()).to(device)
+        label = torch.arange(bs, dtype=torch.long).to(device)
+
+        loss = 2 * self.temperature * self.criterion(score, label)
+
+        if not math.isfinite(loss):
+            print(cl_features_1, '\n', cl_features_2)
+            print(score_all, '\n', score, '\n', mask)
+
+        return loss
+
+
+class MOCOV3Loss(nn.Module):
+    def __init__(self, temperature=0.1):
+        super(MOCOV3Loss, self).__init__()
+        self.temperature = temperature
+
+    def forward(self, cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1]  # [BS, feat_cl_dim]
+
+        # normalize
+        cl_features_1 = nn.functional.normalize(cl_features_1, dim=1)
+        cl_features_2 = nn.functional.normalize(cl_features_2, dim=1)
+        # Einstein sum is more intuitive
+        logits = torch.einsum('nc,mc->nm', [cl_features_1, cl_features_2]) / self.temperature
+        N = logits.shape[0]
+        labels = (torch.arange(N, dtype=torch.long)).cuda()
+        return nn.CrossEntropyLoss()(logits, labels) * (2 * self.temperature)
+
+
+class SupConLoss(nn.Module):
+    """
+    from: https://github.com/HobbitLong/SupContrast
+    Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
+    It also supports the unsupervised contrastive loss in SimCLR"""
+    def __init__(self, temperature=0.1, contrast_mode='all', contrast_sample='all',
+                 base_temperature=0.1):
+        super(SupConLoss, self).__init__()
+        self.temperature = temperature
+        self.contrast_mode = contrast_mode
+        self.contrast_sample = contrast_sample
+        self.base_temperature = base_temperature
+
+    def forward(self, features, labels=None, mask=None):
+        """Compute loss for model. If both `labels` and `mask` are None,
+        it degenerates to SimCLR unsupervised loss:
+        https://arxiv.org/pdf/2002.05709.pdf
+
+        Args:
+            features: hidden vector of shape [bsz, n_views, ...].
+            labels: ground truth of shape [bsz].
+            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
+                has the same class as sample i. Can be asymmetric.
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if features.is_cuda
+                  else torch.device('cpu'))
+
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, ...],'
+                             'at least 3 dimensions are required')
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        batch_size = features.shape[0]
+        if labels is not None and mask is not None:
+            raise ValueError('Cannot define both `labels` and `mask`')
+        elif labels is None and mask is None:
+            mask = torch.eye(batch_size, dtype=torch.float32).to(device)  # torch.Size([BS, BS])
+        elif labels is not None:
+            labels = labels.contiguous().view(-1, 1)
+            if labels.shape[0] != batch_size:
+                raise ValueError('Num of labels does not match num of features')
+            mask = torch.eq(labels, labels.T).float().to(device)
+        else:
+            mask = mask.float().to(device)
+
+        contrast_count = features.shape[1]  # contrast_count(2)
+        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)  # [BS*contrast_count, D]
+        if self.contrast_mode == 'one':
+            anchor_feature = features[:, 0]  # [BS, D]
+            anchor_count = 1
+        elif self.contrast_mode == 'all':
+            anchor_feature = contrast_feature  # [BS*contrast_count, D]
+            anchor_count = contrast_count
+        else:
+            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))
+
+        # compute logits
+        anchor_dot_contrast = torch.div(
+            torch.matmul(anchor_feature, contrast_feature.T),
+            self.temperature)  # [BS*contrast_count, BS*contrast_count]
+        # for numerical stability
+        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)  # [BS*contrast_count, 1]
+        logits = anchor_dot_contrast - logits_max.detach()  # [BS*contrast_count, BS*contrast_count]
+
+        # tile mask
+        mask = mask.repeat(anchor_count, contrast_count)  # [BS*anchor_count, BS*contrast_count]
+        # mask-out self-contrast cases
+        logits_mask = torch.scatter(
+            torch.ones_like(mask),
+            1,
+            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
+            0
+        )  # [BS*anchor_count, BS*contrast_count]
+        mask = mask * logits_mask  # [BS*anchor_count, BS*contrast_count]
+
+        """
+        logits_mask is used to get the denominator(positives and negatives).
+        mask is used to get the numerator(positives). mask is applied to log_prob.
+        """
+
+        # compute log_prob，logits_mask is contrast anchor with both positives and negatives
+        exp_logits = torch.exp(logits) * logits_mask  # [BS*anchor_count, BS*contrast_count]
+        # compute log_prob，logits_mask is contrast anchor with negatives, i.e., denominator only negatives contrast:
+        # exp_logits = torch.exp(logits) * (logits_mask-mask)
+
+        if self.contrast_sample == 'all':
+            log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))  # [BS*anchor_count, BS*anchor_count]
+            # compute mean of log-likelihood over positive
+            mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)  # [BS*anchor_count]
+        elif self.contrast_sample == 'positive':
+            mean_log_prob_pos = (mask * logits).sum(1) / mask.sum(1)
+        else:
+            raise ValueError('Contrastive sample: all{pos&neg} or positive(positive)')
+
+        # loss
+        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+        loss = loss.view(anchor_count, batch_size).mean()
+
+        return loss
+
+
+class InfoNCELossPatchLevel(nn.Module):
+    """
+    test: ref ConMIM: https://github.com/TencentARC/ConMIM.
+    """
+    def __init__(self, temperature=0.1, contrast_sample='all'):
+        """
+        :param temperature: 0.1 0.5 1.0, 1.5 2.0
+        """
+        super(InfoNCELossPatchLevel, self).__init__()
+        self.temperature = temperature
+        self.criterion = nn.CrossEntropyLoss()
+        self.contrast_sample = contrast_sample
+
+        self.facial_region_group = [
+            [2, 3],  # eyebrows
+            [4, 5],  # eyes
+            [6],  # nose
+            [7, 8, 9],  # mouth
+            [10, 1, 0],  # face boundaries
+            [10],  # hair
+            [1],  # facial skin
+            [0]  # background
+        ]
+
+    def forward(self, cl_features, parsing_map=None):
+        """
+        Args:
+            :param parsing_map:
+            :param cl_features: : hidden vector of shape [bsz, n_views, ...]
+        Returns:
+            A loss scalar.
+        """
+        device = (torch.device('cuda')
+                  if cl_features.is_cuda
+                  else torch.device('cpu'))
+
+        if len(cl_features.shape) < 4:
+            raise ValueError('`features` needs to be [bsz, n_views, n_cl_patches, ...],'
+                             'at least 4 dimensions are required')
+        if len(cl_features.shape) > 4:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], cl_features.shape[2], -1)
+            # [BS, 2, num_cl_patches, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0]
+        cl_features_2 = cl_features[:, 1]
+        score = torch.matmul(cl_features_1, cl_features_2.permute(0, 2, 1))  # [BS, num_cl_patches, num_cl_patches]
+        score = score / self.temperature
+        bs = score.size(0)
+        num_cl_patches = score.size(1)
+
+        if self.contrast_sample == 'all':
+            score = score
+        elif self.contrast_sample == 'positive':
+            mask = torch.eye(num_cl_patches, dtype=torch.float32)  # torch.Size([num_cl_patches, num_cl_patches])
+            mask_batch = mask.unsqueeze(0).expand(bs, -1).to(device)  # [bs, num_cl_patches, num_cl_patches]
+            score = score*mask_batch
+        elif self.contrast_sample == 'region':
+            cl_features_1_fr = []
+            cl_features_2_fr = []
+            for facial_region_index in self.facial_region_group:
+                fr_mask = (parsing_map == facial_region_index).unsqueeze(2).expand(-1, -1, cl_features_1.size(-1))
+                cl_features_1_fr.append((cl_features_1 * fr_mask).mean(dim=1, keepdim=False))
+                cl_features_2_fr.append((cl_features_1 * fr_mask).mean(dim=1, keepdim=False))
+            cl_features_1_fr = torch.stack(cl_features_1_fr, dim=1)
+            cl_features_2_fr = torch.stack(cl_features_2_fr, dim=1)
+            score = torch.matmul(cl_features_1_fr, cl_features_2_fr.permute(0, 2, 1))  # [BS, 8, 8]
+            score = score / self.temperature
+            # mask = torch.eye(cl_features_1_fr.size(1), dtype=torch.bool)
+            # torch.Size([cl_features_1_fr.size(1), cl_features_1_fr.size(1)])
+            # mask_batch = mask.unsqueeze(0).expand(bs, -1).to(device)
+            # [bs, cl_features_1_fr.size(1), cl_features_1_fr.size(1)]
+            # score = score*mask_batch
+            label = torch.arange(cl_features_1_fr.size(1), dtype=torch.long).to(device)
+            labels_batch = label.unsqueeze(0).expand(bs, -1)
+            loss = 2 * self.temperature * self.criterion(score, labels_batch)
+            return loss
+        else:
+            raise ValueError('Contrastive sample: all{pos&neg} or positive(positive)')
+
+        # label = (torch.arange(bs, dtype=torch.long) +
+        #          bs * torch.distributed.get_rank()).to(device)
+        label = torch.arange(num_cl_patches, dtype=torch.long).to(device)
+        labels_batch = label.unsqueeze(0).expand(bs, -1)
+
+        loss = 2 * self.temperature * self.criterion(score, labels_batch)
+
+        return loss
+
+
+class MSELoss(nn.Module):
+    """
+    test: unused
+    """
+    def __init__(self):
+        super(MSELoss, self).__init__()
+
+    @staticmethod
+    def forward(cl_features):
+
+        if len(cl_features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, n_patches, ...],'
+                             'at least 3 dimensions are required')
+        if len(cl_features.shape) > 3:
+            cl_features = cl_features.view(cl_features.shape[0], cl_features.shape[1], -1)  # [BS, 2, feat_cl_dim]
+
+        cl_features_1 = cl_features[:, 0].float()  # [BS, feat_cl_dim]
+        cl_features_2 = cl_features[:, 1].float()  # [BS, feat_cl_dim]
+
+        return torch.nn.functional.mse_loss(cl_features_1, cl_features_2, reduction='mean')

util/lr_decay.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import json
+
+
+def param_groups_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
+    """
+    Parameter groups for layer-wise lr decay
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
+    """
+    param_group_names = {}
+    param_groups = {}
+
+    num_layers = len(model.blocks) + 1
+
+    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))
+
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+
+        # no decay: all 1D parameters and model specific ones
+        if p.ndim == 1 or n in no_weight_decay_list:
+            g_decay = "no_decay"
+            this_decay = 0.
+        else:
+            g_decay = "decay"
+            this_decay = weight_decay
+            
+        layer_id = get_layer_id_for_vit(n, num_layers)
+        group_name = "layer_%d_%s" % (layer_id, g_decay)
+
+        if group_name not in param_group_names:
+            this_scale = layer_scales[layer_id]
+
+            param_group_names[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+            param_groups[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+
+        param_group_names[group_name]["params"].append(n)
+        param_groups[group_name]["params"].append(p)
+
+    # print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
+
+    return list(param_groups.values())
+
+
+def get_layer_id_for_vit(name, num_layers):
+    """
+    Assign a parameter with its layer id
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
+    """
+    if name in ['cls_token', 'pos_embed']:
+        return 0
+    elif name.startswith('patch_embed'):
+        return 0
+    elif name.startswith('blocks'):
+        return int(name.split('.')[1]) + 1
+    else:
+        return num_layers

util/lr_sched.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import math
+import numpy as np
+
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < args.warmup_epochs:
+        lr = args.lr * epoch / args.warmup_epochs 
+    else:
+        lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
+            (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr
+    return lr
+
+
+def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0,
+                     start_warmup_value=0, warmup_steps=-1):
+    warmup_schedule = np.array([])
+    warmup_iters = warmup_epochs * niter_per_ep
+    if warmup_steps > 0:
+        warmup_iters = warmup_steps
+    print("Set warmup steps = %d" % warmup_iters)
+    if warmup_epochs > 0:
+        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+    iters = np.arange(epochs * niter_per_ep - warmup_iters)
+    schedule = np.array(
+        [final_value + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters])
+
+    schedule = np.concatenate((warmup_schedule, schedule))
+
+    assert len(schedule) == epochs * niter_per_ep
+    return schedule

util/metrics.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+from sklearn.metrics import roc_auc_score
+from sklearn.metrics import roc_curve
+from sklearn.metrics import auc, accuracy_score, balanced_accuracy_score
+from scipy.optimize import brentq
+from scipy.interpolate import interp1d
+
+
+def frame_level_acc(labels, y_preds):
+    return accuracy_score(labels, y_preds) * 100.
+
+
+def frame_level_balanced_acc(labels, y_preds):
+    return balanced_accuracy_score(labels, y_preds) * 100.
+
+
+def frame_level_auc(labels, preds):
+    return roc_auc_score(labels, preds) * 100.
+
+
+def frame_level_eer(labels, preds):
+    # 推荐；更正确的，MaskRelation（TIFS23也是）
+    fpr, tpr, thresholds = roc_curve(labels, preds, pos_label=1)  # 如果标签不是二进制的，则应显式地给出pos_标签
+    eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
+    # eer_thresh = interp1d(fpr, thresholds)(eer)
+    return eer
+
+
+# def frame_level_eer(labels, preds):
+#     fpr, tpr, thresholds = roc_curve(labels, preds, pos_label=1)
+#     eer_threshold = thresholds[(fpr + (1 - tpr)).argmin()]
+#     fpr_eer = fpr[thresholds == eer_threshold][0]
+#     fnr_eer = 1 - tpr[thresholds == eer_threshold][0]
+#     eer = (fpr_eer + fnr_eer) / 2
+#     metric_logger.meters['eer'].update(eer)
+#     return eer, eer_thresh
+
+
+def get_video_level_label_pred(f_label_list, v_name_list, f_pred_list):
+    """
+    References:
+    CADDM: https://github.com/megvii-research/CADDM
+    """
+    video_res_dict = dict()
+    video_pred_list = list()
+    video_y_pred_list = list()
+    video_label_list = list()
+    # summarize all the results for each video
+    for label, video, score in zip(f_label_list, v_name_list, f_pred_list):
+        if video not in video_res_dict.keys():
+            video_res_dict[video] = {"scores": [score], "label": label}
+        else:
+            video_res_dict[video]["scores"].append(score)
+    # get the score and label for each video
+    for video, res in video_res_dict.items():
+        score = sum(res['scores']) / len(res['scores'])
+        label = res['label']
+        video_pred_list.append(score)
+        video_label_list.append(label)
+        video_y_pred_list.append(score >= 0.5)
+
+    return video_label_list, video_pred_list, video_y_pred_list
+
+
+def video_level_acc(video_label_list, video_y_pred_list):
+    return accuracy_score(video_label_list, video_y_pred_list) * 100.
+
+
+def video_level_balanced_acc(video_label_list, video_y_pred_list):
+    return balanced_accuracy_score(video_label_list, video_y_pred_list) * 100.
+
+
+def video_level_auc(video_label_list, video_pred_list):
+    return roc_auc_score(video_label_list, video_pred_list) * 100.
+
+
+def video_level_eer(video_label_list, video_pred_list):
+    # 推荐；更正确的，MaskRelation（TIFS23也是）
+    fpr, tpr, thresholds = roc_curve(video_label_list, video_pred_list, pos_label=1)  # 如果标签不是二进制的，则应显式地给出pos_标签
+    v_eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
+    # eer_thresh = interp1d(fpr, thresholds)(eer)
+    return v_eer

util/misc.py

@@ -0,0 +1,390 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import os
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch import inf
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('max mem: {memory:.0f}')
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print('[{}] '.format(now), end='')  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if args.dist_on_itp:
+        args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+        args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
+        args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+        args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+        os.environ['LOCAL_RANK'] = str(args.gpu)
+        os.environ['RANK'] = str(args.rank)
+        os.environ['WORLD_SIZE'] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}, gpu {}'.format(
+        args.rank, args.dist_url, args.gpu), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]),
+                                norm_type)
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, tag=None):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        if tag is None:
+            checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+        else:
+            checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % tag)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        if tag is None:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name,
+                                  client_state=client_state)
+        else:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % tag,
+                                  client_state=client_state)
+
+
+def save_model_target_encoder(args, epoch, model, model_target_encoder_without_ddp, optimizer, loss_scaler, tag=None):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        if tag is None:
+            checkpoint_paths = [output_dir / ('checkpoint-te-%s.pth' % epoch_name)]
+        else:
+            checkpoint_paths = [output_dir / ('checkpoint-te-%s.pth' % tag)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_target_encoder_without_ddp.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        if tag is None:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-te-%s" % epoch_name,
+                                  client_state=client_state)
+        else:
+            model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-te-%s" % tag,
+                                  client_state=client_state)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        print("Resume checkpoint %s" % args.resume)
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
+
+def load_model_target_encoder(args, model_target_encoder_without_ddp, optimizer, loss_scaler):
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume_target_encoder, map_location='cpu')
+        model_target_encoder_without_ddp.load_state_dict(checkpoint['model'])
+        print("Resume checkpoint %s" % args.resume_target_encoder)
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
+
+def all_reduce_mean(x):
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x

util/pos_embed.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+# Author: Gaojian Wang@ZJUICSR
+# --------------------------------------------------------
+# This source code is licensed under the Attribution-NonCommercial 4.0 International License.
+# You can find the license in the LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+
+import numpy as np
+import torch
+
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed
+
+
+def interpolate_pos_embed_ema(model, checkpoint_model):
+    if checkpoint_model.pos_embed is not None:
+        pos_embed_checkpoint = checkpoint_model.pos_embed
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model.pos_embed = new_pos_embed