patch_m2d.diff

--- _org/train_audio.py	2023-07-07 10:22:30.225586716 +0900
+++ train_audio.py	2026-01-26 15:02:27.199236431 +0900
@@ -1,3 +1,9 @@
+"""Masked Modeling Duo (M2D) Pre-training Script V2
+
+Masked Modeling Duo: Towards a Universal Audio Pre-training Framework
+https://arxiv.org/abs/2404.06095
+"""
+
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 
@@ -15,69 +21,117 @@
 import os
 import time
 from pathlib import Path
+import subprocess
+import sys
 
 import torch
 import torch.backends.cudnn as cudnn
 from torch.utils.tensorboard import SummaryWriter
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
+from functools import partial
+import matplotlib.pyplot as plt
 
-import timm
-
-assert timm.__version__ == "0.3.2"  # version check
 import timm.optim.optim_factory as optim_factory
 
 import util.misc as misc
 from util.misc import NativeScalerWithGradNormCount as NativeScaler
 
-import models_mae
-
-from engine_pretrain import train_one_epoch
+from m2d import models_mae
+from m2d.engine_pretrain_m2d import train_one_epoch_m2dx
+import audio_dataset
+import common
+from m2d.runtime_audio import RuntimeM2D
 
 
 def get_args_parser():
-    parser = argparse.ArgumentParser('MAE pre-training', add_help=False)
-    parser.add_argument('--batch_size', default=64, type=int,
+    parser = argparse.ArgumentParser('Masked Modeling Duo (M2D) pre-training V2', add_help=False)
+    parser.add_argument('--batch_size', default=512, type=int,
                         help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
-    parser.add_argument('--epochs', default=400, type=int)
+    parser.add_argument('--epochs', default=300, type=int)
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+    parser.add_argument('--eval_after', default=50, type=int)
+    parser.add_argument('--save_freq', default=100, type=int)
+    parser.add_argument('--feature_eval_freq', default=10, type=int, help='Feature label-free evaluation frequency.')
+    parser.add_argument('--stop_at', default=-1, type=int)
 
     # Model parameters
-    parser.add_argument('--model', default='mae_vit_large_patch16', type=str, metavar='MODEL',
-                        help='Name of model to train')
+    parser.add_argument('--model', default='m2d_vit_base', type=str, metavar='MODEL', help='Model name.')
+    parser.add_argument('--decoder_depth', type=int, default=8, metavar='DD', help='Model decoder depth.')
 
-    parser.add_argument('--input_size', default=224, type=int,
-                        help='images input size')
+    parser.add_argument('--input_size', default='80x608', type=str, help='images input size')
+    parser.add_argument('--patch_size', default='16x16', type=str, help='patch size')
+    parser.add_argument('--sr', default='16k', type=str, metavar='SR', help='Sampling rate of the input audio.')
 
-    parser.add_argument('--mask_ratio', default=0.75, type=float,
+    parser.add_argument('--mask_ratio', default=0.7, type=float,
                         help='Masking ratio (percentage of removed patches).')
+    parser.add_argument('--weighted_mask', action='store_true',
+                        help='Use weighted masking based on data contents.')
+    parser.add_argument('--var_weighted_mask', action='store_true',
+                        help='Use weighted masking (var) based on data contents.')
+    parser.add_argument('--block_mask', action='store_true',
+                        help='Use block masking from data2vec 2.0.')
+    parser.add_argument('--obj_cent_mask', action='store_true',
+                        help='Use object centric masking based on data contents.')
+    parser.add_argument('--masking_hint_k', default=2, type=int)
+    parser.add_argument('--masking_start_epoch', default=50, type=int)
+    parser.add_argument('--abs_hint_ratio', default=-1.0, type=float)  # set value >= 0.0 to enable absolute hint ratio
+
+    parser.add_argument('--ema_decay_init', default=0.99995, type=float,
+                        help='Initial EMA decay parameter.')
+    parser.add_argument('--ema_decay', default=0.99999, type=float,
+                        help='EMA decay parameter.')
+    parser.add_argument('--loss_fn', default='norm_mse', type=str,
+                        help='loss function: mse or norm_mse.')
+    parser.add_argument('--loss_m2d', default=1., type=float, help='Loss of M2D masked prediction')
+    parser.add_argument('--loss_off', default=0., type=float, help='Loss of offline target')
+    parser.add_argument('--target_layers', default='', type=str,
+                        help='Experimental: layers to calculate target representations.')
+    parser.add_argument('--bf16', action='store_true')
+    parser.add_argument('--teacher', default='', type=str, help='Weight path of the teacher M2D model.')
 
-    parser.add_argument('--norm_pix_loss', action='store_true',
+    parser.add_argument('--no_norm_pix_loss', action='store_false', dest='norm_pix_loss',
                         help='Use (per-patch) normalized pixels as targets for computing loss')
-    parser.set_defaults(norm_pix_loss=False)
+    parser.set_defaults(norm_pix_loss=True)
+
+    parser.add_argument('--cont_mask', default=0, type=int,
+                        help='Use random 1-d (continuous) masking scheme. 0:off, 0<:# of frames to make them continuous.')
 
     # Optimizer parameters
     parser.add_argument('--weight_decay', type=float, default=0.05,
                         help='weight decay (default: 0.05)')
+    parser.add_argument('--clip_grad', type=float, default=3.0, metavar="NORM",
+                        help="Clip gradient norm (default: None, no clipping)")
+    parser.add_argument('--optim', default='adamw', type=str, help='Optimizer adam or sdg')
 
     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',
+    parser.add_argument('--blr', type=float, default=3e-4, metavar='LR',
                         help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
     parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                         help='lower lr bound for cyclic schedulers that hit 0')
 
-    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
+    parser.add_argument('--warmup_epochs', type=int, default=20, metavar='N',
                         help='epochs to warmup LR')
 
     # Dataset parameters
-    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
+    parser.add_argument('--data_path', default='data', type=str,
                         help='dataset path')
+    parser.add_argument('--csv_main', default='data/files_audioset.csv', type=str,
+                        help='A CSV file to list sample files in the main dataset')
+    parser.add_argument('--csv_bg_noise', default='', type=str,
+                        help='A CSV file to list sample files in the BG noise dataset')
+    parser.add_argument('--csv_val', default='', type=str,
+                        help='A CSV file to list validation sample files')
+    parser.add_argument('--min_ds_size', default=10000, type=int,
+                        help='Inflate the size of the smaller dataset to the desired size')
+    parser.add_argument('--norm_stats', default='None', type=str,  # Will be computed runtime.
+                        help='dataset normalization stats')
+    parser.add_argument('--noise_ratio', default=0., type=float,
+                        help='Noise mixing ratio')
 
-    parser.add_argument('--output_dir', default='./output_dir',
+    parser.add_argument('--output_dir', default='',
                         help='path where to save, empty for no saving')
-    parser.add_argument('--log_dir', default='./output_dir',
+    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')
@@ -87,7 +141,9 @@
 
     parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                         help='start epoch')
-    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--force_start_epoch', default=0, type=int, metavar='N',  # 0=always reset start epoch to 0 even if using --resume
+                        help='start epoch for resuming')
+    parser.add_argument('--num_workers', default=8, 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')
@@ -96,7 +152,7 @@
     # 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('--local-rank' if torch.__version__ >= "2.0.0" else '--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')
@@ -104,9 +160,38 @@
     return parser
 
 
-def main(args):
-    misc.init_distributed_mode(args)
+def ema_decay_sched(step, total_steps, ema_decay_init, ema_decay):
+    interp = step / (total_steps - 1)
+    tau = ema_decay_init + (ema_decay - ema_decay_init) * interp
+    return tau
+
+
+def get_optim(args, param_groups):
+    if args.optim == 'adamw':
+        return torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    elif args.optim == 'sgd':
+        return torch.optim.SGD(param_groups, args.lr, momentum=0.9, weight_decay=0)
+    assert False, f'Unsupported optimizer {args.optim}'
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler, delta_epoch=1, strict=True):
+    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'], strict=strict)
+        print("Resume checkpoint %s" % args.resume)
+        if strict == True and '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'] + delta_epoch
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
 
+
+def main(args):
     print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
     print("{}".format(args).replace(', ', ',\n'))
 
@@ -119,14 +204,21 @@
 
     cudnn.benchmark = True
 
-    # simple augmentation
-    transform_train = transforms.Compose([
-            transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 3 is bicubic
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
-    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
-    print(dataset_train)
+    dataset_train, dataset_val = audio_dataset.build_mixed_dataset(args)
+    print(dataset_train, dataset_val)
+
+    if args.teacher == '':
+        teacher_model, off_emb_dim = None, 3840
+        print(f' ** No teacher with loss_off=={args.loss_off} > 0. **')
+    else:
+        print(f' ** M2D-X **')
+        teacher_model = RuntimeM2D(weight_file=args.teacher)
+        off_emb_dim = teacher_model.cfg.feature_d
+        models_mae.set_requires_grad(teacher_model, False)
+        teacher_model.to(device)
+        teacher_model.eval()
+        print('Teacher weight =', args.teacher, 'feature_d =', off_emb_dim)
+        print("Teacher = %s" % common.short_model_desc(teacher_model))
 
     if True:  # args.distributed:
         num_tasks = misc.get_world_size()
@@ -143,6 +235,8 @@
         log_writer = SummaryWriter(log_dir=args.log_dir)
     else:
         log_writer = None
+    common.PrintLogger(f'{args.log_dir}/console.txt')
+    print(args)
 
     data_loader_train = torch.utils.data.DataLoader(
         dataset_train, sampler=sampler_train,
@@ -151,53 +245,104 @@
         pin_memory=args.pin_mem,
         drop_last=True,
     )
+
+    # for validation loss
+    if args.csv_val != '':
+        sampler_val = torch.utils.data.DistributedSampler(dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+        print("Sampler_val = %s" % str(sampler_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)
+    else:
+        data_loader_val = None
     
     # define the model
-    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
+    model = models_mae.__dict__[args.model](img_size=args.input_size, patch_size=args.patch_size, decoder_depth=args.decoder_depth,
+        norm_pix_loss=args.norm_pix_loss, loss_type=args.loss_fn, target_layers=args.target_layers, loss_m2d=args.loss_m2d, loss_off=args.loss_off,
+        off_emb_dim=off_emb_dim, norm_stats=dataset_train.norm_stats, world_rank=[num_tasks, global_rank])
+
+    if args.weighted_mask:
+        assert (args.cont_mask == 0) and (not args.var_weighted_mask) and (not args.obj_cent_mask), "Masking options conflict."
+        model.set_random_weighted_mask_with_hint()
+    if args.var_weighted_mask:
+        assert (args.cont_mask == 0) and (not args.weighted_mask) and (not args.obj_cent_mask), "Masking options conflict."
+        model.set_random_weighted_mask_var_with_hint()
+    if args.obj_cent_mask:
+        assert (args.cont_mask == 0) and (not args.weighted_mask) and (not args.var_weighted_mask), "Masking options conflict."
+        if args.masking_start_epoch == 0:
+            print(f"⚠️  Warning: Too early to start object centric masking.")
+        model.set_object_centric_mask_with_hint()
+    if args.block_mask:
+        model.set_data2vec2_block_masking()
+    if args.cont_mask > 0:
+        model.set_random_1d_mask(args.cont_mask)
 
     model.to(device)
 
     model_without_ddp = model
     print("Model = %s" % str(model_without_ddp))
 
-    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
-    
+    org_args_lr = args.lr
     if args.lr is None:  # only base_lr is specified
-        args.lr = args.blr * eff_batch_size / 256
+        args.lr = args.blr * args.eff_batch_size / 256
 
-    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
+    print("base lr: %.2e" % (args.lr * 256 / args.eff_batch_size) if org_args_lr is None else 'base lr: not effective')
     print("actual lr: %.2e" % args.lr)
 
     print("accumulate grad iterations: %d" % args.accum_iter)
-    print("effective batch size: %d" % eff_batch_size)
+    print("effective batch size: %d" % args.eff_batch_size)
 
     if args.distributed:
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
         model_without_ddp = model.module
-    
+
     # following timm: set wd as 0 for bias and norm layers
-    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
-    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    try:
+        param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    except:
+        print(' (for compatibility with timm) Switched add_weight_decay() to param_groups_weight_decay()')
+        param_groups = optim_factory.param_groups_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = get_optim(args, param_groups)
     print(optimizer)
     loss_scaler = NativeScaler()
 
-    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
+    load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, delta_epoch=0, strict=False)
+
+    if args.force_start_epoch >= 0:
+        args.start_epoch = args.force_start_epoch
 
     print(f"Start training for {args.epochs} epochs")
     start_time = time.time()
+    last_subprocess = None
     for epoch in range(args.start_epoch, args.epochs):
+        epoch1 = epoch + 1
         if args.distributed:
             data_loader_train.sampler.set_epoch(epoch)
-        train_stats = train_one_epoch(
-            model, data_loader_train,
+        if args.weighted_mask or args.var_weighted_mask or args.block_mask or args.obj_cent_mask:
+            if args.abs_hint_ratio < 0.0:
+                model_without_ddp.set_hint_ratio_by_epoch(cur_epoch=epoch, total_epochs=args.epochs, start_epoch=args.masking_start_epoch, k=args.masking_hint_k)
+            else:
+                model_without_ddp.set_hint_ratio(args.abs_hint_ratio)
+        train_stats = train_one_epoch_m2dx(
+            model, teacher_model, data_loader_train,
             optimizer, device, epoch, loss_scaler,
+            partial(ema_decay_sched, total_steps=len(data_loader_train) * args.epochs,
+                ema_decay_init=args.ema_decay_init, ema_decay=args.ema_decay),
+            val_loader=data_loader_val,
             log_writer=log_writer,
+            do_analysis=(epoch1 % args.feature_eval_freq == 0),
+            autocast_args=dict(dtype=torch.bfloat16) if args.bf16 else {},
             args=args
         )
-        if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
+
+        if args.output_dir and (epoch1 % args.save_freq == 0 or epoch1 == args.epochs):
             misc.save_model(
                 args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                loss_scaler=loss_scaler, epoch=epoch)
+                loss_scaler=loss_scaler, epoch=epoch1)
+            # run the external evaluator
+            if args.eval_after <= epoch1 and epoch1 < args.epochs and misc.is_main_process():
+                abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+                print('quick_eval', abspath)
+                last_subprocess = subprocess.Popen(['/bin/bash', './quick_eval.sh', abspath])
 
         log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                         'epoch': epoch,}
@@ -207,15 +352,80 @@
                 log_writer.flush()
             with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                 f.write(json.dumps(log_stats) + "\n")
+        
+        if args.stop_at > 0 and epoch1 >= args.stop_at:
+            if last_subprocess is not None:
+                last_subprocess.wait()
+            print(f'Stop training by reaching args.stop_at epoch: {args.stop_at}')
+            exit(0)
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
     print('Training time {}'.format(total_time_str))
 
+    del model_without_ddp, model, data_loader_train, optimizer, loss_scaler
+    if misc.is_main_process():
+        abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+        subprocess.call(['/bin/bash', './all_eval.sh', abspath])
+    return epoch1
+
+
+arg_conf_defaults = {
+    'csv_main': ('data/files_audioset.csv', 'M', 'path'),
+    'csv_bg_noise': ('', 'D', 'path'),
+    'ema_decay_init': (0.99995, 'ema', 'z'),
+    'ema_decay': (0.99999, 'ed', 'z'),
+    'decoder_depth': (8, 'dd', 'asis'),
+    'mask_ratio': (0.7, 'mr', 'z'),
+    'seed': (0, 's', 'asis'),
+    'norm_pix_loss':  (True, '~N', 'b'),
+    'loss_fn': ('norm_mse', 'L', 'head'),
+    'optim': ('adamw', 'O', 'asis'),
+    'warmup_epochs': (20, 'wu', 'asis'),
+    'blr': (3e-4, 'blr', 'z'),
+    'lr': (None, 'lr', 'z'),
+    'target_layers': ('', 'T', 'b'),
+    'eff_batch_size': (2048, 'bs', 'asis'),
+    'accum_iter': (1, 'a', 'asis'),
+    'loss_m2d': (1.0, 'lm', 'z'),
+    'loss_off': (0.0, 'lo', 'z'),
+    'noise_ratio': (0.0, 'nr', 'z'),
+    'min_ds_size': (10000, 'dn', 'asis'),
+    'cont_mask': (0, 'C', 'asis'),
+    'weighted_mask': (False, 'Wm', 'b'),
+    'var_weighted_mask': (False, 'Vm', 'b'),
+    'block_mask': (False, 'Bm', 'b'),
+    'obj_cent_mask': (False, 'Om', 'b'),
+    'masking_hint_k': [2, 'k', 'asis'],
+    'masking_start_epoch': [50, 'i', 'asis'],
+    'abs_hint_ratio': [-1.0, 'a', 'asis'],
+    'epochs': (0, '-e', 'asis'),
+}
 
-if __name__ == '__main__':
+
+def complete_args():
     args = get_args_parser()
     args = args.parse_args()
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    _input_size, _patch_size = args.input_size, args.patch_size
+    args.input_size = [int(x) for x in args.input_size.split('x')]
+    args.patch_size = [int(x) for x in args.patch_size.split('x')]
+    args.norm_stats = eval(args.norm_stats) if args.norm_stats else None
+
+    misc.init_distributed_mode(args)
+
+    args.eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
+
+    if not args.output_dir:
+        args.output_dir = f'{args.model}-{_input_size}p{_patch_size}p{args.sr}'
+        args.output_dir += f'-{common.get_timestamp()[:6]}-{common.arg_conf_str(args, defaults=arg_conf_defaults)}'
+
+    if not args.log_dir:
+        args.log_dir = args.output_dir
+    args.target_layers = None if args.target_layers == '' else eval(args.target_layers)
+    return args
+
+
+if __name__ == '__main__':
+    args = complete_args()
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
     main(args)
--- _org/clap/train_clap.py	2024-06-04 08:28:27.348036827 +0900
+++ clap/train_clap.py	2026-01-26 15:02:38.459294984 +0900
@@ -1,3 +1,12 @@
+"""M2D-CLAP Pre-training Script for Stage 1
+
+- M2D-CLAP: Exploring General-purpose Audio-Language Representations Beyond CLAP
+https://ieeexplore.ieee.org/document/11168481
+
+- M2D-CLAP: Masked Modeling Duo Meets CLAP for Learning General-purpose Audio-Language Representation
+https://www.isca-archive.org/interspeech_2024/niizumi24_interspeech.html
+"""
+
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 
@@ -15,79 +24,250 @@
 import os
 import time
 from pathlib import Path
+import subprocess
+import sys
+import random
 
 import torch
 import torch.backends.cudnn as cudnn
 from torch.utils.tensorboard import SummaryWriter
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
-
-import timm
+from functools import partial
+import matplotlib.pyplot as plt
 
-assert timm.__version__ == "0.3.2"  # version check
 import timm.optim.optim_factory as optim_factory
 
 import util.misc as misc
 from util.misc import NativeScalerWithGradNormCount as NativeScaler
 
-import models_mae
+from m2d import models_mae
 
-from engine_pretrain import train_one_epoch
+from m2d.engine_pretrain_m2d import train_one_epoch_m2dx
+from audio_dataset import SpectrogramDataset, get_files_no_sort, log_mixup_exp, pd
+import common
+from m2d.runtime_audio import RuntimeM2D
+
+
+class CapEmbSpecDataset(torch.utils.data.Dataset):
+    def __init__(self, base_folder, files_main, files_bg_noise, file_caption, crop_size, noise_ratio=0.0,
+                 random_crop=True, n_norm_calc=10000) -> None:
+        super().__init__()
+
+        self.ds1 = SpectrogramDataset(folder=base_folder, files=files_main, crop_frames=crop_size[1],
+                random_crop=random_crop, norm_stats=None, n_norm_calc=n_norm_calc//2)
+        self.norm_stats = self.ds1.norm_stats  # for compatibility with SpectrogramDataset
+        self.norm_std = self.ds1.norm_stats[1]
+        self.ds1.norm_stats = (self.ds1.norm_stats[0], 1.0)
+
+        if noise_ratio > 0.0:
+            self.ds2 = SpectrogramDataset(folder=base_folder, files=files_bg_noise, crop_frames=crop_size[1],
+                    random_crop=random_crop, norm_stats=None, n_norm_calc=n_norm_calc//2, repeat_short=True)
+            self.ds2.norm_stats = (self.ds2.norm_stats[0], 1.0) # disable normalizion scaling in the ds2
+        # for BG noise
+        self.noise_ratio = noise_ratio
+        self.bg_index = []
+        # find caption files fron folder
+        if Path(file_caption).is_dir():
+            file_caption = ','.join(sorted([str(f) for f in Path(file_caption).glob('*.npy')]))
+        # load captions
+        self.capembs = []
+        for capfile in file_caption.split(','):
+            embs = np.load(capfile, allow_pickle=True).item()
+            self.capembs.append(embs)
+            print('Caption', Path(capfile).stem, '-> embeddings:', len(embs), ' expample keys:', list(embs.keys())[:5])
+            assert list(self.capembs[0].values())[0].shape[-1] == list(self.capembs[-1].values())[0].shape[-1], 'All captions need to have the same number of embedding dimensions.'
+        self.file_caption = file_caption
+        self.caption_dim = list(self.capembs[0].values())[0].shape[-1]
+        print('Using', len(self.capembs), 'caption definitions with feature dimension:', self.caption_dim)
+
+    def __len__(self):
+        return len(self.ds1)
+
+    def get_random_caption(self, ytid, folder_stem):
+        emb_list = []
+        # by the YouTube ID
+        for capemb in self.capembs:
+            if ytid in capemb:
+                cur = capemb[ytid]
+                cur = [cur] if len(cur.shape) == 1 else [cur[i] for i in range(cur.shape[0])]  # prepare list of captions
+                emb_list.extend(cur)
+        # by the stem of the filename
+        if len(emb_list) == 0:
+            for capemb in self.capembs:
+                if folder_stem in capemb:
+                    cur = capemb[folder_stem]
+                    cur = [cur] if len(cur.shape) == 1 else [cur[i] for i in range(cur.shape[0])]  # prepare list of captions
+                    emb_list.extend(cur)
+        if not emb_list:
+            print('MISSING', ytid, folder_stem)
+        return random.choice(emb_list)
+
+    def __getitem__(self, index, fixed_noise=False):
+        # load index sample
+        clean = self.ds1[index]
+        if self.noise_ratio > 0.0:
+            # load random noise sample ### , while making noise floor zero
+            noise = self.ds2[index if fixed_noise else self.get_next_bgidx()]
+            # mix
+            mixed = log_mixup_exp(noise, clean, self.noise_ratio) if self.noise_ratio < 1.0 else noise
+        else:
+            mixed = clean
+        # finalize normalization. clean and noise were averaged to zero. the following will scale to 1.0 using ds1 std.
+        mixed = mixed / self.norm_std
+
+        # load sample's caption
+        filepath = self.ds1.df.file_name.values[index].split('/')
+        ytid = filepath[-1][:11]
+        stem = Path(filepath[-1]).stem
+        folder = {'AudioSet_SL_flac': 'AudioSet_SL', 'BBC_Sound_Effects_flac': 'BBC_Sound_Effects',
+            'FreeSound_flac': 'FreeSound', 'SoundBible_flac': 'SoundBible',
+        }[filepath[-2]] if filepath[0] == 'wavcaps_lms' else ''  # folder name conversion for wavcaps_lms
+        if filepath[0] == 'clotho_lms':
+            stem = stem.replace(':', '_').replace('?', '_')
+        caption = self.get_random_caption(ytid, folder + stem)
+
+        return mixed, caption
+
+    def get_next_bgidx(self):
+        if len(self.bg_index) == 0:
+            self.bg_index = torch.randperm(len(self.ds2)).tolist()
+            # print(f'Refreshed the bg index list with {len(self.bg_index)} items: {self.bg_index[:5]}...')
+        return self.bg_index.pop(0)
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + f'(crop_frames={self.ds1.crop_frames}, '
+        format_string += f'folder_main={self.ds1.df.file_name.values[0].split("/")[0]}, '
+        if self.noise_ratio > 0.: format_string += f'folder_bg={self.ds2.df.file_name.values[0].split("/")[0]}, '
+        format_string += f'caption={self.file_caption})'
+        return format_string
+
+
+def build_captioned_dataset(cfg):
+    # create dataset of (spectrogram, caption) pairs
+    files_main = get_files_no_sort(cfg.csv_main)
+    files_bg = get_files_no_sort(cfg.csv_bg_noise) if cfg.noise_ratio > 0. else []
+
+    ds = CapEmbSpecDataset(
+        base_folder=cfg.data_path, files_main=files_main,
+        files_bg_noise=files_bg,
+        file_caption=cfg.file_caption,
+        crop_size=cfg.input_size,
+        noise_ratio=cfg.noise_ratio,
+        random_crop=True)
+
+    val_ds = SpectrogramDataset(folder=cfg.data_path, files=get_files_no_sort(cfg.csv_val), crop_frames=cfg.input_size[1], random_crop=True) \
+        if cfg.csv_val else None
+
+    return ds, val_ds
 
 
 def get_args_parser():
-    parser = argparse.ArgumentParser('MAE pre-training', add_help=False)
-    parser.add_argument('--batch_size', default=64, type=int,
+    parser = argparse.ArgumentParser('M2D-CLAP pre-training', add_help=False)
+    parser.add_argument('--batch_size', default=512, type=int,
                         help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
-    parser.add_argument('--epochs', default=400, type=int)
+    parser.add_argument('--epochs', default=300, type=int)
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+    parser.add_argument('--eval_after', default=50, type=int)
+    parser.add_argument('--save_freq', default=100, type=int)
+    parser.add_argument('--feature_eval_freq', default=10, type=int, help='Feature label-free evaluation frequency.')
+    parser.add_argument('--stop_at', default=-1, type=int)
 
     # Model parameters
-    parser.add_argument('--model', default='mae_vit_large_patch16', type=str, metavar='MODEL',
-                        help='Name of model to train')
+    parser.add_argument('--model', default='m2d_clap_vit_base', type=str, metavar='MODEL', help='Model name.')
+    parser.add_argument('--decoder_depth', type=int, default=8, metavar='DD', help='Model decoder depth.')
 
-    parser.add_argument('--input_size', default=224, type=int,
-                        help='images input size')
+    parser.add_argument('--input_size', default='80x608', type=str, help='images input size')
+    parser.add_argument('--patch_size', default='16x16', type=str, help='patch size')
+    parser.add_argument('--sr', default='16k', type=str, metavar='SR', help='Sampling rate of the input audio.')  # just for adding to the output folder name
+    parser.add_argument('--text_encoder', default='B', type=str, metavar='SR', help='Text encoder type: A=GTE base, B=BERT base, N=NVEmbedV2.')
 
-    parser.add_argument('--mask_ratio', default=0.75, type=float,
+    parser.add_argument('--mask_ratio', default=0.7, type=float,
                         help='Masking ratio (percentage of removed patches).')
+    parser.add_argument('--weighted_mask', action='store_true',
+                        help='Use weighted masking based on data contents.')
+    parser.add_argument('--var_weighted_mask', action='store_true',
+                        help='Use weighted masking (var) based on data contents.')
+    parser.add_argument('--block_mask', action='store_true',
+                        help='Use block masking from data2vec 2.0.')
+    parser.add_argument('--obj_cent_mask', action='store_true',
+                        help='Use object centric masking based on data contents.')
+    parser.add_argument('--masking_hint_k', default=2, type=int)
+    parser.add_argument('--masking_start_epoch', default=50, type=int)
+    parser.add_argument('--abs_hint_ratio', default=-1.0, type=float)  # set value >= 0.0 to enable absolute hint ratio
+
+    parser.add_argument('--ema_decay_init', default=0.99995, type=float,
+                        help='Initial EMA decay parameter.')
+    parser.add_argument('--ema_decay', default=0.99999, type=float,
+                        help='EMA decay parameter.')
+    parser.add_argument('--loss_fn', default='norm_mse', type=str,
+                        help='loss function: mse or norm_mse.')
+    parser.add_argument('--loss_m2d', default=1., type=float, help='Loss of M2D masked prediction')
+    parser.add_argument('--loss_m2d_end', default=1., type=float, help='Loss of M2D masked prediction')
+    parser.add_argument('--loss_off', default=.01, type=float, help='Loss of offline target')
+    parser.add_argument('--loss_off_end', default=.01, type=float, help='Offline target loss at the end of training')
+    parser.add_argument('--target_layers', default='', type=str,
+                        help='Experimental: layers to calculate target representations.')
+    parser.add_argument('--bf16', action='store_true')
 
-    parser.add_argument('--norm_pix_loss', action='store_true',
+    parser.add_argument('--no_norm_pix_loss', action='store_false', dest='norm_pix_loss',
                         help='Use (per-patch) normalized pixels as targets for computing loss')
-    parser.set_defaults(norm_pix_loss=False)
+    parser.set_defaults(norm_pix_loss=True)
+
+    parser.add_argument('--cont_mask', default=0, type=int,
+                        help='Use random 1-d (continuous) masking scheme. 0:off, 0<:# of frames to make them continuous.')
+    parser.add_argument('--sem_mode', default=1, type=int,
+                        help='CLAP semantic mode (audio projector blocks).')
 
     # Optimizer parameters
     parser.add_argument('--weight_decay', type=float, default=0.05,
                         help='weight decay (default: 0.05)')
+    parser.add_argument('--clip_grad', type=float, default=3.0, metavar="NORM",
+                        help="Clip gradient norm (default: None, no clipping)")
+    parser.add_argument('--optim', default='adamw', type=str, help='Optimizer adam or sdg')
 
     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',
+    parser.add_argument('--blr', type=float, default=3e-4, metavar='LR',
                         help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
     parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                         help='lower lr bound for cyclic schedulers that hit 0')
 
-    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
+    parser.add_argument('--warmup_epochs', type=int, default=20, metavar='N',
                         help='epochs to warmup LR')
 
     # Dataset parameters
-    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
+    parser.add_argument('--data_path', default='data', type=str,
                         help='dataset path')
+    parser.add_argument('--csv_main', default='data/files_audioset.csv', type=str,
+                        help='A CSV file to list sample files in the main dataset')
+    parser.add_argument('--csv_bg_noise', default='', type=str,
+                        help='A CSV file to list sample files in the BG noise dataset')
+    parser.add_argument('--file_caption', default='data/capemb_GTEbase_Audo_A_C_D.npy', type=str,
+                        help='A npy file for the caption')
+    parser.add_argument('--csv_val', default='', type=str,
+                        help='A CSV file to list validation sample files')
+    parser.add_argument('--min_ds_size', default=10000, type=int,
+                        help='Inflate the size of the smaller dataset to the desired size')
+    parser.add_argument('--norm_stats', default='None', type=str,  # will be computed runtime.
+                        help='dataset normalization stats')
+    parser.add_argument('--noise_ratio', default=0., type=float,
+                        help='Noise mixing ratio')
 
-    parser.add_argument('--output_dir', default='./output_dir',
+    parser.add_argument('--output_dir', default='',
                         help='path where to save, empty for no saving')
-    parser.add_argument('--log_dir', default='./output_dir',
+    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('--resume', default='', help='resume from checkpoint')
+    parser.add_argument('--load', default='', help='load from a checkpoint regardless of model name')
 
     parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                         help='start epoch')
-    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--force_start_epoch', default=0, type=int, metavar='N',  # 0=always reset start epoch to 0 even if using --resume
+                        help='start epoch for resuming')
+    parser.add_argument('--num_workers', default=8, 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')
@@ -96,7 +276,7 @@
     # 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('--local-rank' if torch.__version__ >= "2.0.0" else '--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')
@@ -104,9 +284,50 @@
     return parser
 
 
-def main(args):
-    misc.init_distributed_mode(args)
+def ema_decay_sched(step, total_steps, ema_decay_init, ema_decay):
+    interp = step / (total_steps - 1)
+    tau = ema_decay_init + (ema_decay - ema_decay_init) * interp
+    return tau
+
+
+def get_optim(args, param_groups):
+    if args.optim == 'adamw':
+        return torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    elif args.optim == 'sgd':
+        return torch.optim.SGD(param_groups, args.lr, momentum=0.9, weight_decay=0)
+    assert False, f'Unsupported optimizer {args.optim}'
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler, delta_epoch=1, strict=True):
+    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'], strict=strict)
+        print("Resume checkpoint %s" % args.resume)
+        if strict == True and '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'] + delta_epoch
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+    if args.load:
+        weights = str(args.load).split(',')
+        checkpoint = torch.load(weights[0], map_location='cpu')['model']
+        msg = model_without_ddp.load_state_dict(checkpoint, strict=False)
+        print(f'Load encoder/decoder weights from {weights[0]} -- {msg}')
+
+        if len(weights) > 1:
+            checkpoint = torch.load(weights[1], map_location='cpu')['model']
+            proj_weights =  {k.replace('backbone.', ''): v for k, v in checkpoint.items() if k.startswith('backbone.audio_proj.')}
+            proj_weights |= {k.replace('backbone.', ''): v for k, v in checkpoint.items() if k.startswith('backbone.text_proj.')}
+            msg = model_without_ddp.load_state_dict(proj_weights, strict=False)
+            print(f'Load projector weights {weights[1]} -- loaded keys: {list(proj_weights.keys())}, unexpected keys: {msg.unexpected_keys}')
 
+
+def main(args):
     print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
     print("{}".format(args).replace(', ', ',\n'))
 
@@ -119,14 +340,8 @@
 
     cudnn.benchmark = True
 
-    # simple augmentation
-    transform_train = transforms.Compose([
-            transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 3 is bicubic
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
-    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
-    print(dataset_train)
+    dataset_train, dataset_val = build_captioned_dataset(args)
+    print(dataset_train, dataset_val)
 
     if True:  # args.distributed:
         num_tasks = misc.get_world_size()
@@ -143,6 +358,8 @@
         log_writer = SummaryWriter(log_dir=args.log_dir)
     else:
         log_writer = None
+    common.PrintLogger(f'{args.log_dir}/console.txt')
+    print(args)
 
     data_loader_train = torch.utils.data.DataLoader(
         dataset_train, sampler=sampler_train,
@@ -151,53 +368,104 @@
         pin_memory=args.pin_mem,
         drop_last=True,
     )
+
+    # for validation loss
+    if args.csv_val != '':
+        sampler_val = torch.utils.data.DistributedSampler(dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+        print("Sampler_val = %s" % str(sampler_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)
+    else:
+        data_loader_val = None
     
     # define the model
-    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
+    model = models_mae.__dict__[args.model](img_size=args.input_size, patch_size=args.patch_size, decoder_depth=args.decoder_depth,
+        norm_pix_loss=args.norm_pix_loss, loss_type=args.loss_fn, target_layers=args.target_layers, loss_m2d=args.loss_m2d, loss_off=args.loss_off, norm_stats=dataset_train.norm_stats,
+        off_emb_dim=dataset_train.caption_dim, world_rank=[num_tasks, global_rank], sem_mode=args.sem_mode)
+
+    if args.weighted_mask:
+        assert (args.cont_mask == 0) and (not args.var_weighted_mask) and (not args.obj_cent_mask), "Masking options conflict."
+        model.set_random_weighted_mask_with_hint()
+    if args.var_weighted_mask:
+        assert (args.cont_mask == 0) and (not args.weighted_mask) and (not args.obj_cent_mask), "Masking options conflict."
+        model.set_random_weighted_mask_var_with_hint()
+    if args.obj_cent_mask:
+        assert (args.cont_mask == 0) and (not args.weighted_mask) and (not args.var_weighted_mask), "Masking options conflict."
+        if args.masking_start_epoch == 0:
+            print(f"⚠️  Warning: Too early to start object centric masking.")
+        model.set_object_centric_mask_with_hint()
+    if args.block_mask:
+        model.set_data2vec2_block_masking()
+    if args.cont_mask > 0:
+        model.set_random_1d_mask(args.cont_mask)
 
     model.to(device)
 
     model_without_ddp = model
     print("Model = %s" % str(model_without_ddp))
 
-    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
-    
+    org_args_lr = args.lr
     if args.lr is None:  # only base_lr is specified
-        args.lr = args.blr * eff_batch_size / 256
+        args.lr = args.blr * args.eff_batch_size / 256
 
-    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
+    print("base lr: %.2e" % (args.lr * 256 / args.eff_batch_size) if org_args_lr is None else 'base lr: not effective')
     print("actual lr: %.2e" % args.lr)
 
     print("accumulate grad iterations: %d" % args.accum_iter)
-    print("effective batch size: %d" % eff_batch_size)
+    print("effective batch size: %d" % args.eff_batch_size)
 
     if args.distributed:
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
         model_without_ddp = model.module
-    
+
     # following timm: set wd as 0 for bias and norm layers
-    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
-    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    try:
+        param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    except:
+        print(' (for compatibility with timm) Switched add_weight_decay() to param_groups_weight_decay()')
+        param_groups = optim_factory.param_groups_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = get_optim(args, param_groups)
     print(optimizer)
     loss_scaler = NativeScaler()
 
-    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
+    load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, delta_epoch=0, strict=False)
+
+    if args.force_start_epoch >= 0:
+        args.start_epoch = args.force_start_epoch
 
     print(f"Start training for {args.epochs} epochs")
     start_time = time.time()
+    last_subprocess, cur_loss_m2d, cur_loss_off = None, args.loss_m2d, args.loss_off
     for epoch in range(args.start_epoch, args.epochs):
+        epoch1 = epoch + 1
         if args.distributed:
             data_loader_train.sampler.set_epoch(epoch)
-        train_stats = train_one_epoch(
-            model, data_loader_train,
+        if args.weighted_mask or args.var_weighted_mask or args.block_mask or args.obj_cent_mask:
+            if args.abs_hint_ratio < 0.0:
+                model_without_ddp.set_hint_ratio_by_epoch(cur_epoch=epoch, total_epochs=args.epochs, start_epoch=args.masking_start_epoch, k=args.masking_hint_k)
+            else:
+                model_without_ddp.set_hint_ratio(args.abs_hint_ratio)
+        train_stats = train_one_epoch_m2dx(
+            model, None, data_loader_train,
             optimizer, device, epoch, loss_scaler,
+            partial(ema_decay_sched, total_steps=len(data_loader_train) * args.epochs,
+                ema_decay_init=args.ema_decay_init, ema_decay=args.ema_decay),
+            val_loader=data_loader_val,
             log_writer=log_writer,
+            do_analysis=(epoch1 % args.feature_eval_freq == 0),
+            autocast_args=dict(dtype=torch.bfloat16) if args.bf16 else {},
             args=args
         )
-        if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
+
+        if args.output_dir and (epoch1 % args.save_freq == 0 or epoch1 == args.epochs):
             misc.save_model(
                 args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                loss_scaler=loss_scaler, epoch=epoch)
+                loss_scaler=loss_scaler, epoch=epoch1)
+            # run the external evaluator
+            if args.eval_after <= epoch1 and epoch1 < args.epochs and misc.is_main_process():
+                abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+                print('quick_eval', abspath)
+                last_subprocess = subprocess.Popen(['/bin/bash', './quick_eval.sh', abspath])
 
         log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                         'epoch': epoch,}
@@ -207,15 +475,91 @@
                 log_writer.flush()
             with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                 f.write(json.dumps(log_stats) + "\n")
+        
+        if args.stop_at > 0 and epoch1 >= args.stop_at:
+            if last_subprocess is not None:
+                last_subprocess.wait()
+            print(f'Stop training by reaching args.stop_at epoch: {args.stop_at}')
+            exit(0)
+
+        # loss lambda linear interpolation
+        if args.loss_m2d != args.loss_m2d_end:
+            cur_loss_m2d = (epoch1 * (args.loss_m2d_end - args.loss_m2d))/args.epochs + args.loss_m2d
+            model_without_ddp.loss_m2d = cur_loss_m2d
+        if args.loss_off != args.loss_off_end:
+            cur_loss_off = (epoch1 * (args.loss_off_end - args.loss_off))/args.epochs + args.loss_off
+            model_without_ddp.loss_off = cur_loss_off
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
     print('Training time {}'.format(total_time_str))
 
+    del model_without_ddp, model, data_loader_train, optimizer, loss_scaler
+    if misc.is_main_process():
+        abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+        subprocess.call(['/bin/bash', './all_eval.sh', abspath])
+    return epoch1
+
+
+arg_conf_defaults = {
+    'csv_main': ('data/files_audioset.csv', 'M', 'path'),
+    'csv_bg_noise': ('', 'D', 'path'),
+    'file_caption': ('', 'C', 'path'),
+    'ema_decay_init': (0.99995, 'ema', 'z'),
+    'ema_decay': (0.99999, 'ed', 'z'),
+    'decoder_depth': (8, 'dd', 'asis'),
+    'mask_ratio': (0.7, 'mr', 'z'),
+    'seed': (0, 's', 'asis'),
+    'norm_pix_loss':  (True, '~N', 'b'),
+    'loss_fn': ('norm_mse', 'L', 'head'),
+    'optim': ('adamw', 'O', 'asis'),
+    'warmup_epochs': (20, 'wu', 'asis'),
+    'blr': (3e-4, 'blr', 'z'),
+    'lr': (None, 'lr', 'z'),
+    'eff_batch_size': (2048, 'bs', 'asis'),
+    'accum_iter': (1, 'a', 'asis'),
+    'loss_m2d': (1.0, 'lm', 'z'),
+    'loss_m2d_end': (1.0, 'lme', 'z'),
+    'loss_off': (0.01, 'lo', 'z'),
+    'loss_off_end': (0.01, 'loe', 'z'),
+    'noise_ratio': (0.0, 'nr', 'z'),
+    'min_ds_size': (10000, 'dn', 'asis'),
+    'cont_mask': (0, 'C', 'asis'),
+    'weighted_mask': (False, 'Wm', 'b'),
+    'var_weighted_mask': (False, 'Vm', 'b'),
+    'block_mask': (False, 'Bm', 'b'),
+    'obj_cent_mask': (False, 'Om', 'b'),
+    'masking_hint_k': [2, 'k', 'asis'],
+    'masking_start_epoch': [50, 'i', 'asis'],
+    'abs_hint_ratio': [-1.0, 'a', 'asis'],
+    'sem_mode': (1, 'S', 'asis'),
+    'epochs': (0, '-e', 'asis'),
+}
 
-if __name__ == '__main__':
+
+def complete_args():
     args = get_args_parser()
     args = args.parse_args()
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    _input_size, _patch_size = args.input_size, args.patch_size
+    args.input_size = [int(x) for x in args.input_size.split('x')]
+    args.patch_size = [int(x) for x in args.patch_size.split('x')]
+    args.norm_stats = eval(args.norm_stats) if args.norm_stats else None
+
+    misc.init_distributed_mode(args)
+
+    args.eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
+
+    if not args.output_dir:
+        args.output_dir = f'{args.model}-{_input_size}p{_patch_size}p{args.sr}p{args.text_encoder}'
+        args.output_dir += f'-{common.get_timestamp()[:6]}-{common.arg_conf_str(args, defaults=arg_conf_defaults)}'
+
+    if not args.log_dir:
+        args.log_dir = args.output_dir
+    args.target_layers = None if args.target_layers == '' else eval(args.target_layers)
+    return args
+
+
+if __name__ == '__main__':
+    args = complete_args()
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
     main(args)
--- _org/clap/clap_only.py	2025-09-22 23:39:11.716392388 +0900
+++ clap/clap_only.py	2026-02-21 06:38:05.439452653 +0900
@@ -1,93 +1,235 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
+"""M2D-CLAP Pre-training Script for Stages 2 and 2.1
+
+- M2D-CLAP: Exploring General-purpose Audio-Language Representations Beyond CLAP
+https://ieeexplore.ieee.org/document/11168481
+
+- M2D-CLAP: Masked Modeling Duo Meets CLAP for Learning General-purpose Audio-Language Representation
+https://www.isca-archive.org/interspeech_2024/niizumi24_interspeech.html
+"""
 
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# --------------------------------------------------------
-# References:
-# DeiT: https://github.com/facebookresearch/deit
-# BEiT: https://github.com/microsoft/unilm/tree/master/beit
-# --------------------------------------------------------
 import argparse
 import datetime
 import json
 import numpy as np
 import os
 import time
+import math
 from pathlib import Path
+import subprocess
+import sys
+import random
+from typing import Iterable
 
 import torch
 import torch.backends.cudnn as cudnn
 from torch.utils.tensorboard import SummaryWriter
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
-
-import timm
+from functools import partial
+import matplotlib.pyplot as plt
 
-assert timm.__version__ == "0.3.2"  # version check
+import timm.optim
 import timm.optim.optim_factory as optim_factory
 
 import util.misc as misc
 from util.misc import NativeScalerWithGradNormCount as NativeScaler
 
-import models_mae
+from m2d import models_mae
 
-from engine_pretrain import train_one_epoch
+from m2d.engine_pretrain_m2d import train_one_epoch_m2dx
+from audio_dataset import SpectrogramDataset, get_files_no_sort, log_mixup_exp, pd
+import common
+from m2d.runtime_audio import RuntimeM2D, add_semantic_audio_proj, parse_clap_type
+
+
+class CaptionSpecDataset(torch.utils.data.Dataset):
+    def __init__(self, base_folder, files_main, files_bg_noise, file_caption, audio_frames, noise_ratio=0.0,
+                 random_crop=True, n_norm_calc=10000) -> None:
+        super().__init__()
+
+        self.ds1 = SpectrogramDataset(folder=base_folder, files=files_main, crop_frames=audio_frames,
+                random_crop=random_crop, norm_stats=[0.0, 1.0], n_norm_calc=0)
+        print('ds1 norm_stats:', self.ds1.norm_stats)
+
+        if noise_ratio > 0.0:
+            self.ds2 = SpectrogramDataset(folder=base_folder, files=files_bg_noise, crop_frames=audio_frames,
+                    random_crop=random_crop, norm_stats=[0.0, 1.0], n_norm_calc=0, repeat_short=True)
+            print('ds2 norm_stats:', self.ds2.norm_stats)
+        # for BG noise
+        self.noise_ratio = noise_ratio
+        self.bg_index = []
+        # load captions
+        self.caps = []
+        for capfile in file_caption.split(','):
+            caps = pd.read_csv(capfile).set_index('ytid')
+            self.caps.append(caps)
+            print('Caption', Path(capfile).stem, '-> captions:', len(caps), ' expample keys:', list(caps.index)[:5])
+        self.file_caption = file_caption
+        print('Using', len(self.caps), 'caption datasets.')
+
+    def __len__(self):
+        return len(self.ds1)
+
+    def get_random_caption(self, ytid, folder_stem):
+        cap_list = []
+        # by the YouTube ID
+        for cap in self.caps:
+            if ytid in cap.index:
+                cur = cap.loc[ytid].to_list()
+                cap_list.extend(cur)
+        # by the stem of the filename
+        if len(cap_list) == 0:
+            for cap in self.caps:
+                if folder_stem in cap.index:
+                    cur = cap.loc[folder_stem].to_list()
+                    cap_list.extend(cur)
+        if len(cap_list) == 0:
+            print('MISSING', self.caps, ytid, folder_stem)
+        return random.choice(cap_list)
+
+    def __getitem__(self, index, fixed_noise=False):
+        # load index sample
+        clean = self.ds1[index]
+        if self.noise_ratio > 0.0:
+            # load random noise sample ### , while making noise floor zero
+            noise = self.ds2[index if fixed_noise else self.get_next_bgidx()]
+            # mix
+            mixed = log_mixup_exp(noise, clean, self.noise_ratio) if self.noise_ratio < 1.0 else noise
+        else:
+            mixed = clean
+
+        # load sample's caption
+        filepath = self.ds1.df.file_name.values[index].split('/')
+        ytid = filepath[-1][:11]
+        stem = Path(filepath[-1]).stem
+        folder = {'AudioSet_SL_flac': 'AudioSet_SL', 'BBC_Sound_Effects_flac': 'BBC_Sound_Effects',
+            'FreeSound_flac': 'FreeSound', 'SoundBible_flac': 'SoundBible',
+        }[filepath[-2]] if filepath[0] in ['wavcaps_lms', 'wavcaps_32klms'] else ''  # folder name conversion for wavcaps_lms
+        if filepath[0] in ['clotho_lms', 'clotho_32klms']:
+            stem = stem.replace(':', '_').replace('?', '_')
+        caption = self.get_random_caption(ytid, folder + stem)
+
+        return mixed, caption
+
+    def get_next_bgidx(self):
+        if len(self.bg_index) == 0:
+            self.bg_index = torch.randperm(len(self.ds2)).tolist()
+            # print(f'Refreshed the bg index list with {len(self.bg_index)} items: {self.bg_index[:5]}...')
+        return self.bg_index.pop(0)
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + f'(crop_frames={self.ds1.crop_frames}, '
+        format_string += f'folder_main={self.ds1.df.file_name.values[0].split("/")[0]}, '
+        if self.noise_ratio > 0.: format_string += f'folder_bg={self.ds2.df.file_name.values[0].split("/")[0]}, '
+        format_string += f'caption={self.file_caption})'
+        return format_string
+
+
+def build_captioned_dataset(cfg):
+    # create dataset of (spectrogram, caption) pairs
+    files_main = get_files_no_sort(cfg.csv_main)
+    files_bg = get_files_no_sort(cfg.csv_bg_noise) if cfg.noise_ratio > 0. else []
+    ds = CaptionSpecDataset(
+        base_folder=cfg.data_path, files_main=files_main,
+        files_bg_noise=files_bg,
+        file_caption=cfg.file_caption,
+        audio_frames=cfg.audio_frames,
+        noise_ratio=cfg.noise_ratio,
+        random_crop=True)
+
+    return ds
 
 
 def get_args_parser():
-    parser = argparse.ArgumentParser('MAE pre-training', 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=400, type=int)
-    parser.add_argument('--accum_iter', default=1, type=int,
-                        help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+    """
+    Usage: torchrun --nproc_per_node=4 -m clap.clap_only (arguments)
+
+    Explicitly specified arguments for typical stage 2 pre-training:
+    --epochs 30
+    --eval_after 30
+    --save_freq 30
+    --data_path /path/to/my_data_lms
+    --csv_main data/files_A_S_V_S_W_C_A_C_X_t_E_U.csv
+    --mask_ratio 0.0  # 0.0 for stage 2.1, 0.3 for stage 2
+    --base_model m2d_clap_vit_base-80x1001p16x16p16kpN-2025-stage1.1/weights_ep66it3124-0.49058_loss0.0147.pth
+    --text_encoder B
+    --finetune m2d_clap_vit_base-80x608p16x16p16kpBpTI-2025-stage2/checkpoint-30.pth  # for stage 2.1 only
+    --seed 42
+
+    Other implicit but important arguments:
+    --audio_frames, default=1001
+    --sem_mode 1
+    """
+    parser = argparse.ArgumentParser('CLAP only pre-training for M2D-CLAP', add_help=False)
+    parser.add_argument('--batch_size', default=512, type=int,
+                        help='Batch size per GPU (effective batch size is batch_size * # gpus')
+    parser.add_argument('--epochs', default=100, type=int)
+    parser.add_argument('--eval_after', default=10, type=int)
+    parser.add_argument('--save_freq', default=10, type=int)
+    parser.add_argument('--feature_eval_freq', default=10, type=int, help='Feature label-free evaluation frequency.')
+    parser.add_argument('--stop_at', default=-1, type=int)
 
     # Model parameters
-    parser.add_argument('--model', default='mae_vit_large_patch16', type=str, metavar='MODEL',
-                        help='Name of model to train')
+    parser.add_argument('--base_model', default='m2d_clap_vit_base-80x608p16x16-240128/checkpoint-300.pth',
+        type=str, metavar='MODEL', help='Base M2D-CLAP model.')
 
-    parser.add_argument('--input_size', default=224, type=int,
-                        help='images input size')
+    parser.add_argument('--audio_frames', default=1001, type=int, help='Audio duration: 0=the input dur of the model, 1001=10s, etc.')
 
-    parser.add_argument('--mask_ratio', default=0.75, type=float,
-                        help='Masking ratio (percentage of removed patches).')
+    parser.add_argument('--patch_size', default='16x16', type=str, help='patch size')
+    parser.add_argument('--sr', default='16k', type=str, help='Sampling rate of the input audio.')  # just for adding to the output folder name
+    parser.add_argument('--text_encoder', default='', type=str, metavar='TE', help='Text encoder type: A=GTEbase, B=BertBase, N=NVEmbedV2')
+    parser.add_argument('--sem_mode', default=1, type=int, help='<Re-audio-projector> Replace audio projector with CLAP semantic mode (audio projector blocks).')
+    parser.add_argument('--train_audio', action='store_true', help='Train audio encoder.')   
 
-    parser.add_argument('--norm_pix_loss', action='store_true',
-                        help='Use (per-patch) normalized pixels as targets for computing loss')
-    parser.set_defaults(norm_pix_loss=False)
+    parser.add_argument('--mask_ratio', default=0.3, type=float,
+                        help='Masking ratio (percentage of removed patches).')
+    parser.add_argument('--random_mask', action='store_true')
+    parser.add_argument('--target_layers', default='', type=str,
+                        help='Experimental: layers to calculate target representations.')
+    parser.add_argument('--bf16', action='store_true')
 
     # Optimizer parameters
     parser.add_argument('--weight_decay', type=float, default=0.05,
                         help='weight decay (default: 0.05)')
+    parser.add_argument('--clip_grad', type=float, default=3.0, metavar="NORM",
+                        help="Clip gradient norm (default: None, no clipping)")
+    parser.add_argument('--optim', default='adamw', type=str, help='Optimizer adam or sdg')
 
     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',
+    parser.add_argument('--blr', type=float, default=3e-6, metavar='LR',
                         help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
     parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                         help='lower lr bound for cyclic schedulers that hit 0')
 
-    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
+    parser.add_argument('--warmup_epochs', type=int, default=5, metavar='N',
                         help='epochs to warmup LR')
 
     # Dataset parameters
-    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
+    parser.add_argument('--data_path', default='data', type=str,
                         help='dataset path')
+    parser.add_argument('--csv_main', default='data/files_A_S_V_S_W_C_A_C_X_t_E_U.csv', type=str, help='A CSV file to list sample files in the main dataset')
+
+    parser.add_argument('--csv_bg_noise', default='', type=str,
+                        help='A CSV file to list sample files in the BG noise dataset')
+    parser.add_argument('--file_caption', default='data/rawcap_auto_acd.csv,data/rawcap_auto_acd_vggsound.csv,data/rawcap_wav_caps.csv,data/rawcap_audio_caps.csv,data/rawcap_clotho.csv',
+                        type=str, help='CSV files for captions')
+    parser.add_argument('--noise_ratio', default=0., type=float,
+                        help='Noise mixing ratio')
 
-    parser.add_argument('--output_dir', default='./output_dir',
+    parser.add_argument('--output_dir', default='',
                         help='path where to save, empty for no saving')
-    parser.add_argument('--log_dir', default='./output_dir',
+    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('--finetune', default='', help='Fine-tune checkpoint')
+    parser.add_argument('--load_textenc', default='', help='Load text encoder weight')
 
     parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                         help='start epoch')
-    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--force_start_epoch', default=0, type=int, metavar='N',  # 0=always reset start epoch to 0 even if using --resume
+                        help='start epoch for resuming')
+    parser.add_argument('--num_workers', default=8, 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')
@@ -96,7 +238,7 @@
     # 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('--local-rank' if torch.__version__ >= "2.0.0" else '--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')
@@ -104,11 +246,105 @@
     return parser
 
 
-def main(args):
-    misc.init_distributed_mode(args)
+def get_optim(args, param_groups):
+    if args.optim == 'adamw':
+        return torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    elif args.optim == 'lars':
+        return timm.optim.Lars(param_groups, lr=args.lr, momentum=0.9, weight_decay=0)
+    elif args.optim == 'sgd':
+        return torch.optim.SGD(param_groups, args.lr, momentum=0.9, weight_decay=0)
+    assert False, f'Unsupported optimizer {args.optim}'
+
+
+def set_layers_trainable(layer, trainable=False):
+    for n, p in layer.named_parameters():
+        p.requires_grad = trainable
+
+
+def show_layers_trainable(layer, name='', show_all_trainable=True):
+    total_params = sum(p.numel() for p in layer.parameters())
+    total_trainable_params = sum(p.numel() for p in layer.parameters() if p.requires_grad)
+    print(f'{name}Total number of parameters: {total_params:,} (trainable {total_trainable_params:,})')
+    trainable = [n for n, p in layer.named_parameters() if p.requires_grad]
+    frozen = [n for n, p in layer.named_parameters() if not p.requires_grad]
+    print('Trainable parameters:', trainable if show_all_trainable else f'{trainable[:10]} ...')
+    print('Others are frozen such as:', frozen[:3], '...' if len(frozen) >= 3 else '')
+
+
+import util.lr_sched as lr_sched
+
+def train_one_epoch_clap(model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, loss_fn,
+                    log_writer=None,
+                    do_analysis: bool=False,
+                    autocast_args: dict={},
+                    args=None):
+    """M2D-CLAP training loop."""
+    _model = model if hasattr(model, 'logit_scale') else model.module
+    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
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    for data_iter_step, samples in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+        lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        x_online, x_target = samples[0].to(device, non_blocking=True), samples[1]
+
+        with torch.cuda.amp.autocast(**autocast_args):
+            audio_embs = _model.encode_clap_audio(x_online, is_lms=True)
+            text_embs  = _model.encode_clap_text(x_target).to(device)
+            loss = loss_fn(image_features=audio_embs, text_features=text_embs, logit_scale=_model.logit_scale.exp())
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss_scaler(loss, optimizer, parameters=model.parameters(), clip_grad=args.clip_grad, update_grad=True)
+        optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        # Clip logit scale if a CLAP model own it
+        # Thanks to https://github.com/openai/CLIP/issues/46#issuecomment-793434211
+        # logit scaling set as max 100
+        _model.logit_scale.data = torch.clamp(_model.logit_scale.data, 0, 4.6052)  # 4.6052 == log(100)
+
+        metric_logger.update(loss=loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None:
+            """ 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('train_loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+
+    # for CLIP model: logit_scale
+    _logit_scale = getattr(_model, "logit_scale", None)
+    if _logit_scale is not None:
+        metric_logger.update(logit_scale=float(_logit_scale.detach().cpu()))
+
+    # 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()}
+
 
+def main(args):
     print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
-    print("{}".format(args).replace(', ', ',\n'))
 
     device = torch.device(args.device)
 
@@ -119,31 +355,53 @@
 
     cudnn.benchmark = True
 
-    # simple augmentation
-    transform_train = transforms.Compose([
-            transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 3 is bicubic
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
-    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
-    print(dataset_train)
-
-    if True:  # args.distributed:
-        num_tasks = misc.get_world_size()
-        global_rank = misc.get_rank()
-        sampler_train = torch.utils.data.DistributedSampler(
-            dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
-        )
-        print("Sampler_train = %s" % str(sampler_train))
-    else:
-        sampler_train = torch.utils.data.RandomSampler(dataset_train)
+    # basic matters
+    num_tasks = misc.get_world_size()
+    global_rank = misc.get_rank()
 
     if global_rank == 0 and args.log_dir is not None:
         os.makedirs(args.log_dir, exist_ok=True)
         log_writer = SummaryWriter(log_dir=args.log_dir)
     else:
         log_writer = None
+    common.PrintLogger(f'{args.log_dir}/console.txt')
+
+    # build model first
+    model = RuntimeM2D(weight_file=args.base_model, training_mask=args.mask_ratio, flat_features=True, random_mask=args.random_mask)
+    model.get_clap_text_encoder()
+    model.logit_scale = torch.nn.Parameter(torch.ones([]) * 4.6052)  # 4.6052 == log(100)
+    # replace audio projector
+    if args.sem_mode > 0:
+        print(' ** using semantic audio projector with', args.sem_mode, 'layers **')
+        model.backbone.audio_proj = add_semantic_audio_proj(args.sem_mode, 768)
+    # replace text encoder
+    if re_text_encoder(model, Path(args.base_model).parent, Path(args.output_dir).stem):
+        print(' ** replaced text encoder with', model.text_encoder)
+
+    # set audio spec
+    args.input_size = model.cfg.input_size
+    if args.audio_frames > 0:
+        print(f' using {args.audio_frames} audio frames from each data samples')
+    args.audio_frames = args.audio_frames if args.audio_frames > 0 else model.cfg.input_size[1]
 
+    # build dataset
+    dataset_train = build_captioned_dataset(args)
+    print(dataset_train)
+
+    org_args_lr = args.lr
+    if args.lr is None:  # only base_lr is specified
+        args.lr = args.blr * args.eff_batch_size / 256
+
+    print("base lr: %.2e" % (args.lr * 256 / args.eff_batch_size) if org_args_lr is None else 'base lr: not effective')
+    print("actual lr: %.2e" % args.lr)
+
+    print("effective batch size: %d" % args.eff_batch_size)
+
+    # build data loader
+    sampler_train = torch.utils.data.DistributedSampler(
+        dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+    )
+    print("Sampler_train = %s" % str(sampler_train))
     data_loader_train = torch.utils.data.DataLoader(
         dataset_train, sampler=sampler_train,
         batch_size=args.batch_size,
@@ -151,53 +409,79 @@
         pin_memory=args.pin_mem,
         drop_last=True,
     )
-    
-    # define the model
-    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
 
+    # prepare the model
     model.to(device)
+    print(model.cfg)
 
     model_without_ddp = model
     print("Model = %s" % str(model_without_ddp))
 
-    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
-    
-    if args.lr is None:  # only base_lr is specified
-        args.lr = args.blr * eff_batch_size / 256
-
-    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
-    print("actual lr: %.2e" % args.lr)
-
-    print("accumulate grad iterations: %d" % args.accum_iter)
-    print("effective batch size: %d" % eff_batch_size)
+    # freeze audio encoder
+    set_layers_trainable(model.backbone, trainable=args.train_audio)
+    set_layers_trainable(model.backbone.audio_proj, trainable=True)
+    show_layers_trainable(model.backbone, name='<<Audio encoder>> ', show_all_trainable=False)
+    show_layers_trainable(model.backbone.audio_proj, name='<<audio proj>> ', show_all_trainable=False)
+    show_layers_trainable(model.text_encoder, name='<<Text encoder>> ', show_all_trainable=False)
 
+    # set model as distributed training
     if args.distributed:
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
         model_without_ddp = model.module
-    
+
     # following timm: set wd as 0 for bias and norm layers
-    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
-    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    try:
+        param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    except:
+        print(' (for compatibility with timm) Switched add_weight_decay() to param_groups_weight_decay()')
+        param_groups = timm.optim.param_groups_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = get_optim(args, param_groups)
     print(optimizer)
     loss_scaler = NativeScaler()
 
-    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
+    if args.finetune:
+        checkpoint = torch.load(args.finetune, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
+        print(f'Fine-tunning the checkpoint {args.finetune}')
+
+    if args.load_textenc:
+        checkpoint = torch.load(args.load_textenc, map_location='cpu')
+        model_without_ddp.text_encoder.load_state_dict(checkpoint['model'], strict=True)
+        print(f'Loaded text encoder weight {args.load_textenc}')
+
+    if args.force_start_epoch >= 0:
+        args.start_epoch = args.force_start_epoch
 
+    print("{}".format(args).replace(', ', ',\n'))
     print(f"Start training for {args.epochs} epochs")
     start_time = time.time()
+    last_subprocess = None
+
+    # CLIP loss function
+    loss_fn = models_mae.ClipLoss(rank=global_rank, world_size=num_tasks)
+
     for epoch in range(args.start_epoch, args.epochs):
+        epoch1 = epoch + 1
         if args.distributed:
             data_loader_train.sampler.set_epoch(epoch)
-        train_stats = train_one_epoch(
+        train_stats = train_one_epoch_clap(
             model, data_loader_train,
-            optimizer, device, epoch, loss_scaler,
+            optimizer, device, epoch, loss_scaler, loss_fn,
             log_writer=log_writer,
+            do_analysis=(epoch1 % args.feature_eval_freq == 0),
+            autocast_args=dict(dtype=torch.bfloat16) if args.bf16 else {},
             args=args
         )
-        if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
+
+        if args.output_dir and (epoch1 % args.save_freq == 0 or epoch1 == args.epochs):
             misc.save_model(
                 args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                loss_scaler=loss_scaler, epoch=epoch)
+                loss_scaler=loss_scaler, epoch=epoch1)
+            # run the external evaluator
+            if args.eval_after <= epoch1 and epoch1 < args.epochs and misc.is_main_process():
+                abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+                print('quick_eval', abspath)
+                last_subprocess = subprocess.Popen(['/bin/bash', './quick_eval.sh', abspath])
 
         log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                         'epoch': epoch,}
@@ -207,15 +491,82 @@
                 log_writer.flush()
             with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                 f.write(json.dumps(log_stats) + "\n")
+        
+        if args.stop_at > 0 and epoch1 >= args.stop_at:
+            if last_subprocess is not None:
+                last_subprocess.wait()
+            print(f'Stop training by reaching args.stop_at epoch: {args.stop_at}')
+            exit(0)
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
     print('Training time {}'.format(total_time_str))
 
+    del model_without_ddp, model, data_loader_train, optimizer, loss_scaler
+    if misc.is_main_process():
+        abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+        subprocess.call(['/bin/bash', 'clap/clap_eval.sh', abspath])
+    return epoch1
+
+
+arg_conf_defaults = {
+    'csv_main': ('data/files_A_S_V_S_W_C_A_C_X_t_E_U.csv', 'M', 'path'),
+    'mask_ratio': (0.3, 'mr', 'z'),
+    'seed': (0, 's', 'asis'),
+    'optim': ('adamw', 'O', 'asis'),
+    'blr': (3e-6, 'blr', 'z'),
+    'eff_batch_size': (2048, 'bs', 'asis'),
+    'load_textenc': ('', 'TE', 'path'),
+    'audio_frames': (1001, 'Dur', 'asis'),
+    'random_mask': (False, 'R', 'b'),
+    'train_audio': (False, 'trnA', 'b'),
+    'sem_mode': (1, 'S', 'asis'),
+    'finetune': ('', 'FT', 'b'),
+}
+
+
+def replace_text_encoder_spec(spec, new_text_encoder):
+    if new_text_encoder != '':
+        specs = spec.split('p')
+        cur_text_encoder = specs[3]
+        if new_text_encoder != cur_text_encoder:
+            return 'p'.join(specs[:3] + [new_text_encoder] + specs[4:])
+    return spec
+
+
+def re_text_encoder(model, cur_name, new_name):
+    cur_te = parse_clap_type(cur_name)
+    new_te = parse_clap_type(new_name)
+    if cur_te == new_te:
+        return False
+    model.get_clap_text_encoder(weight_name=new_name, remove_text_proj=True)
+    return True
 
-if __name__ == '__main__':
+
+def complete_args():
     args = get_args_parser()
     args = args.parse_args()
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+    misc.init_distributed_mode(args)
+
+    args.eff_batch_size = args.batch_size * misc.get_world_size()
+
+    # m2d_clap_vit_base-80x608p16x16 -> m2d_clap_vit_base-80x608p16x16p16kpApTI
+    parent = Path(args.base_model).parent.name.split('-')
+    # For backward compatibility
+    parent[1] = parent[1] if len(parent[1].split('p')) > 3 else (
+        parent[1] + 'pA' if len(parent[1].split('p')) == 3 else parent[1] + 'p16kpA')
+    # Replace text encoder model if specified
+    parent[1] = replace_text_encoder_spec(parent[1], args.text_encoder)
+    args.output_dir = f'{parent[0]}-{parent[1]}pTI-{"-".join(parent[2:])}'
+    args.output_dir += f'-{common.arg_conf_str(args, defaults=arg_conf_defaults)}-{common.get_timestamp()[:6]}'
+
+    if not args.log_dir:
+        args.log_dir = args.output_dir
+    return args
+
+
+if __name__ == '__main__':
+    args = complete_args()
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
     main(args)
--- _org/speech/train_speech.py	2024-03-25 20:02:30.881776628 +0900
+++ speech/train_speech.py	2024-04-16 10:51:29.818134693 +0900
@@ -1,3 +1,9 @@
+"""M2D-S Pre-training Script
+
+Masked Modeling Duo for Speech: Specializing General-Purpose Audio Representation to Speech using Denoising Distillation
+https://arxiv.org/abs/2305.14079
+"""
+
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 
@@ -15,69 +21,105 @@
 import os
 import time
 from pathlib import Path
+import subprocess
 
 import torch
 import torch.backends.cudnn as cudnn
 from torch.utils.tensorboard import SummaryWriter
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
-
-import timm
+from functools import partial
+import matplotlib.pyplot as plt
 
-assert timm.__version__ == "0.3.2"  # version check
 import timm.optim.optim_factory as optim_factory
 
 import util.misc as misc
 from util.misc import NativeScalerWithGradNormCount as NativeScaler
 
-import models_mae
+from m2d import models_mae
 
-from engine_pretrain import train_one_epoch
+from m2d.engine_pretrain_m2d import train_one_epoch
+import speech.speech_dataset as audio_speech_dataset
+import common
 
 
 def get_args_parser():
-    parser = argparse.ArgumentParser('MAE pre-training', add_help=False)
-    parser.add_argument('--batch_size', default=64, type=int,
+    parser = argparse.ArgumentParser('Masked Modeling Duo (M2D) pre-training', add_help=False)
+    parser.add_argument('--batch_size', default=512, type=int,
                         help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
-    parser.add_argument('--epochs', default=400, type=int)
+    parser.add_argument('--epochs', default=1000, type=int)
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+    parser.add_argument('--eval_after', default=500, type=int)
+    parser.add_argument('--save_freq', default=500, type=int)
+    parser.add_argument('--eval_freq', default=10, type=int, help='Feature label-free evaluation frequency.')
+    parser.add_argument('--stop_at', default=-1, type=int)
 
     # Model parameters
-    parser.add_argument('--model', default='mae_vit_large_patch16', type=str, metavar='MODEL',
+    parser.add_argument('--model', default='m2d_s_vit_base', type=str, metavar='MODEL',
                         help='Name of model to train')
+    parser.add_argument('--decoder_depth', type=int, default=8, metavar='DD',
+                        help='model decoder depth')
 
-    parser.add_argument('--input_size', default=224, type=int,
-                        help='images input size')
+    parser.add_argument('--input_size', default='80x208', type=str, help='images input size')
+    parser.add_argument('--patch_size', default='80x4', type=str, help='patch size')
 
-    parser.add_argument('--mask_ratio', default=0.75, type=float,
+    parser.add_argument('--mask_ratio', default=0.6, type=float,
                         help='Masking ratio (percentage of removed patches).')
+    parser.add_argument('--ema_decay_init', default=0.99995, type=float,
+                        help='Initial EMA decay parameter.')
+    parser.add_argument('--ema_decay', default=0.99999, type=float,
+                        help='EMA decay parameter.')
+    parser.add_argument('--loss_fn', default='norm_mse', type=str,
+                        help='loss function: mse or norm_mse.')
+    parser.add_argument('--loss_m2d', default=1., type=float,
+                        help='Loss of M2D masked prediction')
+    parser.add_argument('--loss_off', default=0., type=float,
+                        help='Loss of offline target')
+    parser.add_argument('--target_layers', default='', type=str,
+                        help='Experimental: layers to calculate target representations.')
 
-    parser.add_argument('--norm_pix_loss', action='store_true',
+    parser.add_argument('--no_norm_pix_loss', action='store_false', dest='norm_pix_loss',
                         help='Use (per-patch) normalized pixels as targets for computing loss')
-    parser.set_defaults(norm_pix_loss=False)
+    parser.set_defaults(norm_pix_loss=True)
+
+    parser.add_argument('--cont_mask', default=0, type=int,
+                        help='Use random 1-d (continuous) masking scheme. 0:off, 0<:# of frames to make them continuous.')
 
     # Optimizer parameters
     parser.add_argument('--weight_decay', type=float, default=0.05,
                         help='weight decay (default: 0.05)')
+    parser.add_argument('--clip_grad', type=float, default=3.0, metavar="NORM",
+                        help="Clip gradient norm (default: None, no clipping)")
+    parser.add_argument('--optim', default='adamw', type=str, help='Optimizer adam or sdg')
 
     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',
+    parser.add_argument('--blr', type=float, default=3e-4, metavar='LR',
                         help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
     parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                         help='lower lr bound for cyclic schedulers that hit 0')
 
-    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
+    parser.add_argument('--warmup_epochs', type=int, default=60, metavar='N',
                         help='epochs to warmup LR')
 
     # Dataset parameters
-    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
+    parser.add_argument('--data_path', default='data', type=str,
                         help='dataset path')
+    parser.add_argument('--csv_main', default='data/files_ls960_hybrid.csv', type=str,
+                        help='A CSV file to list sample files in the main dataset')
+    parser.add_argument('--csv_bg_noise', default='data/files_audioset.csv', type=str,
+                        help='A CSV file to list sample files in the BG noise dataset')
+    parser.add_argument('--csv_val', default='', type=str,
+                        help='A CSV file to list validation sample files')
+    parser.add_argument('--norm_stats', default='None', type=str,  # Will be computed runtime.
+                        help='dataset normalization stats')
+    parser.add_argument('--noise_ratio', default=0.2, type=float,
+                        help='Noise mixing ratio')
+    parser.add_argument('--swap_inputs', action='store_true', default=False,
+                        help='Swap input samples between clean and mixed for feeding clean samples to M2D and mixed samples to the teacher.')
 
-    parser.add_argument('--output_dir', default='./output_dir',
+    parser.add_argument('--output_dir', default='',
                         help='path where to save, empty for no saving')
-    parser.add_argument('--log_dir', default='./output_dir',
+    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')
@@ -87,7 +129,9 @@
 
     parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                         help='start epoch')
-    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--force_start_epoch', default=-1, type=int, metavar='N',
+                        help='start epoch for resuming')
+    parser.add_argument('--num_workers', default=8, 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')
@@ -104,9 +148,46 @@
     return parser
 
 
-def main(args):
-    misc.init_distributed_mode(args)
+def ema_decay_sched(step, total_steps, ema_decay_init, ema_decay):
+    interp = step / (total_steps - 1)
+    tau = ema_decay_init + (ema_decay - ema_decay_init) * interp
+    return tau
+
+
+def get_optim(args, param_groups):
+    if args.optim == 'adamw':
+        return torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    elif args.optim == 'sgd':
+        return torch.optim.SGD(param_groups, args.lr, momentum=0.9, weight_decay=0)
+    assert False, f'Unsupported optimizer {args.optim}'
+
+
+def visualize_reconstruction(args, device, model, save_path):
+    ds, files = audio_speech_dataset.build_viz_dataset(args)
+    if (ds is None) or (len(ds) == 0):
+        print(f'--Skipped visualization (No samples in {args.data_path}/vis_samples folder.)')
+        return
+    batch = torch.stack([ds[i][0] for i in range(len(ds))])
+    model.eval()
+    with torch.no_grad():
+        recons, _, masks = model.forward_viz(batch.to(device))
+    save_path.mkdir(parents=True, exist_ok=True)
+    viz_imgs = [batch[:, 0], masks] if recons is None else [batch[:, 0], recons[:, 0], masks]
+
+    for i, file in enumerate(files):
+        # as .npy
+        if recons is not None: np.save(f"{save_path}/recon_{Path(file).name}", recons[i].cpu().numpy())
+        # as .png
+        cur_imgs = [img[i] for img in viz_imgs]
+        fig = plt.figure(figsize=[12, 8 if batch[0].shape[-1] < 310 else 6])
+        for j, img in enumerate(cur_imgs):
+            ax = fig.add_subplot(3, 1, j + 1)
+            ax.imshow(img.cpu().numpy(), origin='lower')
+        plt.margins(x=0, y=0)
+        fig.savefig(f'{save_path}/recon_{Path(file).stem}.png', bbox_inches = 'tight')
 
+
+def main(args):
     print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
     print("{}".format(args).replace(', ', ',\n'))
 
@@ -119,13 +200,7 @@
 
     cudnn.benchmark = True
 
-    # simple augmentation
-    transform_train = transforms.Compose([
-            transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 3 is bicubic
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
-    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
+    dataset_train, dataset_val = audio_speech_dataset.build_mixed_speech_dataset(args)
     print(dataset_train)
 
     if True:  # args.distributed:
@@ -143,6 +218,8 @@
         log_writer = SummaryWriter(log_dir=args.log_dir)
     else:
         log_writer = None
+    common.PrintLogger(f'{args.log_dir}/console.txt')
+    print(args)
 
     data_loader_train = torch.utils.data.DataLoader(
         dataset_train, sampler=sampler_train,
@@ -151,53 +228,86 @@
         pin_memory=args.pin_mem,
         drop_last=True,
     )
-    
+
+    # for validation loss
+    if args.csv_val != '':
+        sampler_val = torch.utils.data.DistributedSampler(dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+        print("Sampler_val = %s" % str(sampler_val))
+        data_loader_val = torch.utils.data.DataLoader(dataset_val, sampler=sampler_val, batch_size=args.batch_size,
+            num_workers=args.num_workers//2, pin_memory=args.pin_mem, drop_last=False)
+    else:
+        data_loader_val = None
+
     # define the model
-    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
+    model = models_mae.__dict__[args.model](img_size=args.input_size, patch_size=args.patch_size, decoder_depth=args.decoder_depth,
+        norm_pix_loss=args.norm_pix_loss, loss_type=args.loss_fn, target_layers=args.target_layers, loss_m2d=args.loss_m2d, loss_off=args.loss_off)
+
+    if args.cont_mask > 0:
+        model.set_random_1d_mask(args.cont_mask)
 
     model.to(device)
 
     model_without_ddp = model
     print("Model = %s" % str(model_without_ddp))
 
-    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
-    
+    org_args_lr = args.lr
     if args.lr is None:  # only base_lr is specified
-        args.lr = args.blr * eff_batch_size / 256
+        args.lr = args.blr * args.eff_batch_size / 256
 
-    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
+    print("base lr: %.2e" % (args.lr * 256 / args.eff_batch_size) if org_args_lr is None else 'base lr: not effective')
     print("actual lr: %.2e" % args.lr)
 
     print("accumulate grad iterations: %d" % args.accum_iter)
-    print("effective batch size: %d" % eff_batch_size)
+    print("effective batch size: %d" % args.eff_batch_size)
 
     if args.distributed:
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
         model_without_ddp = model.module
-    
+
     # following timm: set wd as 0 for bias and norm layers
-    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
-    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    try:
+        param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    except:
+        print(' (for compatibility with timm) Switched add_weight_decay() to param_groups_weight_decay()')
+        param_groups = optim_factory.param_groups_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = get_optim(args, param_groups)
     print(optimizer)
     loss_scaler = NativeScaler()
 
-    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
+    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, delta_epoch=0)
+
+    if args.force_start_epoch >= 0:
+        args.start_epoch = args.force_start_epoch
 
     print(f"Start training for {args.epochs} epochs")
     start_time = time.time()
+    last_subprocess = None
     for epoch in range(args.start_epoch, args.epochs):
+        epoch1 = epoch + 1
         if args.distributed:
             data_loader_train.sampler.set_epoch(epoch)
         train_stats = train_one_epoch(
             model, data_loader_train,
             optimizer, device, epoch, loss_scaler,
+            partial(ema_decay_sched, total_steps=len(data_loader_train) * args.epochs,
+                ema_decay_init=args.ema_decay_init, ema_decay=args.ema_decay),
+            val_loader=data_loader_val,
             log_writer=log_writer,
+            do_analysis=(epoch1 % args.eval_freq == 0),
             args=args
         )
-        if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
+        if args.output_dir and (epoch1 % args.save_freq == 0 or epoch1 == args.epochs):
             misc.save_model(
                 args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                loss_scaler=loss_scaler, epoch=epoch)
+                loss_scaler=loss_scaler, epoch=epoch1)
+            # visualize reconstructions
+            out_dir = Path(args.output_dir)/str(epoch1)
+            visualize_reconstruction(args, device, model_without_ddp, out_dir)
+            # run the external evaluator
+            if args.eval_after <= epoch1 and epoch1 < args.epochs and misc.is_main_process():
+                abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+                print('quick_eval', abspath)
+                last_subprocess = subprocess.Popen(['/bin/bash', './quick_eval.sh', abspath])
 
         log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                         'epoch': epoch,}
@@ -207,15 +317,64 @@
                 log_writer.flush()
             with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                 f.write(json.dumps(log_stats) + "\n")
+        
+        if args.stop_at > 0 and epoch1 >= args.stop_at:
+            if last_subprocess is not None:
+                last_subprocess.wait()
+            print(f'Stop training by reaching args.stop_at epoch: {args.stop_at}')
+            exit(0)
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
     print('Training time {}'.format(total_time_str))
 
+    del model_without_ddp, model, data_loader_train, optimizer, loss_scaler
+    if misc.is_main_process():
+        abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+        subprocess.call(['/bin/bash', './all_eval.sh', abspath])
+
+
+arg_conf_defaults = {
+    'csv_main': ('', 'M', 'path'),  # data/files_ls960_hybrid.csv
+    'csv_bg_noise': ('', 'D', 'path'),  # data/files_audioset.csv
+    'ema_decay_init': (0.99995, 'ema', 'z'),
+    'ema_decay': (0.99999, 'ed', 'z'),
+    'decoder_depth': (8, 'dd', 'asis'),
+    'mask_ratio': (0.6, 'mr', 'z'),
+    'seed': (0, 's', 'asis'),
+    'norm_pix_loss':  (True, '~N', 'b'),
+    'loss_fn': ('norm_mse', 'L', 'head'),
+    'optim': ('adamw', 'O', 'asis'),
+    'blr': (3e-4, 'blr', 'z'),
+    'lr': (None, 'lr', 'z'),
+    'eff_batch_size': (2048, 'bs', 'asis'),
+    'accum_iter': (1, 'a', 'asis'),
+    'loss_m2d': (1.0, 'lm', 'z'),
+    'loss_off': (0.0, 'lo', 'z'),
+    'noise_ratio': (0.2, 'nr', 'z'),
+    'cont_mask': (0, 'C', 'asis'),
+}
+
 
 if __name__ == '__main__':
     args = get_args_parser()
     args = args.parse_args()
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    _input_size, _patch_size = args.input_size, args.patch_size
+    args.input_size = [int(x) for x in args.input_size.split('x')]
+    args.patch_size = [int(x) for x in args.patch_size.split('x')]
+    args.norm_stats = eval(args.norm_stats) if args.norm_stats else None
+
+    misc.init_distributed_mode(args)
+
+    args.eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
+
+    if not args.output_dir:
+        args.output_dir = f'{args.model}-{_input_size}p{_patch_size}'
+        args.output_dir += f'-{common.get_timestamp()[:6]}-{common.arg_conf_str(args, defaults=arg_conf_defaults)}'
+
+    if not args.log_dir:
+        args.log_dir = args.output_dir
+    args.target_layers = None if args.target_layers == '' else eval(args.target_layers)
+
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
     main(args)
--- _org/audioset/train_as.py	2024-04-10 20:27:28.654944006 +0900
+++ audioset/train_as.py	2024-05-13 07:42:18.383900880 +0900
@@ -1,3 +1,9 @@
+"""Masked Modeling Duo (M2D) Pre-training Script for M2D-AS
+
+Masked Modeling Duo: Towards a Universal Audio Pre-training Framework
+https://arxiv.org/abs/2404.06095
+"""
+
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 
@@ -15,69 +21,187 @@
 import os
 import time
 from pathlib import Path
+import subprocess
+import sys
 
 import torch
 import torch.backends.cudnn as cudnn
 from torch.utils.tensorboard import SummaryWriter
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
+from functools import partial
+import matplotlib.pyplot as plt
 
-import timm
-
-assert timm.__version__ == "0.3.2"  # version check
 import timm.optim.optim_factory as optim_factory
 
-import util.misc as misc
+from util import misc
 from util.misc import NativeScalerWithGradNormCount as NativeScaler
 
-import models_mae
+from m2d import models_mae
+
+from m2d.engine_pretrain_m2d import train_one_epoch_m2dx
+from audio_dataset import SpectrogramDataset, get_files_no_sort, log_mixup_exp, pd
+import common
+from m2d.runtime_audio import RuntimeM2D
+from sklearn.preprocessing import MultiLabelBinarizer
+
+
+class LabelEmbSpecDataset(torch.utils.data.Dataset):
+    def __init__(self, base_folder, files_main, files_bg_noise, labels, crop_size, noise_ratio=0.0,
+                 random_crop=True, n_norm_calc=10000) -> None:
+        super().__init__()
+
+        self.ds1 = SpectrogramDataset(folder=base_folder, files=files_main, crop_frames=crop_size[1],
+                random_crop=random_crop, norm_stats=None, n_norm_calc=n_norm_calc//2)
+        self.norm_stats = self.ds1.norm_stats  # for compatibility with SpectrogramDataset
+        self.norm_std = self.ds1.norm_stats[1]
+        self.ds1.norm_stats = (self.norm_stats[0], 1.0)
+
+        if noise_ratio > 0.0:
+            self.ds2 = SpectrogramDataset(folder=base_folder, files=files_bg_noise, crop_frames=crop_size[1],
+                    random_crop=random_crop, norm_stats=None, n_norm_calc=n_norm_calc//2, repeat_short=True)
+            self.ds2.norm_stats = (self.ds2.norm_stats[0], 1.0) # disable normalizion scaling in the ds2
+
+        # for BG noise
+        self.noise_ratio = noise_ratio
+        self.bg_index = []
+        # label to one-hot
+        oh_enc = MultiLabelBinarizer()
+        self.labels = torch.tensor(oh_enc.fit_transform([str(ls).split(',') for ls in labels]), dtype=torch.float32)
+        self.classes = oh_enc.classes_
+        self.one_hot_dim = self.labels.shape[-1]
+        print('Labels:', len(self.labels), ' classes:', self.one_hot_dim)
+
+    def __len__(self):
+        return len(self.ds1)
+
+    def __getitem__(self, index, fixed_noise=False):
+        # load index sample
+        clean = self.ds1[index]
+        if self.noise_ratio > 0.0:
+            # load random noise sample ### , while making noise floor zero
+            noise = self.ds2[index if fixed_noise else self.get_next_bgidx()]
+            # mix
+            mixed = log_mixup_exp(noise, clean, self.noise_ratio) if self.noise_ratio < 1.0 else noise
+        else:
+            mixed = clean
+        # finalize normalization. clean and noise were averaged to zero. the following will scale to 1.0 using ds1 std.
+        mixed = mixed / self.norm_std
+
+        # load sample's label
+        label = self.labels[index]
+
+        return mixed, label
+
+    def get_next_bgidx(self):
+        if len(self.bg_index) == 0:
+            self.bg_index = torch.randperm(len(self.ds2)).tolist()
+            # print(f'Refreshed the bg index list with {len(self.bg_index)} items: {self.bg_index[:5]}...')
+        return self.bg_index.pop(0)
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + f'(crop_frames={self.ds1.crop_frames}, '
+        format_string += f'folder_main={self.ds1.df.file_name.values[0].split("/")[0]}, '
+        if self.noise_ratio > 0.: format_string += f'folder_bg={self.ds2.df.file_name.values[0].split("/")[0]}, '
+        return format_string
+
+
+def build_as_clf_dataset(cfg):
+    # get files and inflate the number of files (by repeating the list) if needed
+    main_df = pd.read_csv(str(cfg.csv_main))
+    files_main = main_df.file_name.values
+    files_bg = get_files_no_sort(cfg.csv_bg_noise) if cfg.noise_ratio > 0. else []
+
+    ds = LabelEmbSpecDataset(
+        base_folder=cfg.data_path, files_main=files_main,
+        files_bg_noise=files_bg,
+        labels=main_df.label,
+        crop_size=cfg.input_size,
+        noise_ratio=cfg.noise_ratio,
+        random_crop=True)
 
-from engine_pretrain import train_one_epoch
+    val_ds = SpectrogramDataset(folder=cfg.data_path, files=get_files_no_sort(cfg.csv_val), crop_frames=cfg.input_size[1], random_crop=True) \
+        if cfg.csv_val else None
+
+    return ds, val_ds
 
 
 def get_args_parser():
-    parser = argparse.ArgumentParser('MAE pre-training', add_help=False)
-    parser.add_argument('--batch_size', default=64, type=int,
+    parser = argparse.ArgumentParser('Masked Modeling Duo (M2D) pre-training V2', add_help=False)
+    parser.add_argument('--batch_size', default=512, type=int,
                         help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
-    parser.add_argument('--epochs', default=400, type=int)
+    parser.add_argument('--epochs', default=300, type=int)
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+    parser.add_argument('--eval_after', default=50, type=int)
+    parser.add_argument('--save_freq', default=100, type=int)
+    parser.add_argument('--feature_eval_freq', default=10, type=int, help='Feature label-free evaluation frequency.')
+    parser.add_argument('--stop_at', default=-1, type=int)
 
     # Model parameters
-    parser.add_argument('--model', default='mae_vit_large_patch16', type=str, metavar='MODEL',
-                        help='Name of model to train')
+    parser.add_argument('--model', default='m2d_as_vit_base', type=str, metavar='MODEL', help='Name of model to train')
+    parser.add_argument('--decoder_depth', type=int, default=8, metavar='DD', help='Model decoder depth.')
 
-    parser.add_argument('--input_size', default=224, type=int,
-                        help='images input size')
+    parser.add_argument('--input_size', default='80x608', type=str, help='images input size')
+    parser.add_argument('--patch_size', default='16x16', type=str, help='patch size')
+    parser.add_argument('--sr', default='16k', type=str, metavar='SR', help='Sampling rate of the input audio.')
 
-    parser.add_argument('--mask_ratio', default=0.75, type=float,
+    parser.add_argument('--mask_ratio', default=0.7, type=float,
                         help='Masking ratio (percentage of removed patches).')
+    parser.add_argument('--ema_decay_init', default=0.99995, type=float,
+                        help='Initial EMA decay parameter.')
+    parser.add_argument('--ema_decay', default=0.99999, type=float,
+                        help='EMA decay parameter.')
+    parser.add_argument('--loss_fn', default='norm_mse', type=str,
+                        help='loss function: mse or norm_mse.')
+    parser.add_argument('--loss_m2d', default=1., type=float, help='Loss of M2D masked prediction')
+    parser.add_argument('--loss_off', default=.1, type=float, help='Loss of offline target')
+    parser.add_argument('--target_layers', default='', type=str,
+                        help='Experimental: layers to calculate target representations.')
+    parser.add_argument('--bf16', action='store_true')
+    parser.add_argument('--teacher', default='', type=str, help='Weight path of the teacher M2D model.')
 
-    parser.add_argument('--norm_pix_loss', action='store_true',
+    parser.add_argument('--no_norm_pix_loss', action='store_false', dest='norm_pix_loss',
                         help='Use (per-patch) normalized pixels as targets for computing loss')
-    parser.set_defaults(norm_pix_loss=False)
+    parser.set_defaults(norm_pix_loss=True)
+
+    parser.add_argument('--cont_mask', default=0, type=int,
+                        help='Use random 1-d (continuous) masking scheme. 0:off, 0<:# of frames to make them continuous.')
 
     # Optimizer parameters
     parser.add_argument('--weight_decay', type=float, default=0.05,
                         help='weight decay (default: 0.05)')
+    parser.add_argument('--clip_grad', type=float, default=3.0, metavar="NORM",
+                        help="Clip gradient norm (default: None, no clipping)")
+    parser.add_argument('--optim', default='adamw', type=str, help='Optimizer adam or sdg')
 
     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',
+    parser.add_argument('--blr', type=float, default=3e-4, metavar='LR',
                         help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
     parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                         help='lower lr bound for cyclic schedulers that hit 0')
 
-    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
+    parser.add_argument('--warmup_epochs', type=int, default=20, metavar='N',
                         help='epochs to warmup LR')
 
     # Dataset parameters
-    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
+    parser.add_argument('--data_path', default='data', type=str,
                         help='dataset path')
+    parser.add_argument('--csv_main', default='data/files_as_weighted.csv', type=str,
+                        help='A CSV file to list sample files in the main dataset')
+    parser.add_argument('--csv_bg_noise', default='', type=str,
+                        help='A CSV file to list sample files in the BG noise dataset')
+    parser.add_argument('--csv_val', default='', type=str,
+                        help='A CSV file to list validation sample files')
+    parser.add_argument('--min_ds_size', default=10000, type=int,
+                        help='Inflate the size of the smaller dataset to the desired size')
+    parser.add_argument('--norm_stats', default='None', type=str,  # Will be computed runtime.
+                        help='dataset normalization stats')
+    parser.add_argument('--noise_ratio', default=0., type=float,
+                        help='Noise mixing ratio')
 
-    parser.add_argument('--output_dir', default='./output_dir',
+    parser.add_argument('--output_dir', default='',
                         help='path where to save, empty for no saving')
-    parser.add_argument('--log_dir', default='./output_dir',
+    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')
@@ -87,7 +211,9 @@
 
     parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                         help='start epoch')
-    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--force_start_epoch', default=0, type=int, metavar='N',  # 0=always reset start epoch to 0 even if using --resume
+                        help='start epoch for resuming')
+    parser.add_argument('--num_workers', default=8, 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')
@@ -96,7 +222,7 @@
     # 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('--local-rank' if torch.__version__ >= "2.0.0" else '--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')
@@ -104,9 +230,38 @@
     return parser
 
 
-def main(args):
-    misc.init_distributed_mode(args)
+def ema_decay_sched(step, total_steps, ema_decay_init, ema_decay):
+    interp = step / (total_steps - 1)
+    tau = ema_decay_init + (ema_decay - ema_decay_init) * interp
+    return tau
+
+
+def get_optim(args, param_groups):
+    if args.optim == 'adamw':
+        return torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    elif args.optim == 'sgd':
+        return torch.optim.SGD(param_groups, args.lr, momentum=0.9, weight_decay=0)
+    assert False, f'Unsupported optimizer {args.optim}'
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler, delta_epoch=1, strict=True):
+    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'], strict=strict)
+        print("Resume checkpoint %s" % args.resume)
+        if strict == True and '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'] + delta_epoch
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
 
+def main(args):
     print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
     print("{}".format(args).replace(', ', ',\n'))
 
@@ -119,14 +274,8 @@
 
     cudnn.benchmark = True
 
-    # simple augmentation
-    transform_train = transforms.Compose([
-            transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 3 is bicubic
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
-    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
-    print(dataset_train)
+    dataset_train, dataset_val = build_as_clf_dataset(args)
+    print(dataset_train, dataset_val)
 
     if True:  # args.distributed:
         num_tasks = misc.get_world_size()
@@ -143,6 +292,8 @@
         log_writer = SummaryWriter(log_dir=args.log_dir)
     else:
         log_writer = None
+    common.PrintLogger(f'{args.log_dir}/console.txt')
+    print(args)
 
     data_loader_train = torch.utils.data.DataLoader(
         dataset_train, sampler=sampler_train,
@@ -151,53 +302,86 @@
         pin_memory=args.pin_mem,
         drop_last=True,
     )
+
+    # for validation loss
+    if args.csv_val != '':
+        sampler_val = torch.utils.data.DistributedSampler(dataset_val, num_replicas=num_tasks, rank=global_rank, shuffle=False)
+        print("Sampler_val = %s" % str(sampler_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)
+    else:
+        data_loader_val = None
     
     # define the model
-    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
+    model = models_mae.__dict__[args.model](img_size=args.input_size, patch_size=args.patch_size, decoder_depth=args.decoder_depth,
+        norm_pix_loss=args.norm_pix_loss, loss_type=args.loss_fn, target_layers=args.target_layers, loss_m2d=args.loss_m2d, loss_off=args.loss_off,
+        off_emb_dim=dataset_train.one_hot_dim, norm_stats=dataset_train.norm_stats)
+
+    if args.cont_mask > 0:
+        model.set_random_1d_mask(args.cont_mask)
 
     model.to(device)
 
     model_without_ddp = model
     print("Model = %s" % str(model_without_ddp))
 
-    eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
-    
+    org_args_lr = args.lr
     if args.lr is None:  # only base_lr is specified
-        args.lr = args.blr * eff_batch_size / 256
+        args.lr = args.blr * args.eff_batch_size / 256
 
-    print("base lr: %.2e" % (args.lr * 256 / eff_batch_size))
+    print("base lr: %.2e" % (args.lr * 256 / args.eff_batch_size) if org_args_lr is None else 'base lr: not effective')
     print("actual lr: %.2e" % args.lr)
 
     print("accumulate grad iterations: %d" % args.accum_iter)
-    print("effective batch size: %d" % eff_batch_size)
+    print("effective batch size: %d" % args.eff_batch_size)
 
     if args.distributed:
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
         model_without_ddp = model.module
-    
+
     # following timm: set wd as 0 for bias and norm layers
-    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
-    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    try:
+        param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    except:
+        print(' (for compatibility with timm) Switched add_weight_decay() to param_groups_weight_decay()')
+        param_groups = optim_factory.param_groups_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = get_optim(args, param_groups)
     print(optimizer)
     loss_scaler = NativeScaler()
 
-    misc.load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler)
+    load_model(args=args, model_without_ddp=model_without_ddp, optimizer=optimizer, loss_scaler=loss_scaler, delta_epoch=0, strict=False)
+
+    if args.force_start_epoch >= 0:
+        args.start_epoch = args.force_start_epoch
 
     print(f"Start training for {args.epochs} epochs")
     start_time = time.time()
+    last_subprocess = None
     for epoch in range(args.start_epoch, args.epochs):
+        epoch1 = epoch + 1
         if args.distributed:
             data_loader_train.sampler.set_epoch(epoch)
-        train_stats = train_one_epoch(
-            model, data_loader_train,
+        train_stats = train_one_epoch_m2dx(
+            model, None, data_loader_train,
             optimizer, device, epoch, loss_scaler,
+            partial(ema_decay_sched, total_steps=len(data_loader_train) * args.epochs,
+                ema_decay_init=args.ema_decay_init, ema_decay=args.ema_decay),
+            val_loader=data_loader_val,
             log_writer=log_writer,
+            do_analysis=(epoch1 % args.feature_eval_freq == 0),
+            autocast_args=dict(dtype=torch.bfloat16) if args.bf16 else {},
             args=args
         )
-        if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
+
+        if args.output_dir and (epoch1 % args.save_freq == 0 or epoch1 == args.epochs):
             misc.save_model(
                 args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                loss_scaler=loss_scaler, epoch=epoch)
+                loss_scaler=loss_scaler, epoch=epoch1)
+            # run the external evaluator
+            if args.eval_after <= epoch1 and epoch1 < args.epochs and misc.is_main_process():
+                abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+                print('quick_eval', abspath)
+                last_subprocess = subprocess.Popen(['/bin/bash', './quick_eval.sh', abspath])
 
         log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                         'epoch': epoch,}
@@ -207,15 +391,72 @@
                 log_writer.flush()
             with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                 f.write(json.dumps(log_stats) + "\n")
+        
+        if args.stop_at > 0 and epoch1 >= args.stop_at:
+            if last_subprocess is not None:
+                last_subprocess.wait()
+            print(f'Stop training by reaching args.stop_at epoch: {args.stop_at}')
+            exit(0)
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
     print('Training time {}'.format(total_time_str))
 
+    del model_without_ddp, model, data_loader_train, optimizer, loss_scaler
+    if misc.is_main_process():
+        abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+        subprocess.call(['/bin/bash', './all_eval.sh', abspath])
+    return epoch1
+
+
+arg_conf_defaults = {
+    'csv_main': ('data/files_as_weighted.csv', 'M', 'path'),
+    'csv_bg_noise': ('', 'D', 'path'),
+    'ema_decay_init': (0.99995, 'ema', 'z'),
+    'ema_decay': (0.99999, 'ed', 'z'),
+    'decoder_depth': (8, 'dd', 'asis'),
+    'mask_ratio': (0.7, 'mr', 'z'),
+    'seed': (0, 's', 'asis'),
+    'norm_pix_loss':  (True, '~N', 'b'),
+    'loss_fn': ('norm_mse', 'L', 'head'),
+    'optim': ('adamw', 'O', 'asis'),
+    'warmup_epochs': (20, 'wu', 'asis'),
+    'blr': (3e-4, 'blr', 'z'),
+    'lr': (None, 'lr', 'z'),
+    'eff_batch_size': (2048, 'bs', 'asis'),
+    'accum_iter': (1, 'a', 'asis'),
+    'loss_m2d': (1.0, 'lm', 'z'),
+    'loss_off': (1.0, 'lo', 'z'),
+    'noise_ratio': (0.0, 'nr', 'z'),
+    'min_ds_size': (10000, 'dn', 'asis'),
+    'cont_mask': (0, 'C', 'asis'),
+    'epochs': (0, '-e', 'asis'),
+}
 
-if __name__ == '__main__':
+
+def complete_args():
     args = get_args_parser()
     args = args.parse_args()
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    _input_size, _patch_size = args.input_size, args.patch_size
+    args.input_size = [int(x) for x in args.input_size.split('x')]
+    args.patch_size = [int(x) for x in args.patch_size.split('x')]
+    args.norm_stats = eval(args.norm_stats) if args.norm_stats else None
+
+    misc.init_distributed_mode(args)
+
+    args.eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
+
+    if not args.output_dir:
+        args.output_dir = f'{args.model}-{_input_size}p{_patch_size}p{args.sr}'
+        args.output_dir += f'-{common.get_timestamp()[:6]}-{common.arg_conf_str(args, defaults=arg_conf_defaults)}'
+
+    if not args.log_dir:
+        args.log_dir = args.output_dir
+    args.target_layers = None if args.target_layers == '' else eval(args.target_layers)
+    return args
+
+
+if __name__ == '__main__':
+    args = complete_args()
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
     main(args)
--- _org/util/misc.py	2023-07-07 10:22:30.225586716 +0900
+++ util/misc.py	2023-09-22 17:24:28.791494691 +0900
@@ -18,7 +18,10 @@
 
 import torch
 import torch.distributed as dist
-from torch._six import inf
+try:
+    from torch._six import inf
+except:
+    from torch import inf
 
 
 class SmoothedValue(object):
@@ -312,7 +315,7 @@
         model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)
 
 
-def load_model(args, model_without_ddp, optimizer, loss_scaler):
+def load_model(args, model_without_ddp, optimizer, loss_scaler, delta_epoch=1):
     if args.resume:
         if args.resume.startswith('https'):
             checkpoint = torch.hub.load_state_dict_from_url(
@@ -323,7 +326,7 @@
         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
+            args.start_epoch = checkpoint['epoch'] + delta_epoch
             if 'scaler' in checkpoint:
                 loss_scaler.load_state_dict(checkpoint['scaler'])
             print("With optim & sched!")
--- _org/m2d/engine_pretrain_m2d.py	2023-07-07 10:22:30.217586685 +0900
+++ m2d/engine_pretrain_m2d.py	2026-02-21 09:27:52.857061385 +0900
@@ -16,12 +16,156 @@
 
 import util.misc as misc
 import util.lr_sched as lr_sched
+from util.analyze_repr import get_cumulative_explained_variance_auc, rank_me
+from einops import rearrange
+
+
+def evaluate_feature(data_loader_train, data_loader_val, device, model, n_samples=1000):
+    embs = []
+    with torch.no_grad():
+        for batch in data_loader_train:  # sample only or (mixed, clean) or (sample, label)
+            if isinstance(batch, (list, tuple)):
+                if batch[1].shape == batch[0].shape:
+                    batch = batch[1]  # clean from (mixed, clean)
+                else:
+                    batch = batch[0]  # sample from (sample, label)
+            sample = batch.to(device, non_blocking=True)
+            B = sample.shape[0]
+            _model = model if hasattr(model, 'forward_encoder') else model.module
+            x, *_ = _model.forward_encoder(sample, mask_ratio=0.0)
+
+            x = x[:, 1:]   # remove CLS
+            x = rearrange(x, 'b l d -> (b l) d')
+            embs.append(x)
+            if B * len(embs) > n_samples:
+                break # print(f'Stopped loading samples ({B*len(embs)}) > {n_samples}.')
+    embs = torch.vstack(embs)
+    n_actual_features = embs.shape[0]
+    rank = rank_me(embs).cpu()
+    auc = get_cumulative_explained_variance_auc(embs).cpu()
+
+    val_message, val_rank, val_auc, embs = '', -1., -1., []
+    if data_loader_val is not None:
+        with torch.no_grad():
+            for sample in data_loader_val:
+                sample = sample.to(device, non_blocking=True)
+                x, *_ = _model.forward_encoder(sample, mask_ratio=0.0)
+
+                x = x[:, 1:]   # remove CLS
+                x = rearrange(x, 'b l d -> (b l) d')
+                embs.append(x)
+                if B * len(embs) > n_samples:
+                    break # print(f'Stopped loading samples ({B*len(embs)}) > {n_samples}.')
+        embs = torch.vstack(embs)
+        n_val_features = embs.shape[0]
+        val_rank = rank_me(embs).cpu()
+        val_auc = get_cumulative_explained_variance_auc(embs).cpu()
+        val_message = f', val_rank={val_rank:.1f} val_auc={val_auc:.5f} ({n_val_features} features)'
+
+    print(f'effective rank={rank:.1f} train_auc={auc:.5f} ({n_actual_features} features)' + val_message)
+    return rank, auc, val_rank, val_auc
 
 
 def train_one_epoch(model: torch.nn.Module,
                     data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, ema_scheduler,
+                    log_writer=None,
+                    val_loader: Iterable=None,
+                    do_analysis: bool=False,
+                    autocast_args: dict={},
+                    args=None):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    metric_logger.add_meter('ema_decay', 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 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)
+
+        if isinstance(samples, (list, tuple)):
+            (samples, targs) = samples
+        samples = samples.to(device, non_blocking=True)
+
+        with torch.cuda.amp.autocast(**autocast_args):
+            if args.model.startswith('m2d_vit'):
+                loss, *_ = model(samples, mask_ratio=args.mask_ratio)
+            elif args.model.startswith('m2d_s_vit'):
+                targs = targs.to(device, non_blocking=True)
+                loss, _, (_, _, loss_online, loss_offline) = model(samples, targs, mask_ratio=args.mask_ratio)
+            else:
+                assert False, f'Unknown model: {args.model}'
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        is_update_grad = (data_iter_step + 1) % accum_iter == 0
+        loss = loss / accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(), clip_grad=args.clip_grad, update_grad=is_update_grad)
+        if is_update_grad:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        # Update target network
+        global_step = len(data_loader) * epoch + data_iter_step
+        ema_decay = ema_scheduler(global_step)
+        if is_update_grad:
+            _model = model if hasattr(model, 'update_target_network') else model.module
+            _model.update_target_network(ema_decay)
+
+        metric_logger.update(loss=loss_value)
+        if args.model.startswith('m2d_s_vit'):
+            metric_logger.update(l_on=loss_online, l_off=loss_offline)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+        metric_logger.update(ema_decay=ema_decay)
+
+        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('train_loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+            log_writer.add_scalar('ema_decay', ema_decay, epoch_1000x)
+
+    # feature analysis
+    rank, auc, val_rank, val_auc = evaluate_feature(data_loader, val_loader, device, model) if do_analysis else (-1., -1., -1., -1.)
+    metric_logger.update(rank=float(rank), auc=float(auc))
+    if val_loader is not None:
+        metric_logger.update(val_rank=float(val_rank), val_auc=float(val_auc))
+
+    # 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()}
+
+
+# The original MAE training loop.
+def train_one_epoch_mae(model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
                     device: torch.device, epoch: int, loss_scaler,
                     log_writer=None,
+                    val_loader: Iterable=None,
+                    do_analysis: bool=False,
+                    autocast_args: dict={},
                     args=None):
     model.train(True)
     metric_logger = misc.MetricLogger(delimiter="  ")
@@ -36,15 +180,17 @@
     if log_writer is not None:
         print('log_dir: {}'.format(log_writer.log_dir))
 
-    for data_iter_step, (samples, _) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+    for data_iter_step, samples 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)
 
+        if isinstance(samples, (list, tuple)):
+            (samples, targs) = samples
         samples = samples.to(device, non_blocking=True)
 
-        with torch.cuda.amp.autocast():
+        with torch.cuda.amp.autocast(**autocast_args):
             loss, _, _ = model(samples, mask_ratio=args.mask_ratio)
 
         loss_value = loss.item()
@@ -53,8 +199,9 @@
             print("Loss is {}, stopping training".format(loss_value))
             sys.exit(1)
 
-        loss /= accum_iter
+        loss = loss / accum_iter
         loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    clip_grad=args.clip_grad,
                     update_grad=(data_iter_step + 1) % accum_iter == 0)
         if (data_iter_step + 1) % accum_iter == 0:
             optimizer.zero_grad()
@@ -75,8 +222,159 @@
             log_writer.add_scalar('train_loss', loss_value_reduce, epoch_1000x)
             log_writer.add_scalar('lr', lr, epoch_1000x)
 
+    # validation loss
+    if val_loader is not None:
+        model.eval()
+        N_TOTAL_VAL = 2048  # Total number of validation samples -- keep it small to train faster.
+        N_DL = len(val_loader)
+        N_VAL = max(1, (N_TOTAL_VAL + N_DL - 1)//N_DL)
+        for sample in val_loader:
+            sample = sample[torch.randperm(sample.shape[0])[:N_VAL]]  # random sample N_VAL items.
+            sample = sample.to(device, non_blocking=True)
+            with torch.cuda.amp.autocast(**autocast_args):
+                with torch.no_grad():
+                    loss, *_ = model(sample, mask_ratio=args.mask_ratio)
+                    metric_logger.update(val_loss=loss.item())
+
+    # feature analysis
+    rank, auc, val_rank, val_auc = evaluate_feature(data_loader, val_loader, device, model) if do_analysis else (-1., -1., -1., -1.)
+    metric_logger.update(rank=float(rank), auc=float(auc))
+    if val_loader is not None:
+        metric_logger.update(val_rank=float(val_rank), val_auc=float(val_auc))
+
+    # 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()}
+
+
+def train_one_epoch_m2dx(model: torch.nn.Module, teacher_model: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, ema_scheduler,
+                    log_writer=None,
+                    val_loader: Iterable=None,
+                    do_analysis: bool=False,
+                    autocast_args: dict={},
+                    args=None):
+    """M2D-X training loop."""
+    if teacher_model is not None: teacher_model.eval()
+    model.train(True)
+    _model = model if hasattr(model, 'update_target_network') else model.module
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    metric_logger.add_meter('ema_decay', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+
+    use_offline_target = model.use_offline_target if hasattr(model, 'use_offline_target') else model.module.use_offline_target
+    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 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)
+
+        x_online, x_target = samples[0].to(device, non_blocking=True), samples[1].to(device, non_blocking=True)
+
+        with torch.cuda.amp.autocast(**autocast_args):
+            if teacher_model is None:
+                loss, *others = model(x_online, x_target, mask_ratio=args.mask_ratio) if use_offline_target else model(x_online, mask_ratio=args.mask_ratio)
+                try:
+                    _, (_, _, loss_online, loss_offline) = others
+                except:
+                    loss_online, loss_offline = loss, None
+            else:
+                with torch.no_grad():
+                    z_target = teacher_model.encode_lms(x_target)
+                loss, _, (_, _, loss_online, loss_offline) = model(x_online, z_target, mask_ratio=args.mask_ratio)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        is_update_grad = (data_iter_step + 1) % accum_iter == 0
+        loss = loss / accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(), clip_grad=args.clip_grad, update_grad=is_update_grad)
+        if is_update_grad:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        # Update target network
+        global_step = len(data_loader) * epoch + data_iter_step
+        ema_decay = ema_scheduler(global_step)
+        if is_update_grad:
+            _model.update_target_network(ema_decay)
+            # Clip logit scale if a CLAP model owns it
+            _clip_logit_scale = getattr(_model, "clip_logit_scale", None)
+            if callable(_clip_logit_scale):
+                _clip_logit_scale()
+
+        metric_logger.update(loss=loss_value)
+        if loss_offline is not None:
+            metric_logger.update(l_on=loss_online, l_off=loss_offline)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+        metric_logger.update(ema_decay=ema_decay)
+
+        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('train_loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+            log_writer.add_scalar('ema_decay', ema_decay, epoch_1000x)
+
+    # validation loss
+    if val_loader is not None:
+        if teacher_model is not None: teacher_model.eval()
+        model.eval()
+        N_TOTAL_VAL = 2048  # Total number of validation samples -- keep it small to train faster.
+        N_DL = len(val_loader)
+        N_VAL = max(1, (N_TOTAL_VAL + N_DL - 1)//N_DL)
+        for sample in val_loader:
+            sample = sample[torch.randperm(sample.shape[0])[:N_VAL]]  # random sample N_VAL items.
+            sample = sample.to(device, non_blocking=True)
+            with torch.cuda.amp.autocast(**autocast_args):
+                with torch.no_grad():
+                    if teacher_model is None:
+                        loss, *others = model(sample, None, mask_ratio=args.mask_ratio) if use_offline_target else model(sample, mask_ratio=args.mask_ratio)
+                        try:
+                            _, (_, _, loss_online, loss_offline) = others
+                            metric_logger.update(val_loss=loss.item(), val_l_on = loss_online.item(), val_l_off=loss_offline.item())
+                        except:
+                            metric_logger.update(val_loss=loss.item())
+                    else:
+                        z_target = teacher_model.encode_lms(sample)
+                        loss, _, (_, _, loss_online, loss_offline) = model(sample, z_target, mask_ratio=args.mask_ratio)
+                        metric_logger.update(val_loss=loss.item(), val_l_on = loss_online.item(), val_l_off=loss_offline.item())
+
+    # feature analysis
+    rank, auc, val_rank, val_auc = evaluate_feature(data_loader, val_loader, device, model) if do_analysis else (-1., -1., -1., -1.)
+    metric_logger.update(rank=float(rank), auc=float(auc))
+    if val_loader is not None:
+        metric_logger.update(val_rank=float(val_rank), val_auc=float(val_auc))
+
+    # for CLIP model: logit_scale
+    _logit_scale = getattr(_model, "logit_scale", None)
+    if _logit_scale is not None:
+        metric_logger.update(logit_scale=float(_logit_scale.detach().cpu()))
+
+    if hasattr(_model, 'loss_m2d'):
+        metric_logger.update(lamda_m2d=_model.loss_m2d, lamda_off=_model.loss_off)
 
     # 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()}
\ ファイル末尾に改行がありません
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
--- _org/m2d/models_mae.py	2023-07-07 10:22:30.217586685 +0900
+++ m2d/models_mae.py	2026-02-21 08:21:56.724331273 +0900
@@ -9,14 +9,55 @@
 # DeiT: https://github.com/facebookresearch/deit
 # --------------------------------------------------------
 
+import os
+import sys
+sys.path.append(os.path.abspath(os.path.dirname(__file__)))
 from functools import partial
 
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
+from torch import distributed as dist
+import numpy as np
 
-from timm.models.vision_transformer import PatchEmbed, Block
+import timm
+from timm.models.vision_transformer import Block, Mlp, SwiGLU
 
-from util.pos_embed import get_2d_sincos_pos_embed
+from pos_embed import get_2d_sincos_pos_embed
+
+
+def expand_size(sz):
+    if isinstance(sz, int):
+        return [sz, sz]
+    return sz
+
+
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding -- borrowed from https://pypi.org/project/timm/0.4.12/
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+        img_size = expand_size(img_size)
+        patch_size = expand_size(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+from .masking import (data2vec2_block_masking_wrapper, random_unstructured_mask, object_centric_mask_with_hint,
+    random_weighted_mask_with_hint, random_weighted_mask_var_with_hint, determ_weighted_mask_with_hint, random_structured_mask, random_1d_mask)
 
 
 class MaskedAutoencoderViT(nn.Module):
@@ -25,8 +66,11 @@
     def __init__(self, img_size=224, patch_size=16, in_chans=3,
                  embed_dim=1024, depth=24, num_heads=16,
                  decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
-                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False):
+                 mlp_ratio=4., mlp_layer=Mlp, norm_layer=nn.LayerNorm, norm_pix_loss=False, norm_stats=None):
         super().__init__()
+        self.in_chans = in_chans
+        img_size, patch_size = expand_size(img_size), expand_size(patch_size)
+        self.norm_stats = nn.Parameter(torch.tensor([-7.1, 4.2] if norm_stats is None else norm_stats), requires_grad=False)
 
         # --------------------------------------------------------------------------
         # MAE encoder specifics
@@ -37,7 +81,7 @@
         self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding
 
         self.blocks = nn.ModuleList([
-            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, mlp_layer=mlp_layer)
             for i in range(depth)])
         self.norm = norm_layer(embed_dim)
         # --------------------------------------------------------------------------
@@ -51,24 +95,75 @@
         self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding
 
         self.decoder_blocks = nn.ModuleList([
-            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
+            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, mlp_layer=mlp_layer)
             for i in range(decoder_depth)])
 
         self.decoder_norm = norm_layer(decoder_embed_dim)
-        self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True) # decoder to patch
+        self.decoder_pred = nn.Linear(decoder_embed_dim, self.img_patch_dim(), bias=True) # decoder to patch
         # --------------------------------------------------------------------------
 
         self.norm_pix_loss = norm_pix_loss
 
+        print(f'{self.__class__.__name__}(in_chans={self.in_chans}, patch size={self.patch_size()}, grid_size={self.grid_size()},\n'
+              f'  embed_dim={embed_dim}, depth={depth}, num_heads={num_heads}, decoder_embed_dim={decoder_embed_dim},\n'
+              f'  decoder_depth={decoder_depth}, decoder_num_heads={decoder_num_heads}, mlp_ratio={mlp_ratio},\n'
+              f'  norm_pix_loss={norm_pix_loss}, norm_stats={self.norm_stats})')
+
         self.initialize_weights()
 
+        self._random_mask_fn = random_unstructured_mask
+        self.set_hint_ratio(1.0)
+
+    def set_data2vec2_block_masking(self):
+        print('using data2vec2_block_masking().')
+        self._random_mask_fn = data2vec2_block_masking_wrapper
+
+    def set_object_centric_mask_with_hint(self):
+        print('using object_centric_masking_hint().')
+        self._random_mask_fn = object_centric_mask_with_hint
+
+    def set_random_weighted_mask_with_hint(self):
+        print('using random_weighted_mask_with_hint().')
+        self._random_mask_fn = random_weighted_mask_with_hint
+
+    def set_random_weighted_mask_var_with_hint(self):
+        print('using random_weighted_mask_var_with_hint().')
+        self._random_mask_fn = random_weighted_mask_var_with_hint
+
+    def set_determ_weighted_mask_with_hint(self):
+        print('using determ_weighted_mask_with_hint().')
+        self._random_mask_fn = determ_weighted_mask_with_hint
+
+    def set_random_structured_mask(self):
+        print('using random_structured_mask().')
+        self._random_mask_fn = random_structured_mask
+
+    def set_random_1d_mask(self, one_mask_frames=20//2):
+        print(f'using random_1d_mask(one_mask_frames={one_mask_frames}).')
+        self._random_mask_fn = partial(random_1d_mask, one_mask_frames=one_mask_frames)
+
+    def patch_size(self):
+        return self.patch_embed.proj.kernel_size
+
+    def grid_size(self):
+        # This fails with timm 0.4.5 -> return self.patch_embed.grid_size
+        # Workaround for avoid compatibility issue
+        img_size = np.array(self.patch_embed.img_size)
+        patch_size = np.array(self.patch_embed.patch_size)
+        grid_size = img_size // patch_size
+        return grid_size
+
+    def img_patch_dim(self):
+        patch_size = self.patch_size()
+        return patch_size[0] * patch_size[1] * self.in_chans
+
     def initialize_weights(self):
         # initialization
         # initialize (and freeze) pos_embed by sin-cos embedding
-        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], self.grid_size(), cls_token=True)
         self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
 
-        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], self.grid_size(), cls_token=True)
         self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
 
         # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
@@ -94,46 +189,78 @@
 
     def patchify(self, imgs):
         """
-        imgs: (N, 3, H, W)
-        x: (N, L, patch_size**2 *3)
+        imgs: (N, C, H, W)
+        x: (N, L, patch_size[0]*patch_size[0]*in_chans)
         """
-        p = self.patch_embed.patch_size[0]
-        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
-
-        h = w = imgs.shape[2] // p
-        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
+        ph, pw = self.patch_size()
+        h, w = self.grid_size()
+        x = imgs.reshape(shape=(imgs.shape[0], self.in_chans, h, ph, w, pw))
         x = torch.einsum('nchpwq->nhwpqc', x)
-        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
+        x = x.reshape(shape=(imgs.shape[0], h * w, self.img_patch_dim()))
         return x
 
     def unpatchify(self, x):
         """
-        x: (N, L, patch_size**2 *3)
-        imgs: (N, 3, H, W)
+        x: (N, L, patch_size[0]*patch_size[0]*in_chans)
+        imgs: (N, C, H, W)
         """
-        p = self.patch_embed.patch_size[0]
-        h = w = int(x.shape[1]**.5)
+        ph, pw = self.patch_size()
+        h, w = self.grid_size()
         assert h * w == x.shape[1]
         
-        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
+        x = x.reshape(shape=(x.shape[0], h, w, ph, pw, self.in_chans))
         x = torch.einsum('nhwpqc->nchpwq', x)
-        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
+        imgs = x.reshape(shape=(x.shape[0], self.in_chans, h * ph, w * pw))
         return imgs
 
-    def random_masking(self, x, mask_ratio):
+    def set_hint_ratio_by_epoch(self, cur_epoch, total_epochs=300, start_epoch=0, k=2):
+        """Compute hint ratio h based on current epoch and total epochs.
+        h = 1.0 - (cur_epoch / total_epochs) ** 2
+        https://arxiv.org/pdf/2206.10207
+        """
+        ratio = 1.0 - (cur_epoch / total_epochs) ** k if cur_epoch >= start_epoch else 1.0
+        self.set_hint_ratio(ratio)
+    def set_hint_ratio(self, hint_ratio):
+        self.hint_ratio_value = hint_ratio
+        print(f'Set masking hint ratio to {self.hint_ratio_value:.4f}.')
+    def hint_ratio(self):
+        return self.hint_ratio_value
+
+    def random_masking(self, x, mask_ratio, adjust_short=False, raw_input=None):
         """
         Perform per-sample random masking by per-sample shuffling.
         Per-sample shuffling is done by argsort random noise.
         x: [N, L, D], sequence
         """
         N, L, D = x.shape  # batch, length, dim
-        len_keep = int(L * (1 - mask_ratio))
-        
-        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
-        
-        # sort noise for each sample
-        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
-        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        if isinstance(mask_ratio, (torch.Tensor, np.ndarray)):
+            # Prefixed mask. `mask` shall be 2x sized.
+            mask = mask_ratio.clone().detach()
+            #ids_shuffle = torch.where(mask.reshape(N, -1) == 0)[1].reshape(N, -1)
+            ids_shuffle = torch.argsort(mask.reshape(N, -1), dim=1)
+            ids_restore = torch.argsort(ids_shuffle, dim=1)
+            len_keep = (mask[0] == 0).sum() // 2
+        elif isinstance(mask_ratio, (list, tuple)):
+            # Prefixed ids_restore & len_keep.
+            ids_restore = mask_ratio[0]
+            ids_shuffle = torch.argsort(ids_restore, dim=1)
+            len_keep = mask_ratio[1]
+        elif mask_ratio == 0:
+            # No mask
+            mask = torch.zeros([N, L], device=x.device)
+            ids_restore = torch.tensor(list(range(L))).to(torch.int)
+            return x, mask, ids_restore
+        else:
+            # Random mask
+            HorF, WorT = self.grid_size()
+            if adjust_short and L < HorF * WorT:
+                # audio: shorten pos_embed for a short input
+                WorT = L // HorF
+            ids_shuffle, len_keep = self._random_mask_fn((N, HorF, WorT), mask_ratio, x.device,
+                                                         data=raw_input, patch_size=self.patch_size(),
+                                                         hint_ratio=self.hint_ratio(), model=self)
+            ids_restore = torch.argsort(ids_shuffle, dim=1)
 
         # keep the first subset
         ids_keep = ids_shuffle[:, :len_keep]
@@ -147,15 +274,22 @@
 
         return x_masked, mask, ids_restore
 
-    def forward_encoder(self, x, mask_ratio):
+    def forward_encoder(self, org_x, mask_ratio, return_layers=False, adjust_short=False):
         # embed patches
-        x = self.patch_embed(x)
+        x = self.patch_embed(org_x)
 
         # add pos embed w/o cls token
-        x = x + self.pos_embed[:, 1:, :]
+        pos_embed = self.pos_embed[:, 1:, :]
+        if adjust_short and x.shape[1] < pos_embed.shape[1]:
+            # audio: shorten pos_embed for a short input
+            dims = pos_embed.shape[-1]
+            fbins = self.grid_size()[0]
+            frames = x.shape[1] // fbins
+            pos_embed = pos_embed.reshape(1, fbins, -1, dims)[:, :, :frames, :].reshape(1, fbins*frames, dims)
+        x = x + pos_embed
 
         # masking: length -> length * mask_ratio
-        x, mask, ids_restore = self.random_masking(x, mask_ratio)
+        x, mask, ids_restore = self.random_masking(x, mask_ratio, adjust_short=adjust_short, raw_input=org_x)
 
         # append cls token
         cls_token = self.cls_token + self.pos_embed[:, :1, :]
@@ -163,12 +297,31 @@
         x = torch.cat((cls_tokens, x), dim=1)
 
         # apply Transformer blocks
+        layers = []
         for blk in self.blocks:
             x = blk(x)
+            if return_layers: layers.append(x)
         x = self.norm(x)
+        if return_layers:
+            layers.pop() # replace the last feature with the normalized one.
+            layers.append(x)
 
+        if return_layers:
+            return torch.stack(layers), mask, ids_restore
         return x, mask, ids_restore
 
+    def drop_cls_token(self, latent):
+        # remove cls token [B, 1+H*W: D] -> [B, H*W, D]
+        # L = latent.shape[-2]
+        # assert L == 1 + self.grid_size()[0]*self.grid_size()[1], f'Already no class token...? {L} vs {self.grid_size()}'
+        return  latent[:, 1:, :]
+
+    def get_cls_token(self, latent):
+        # return cls token only [B, 1+H*W: D] -> [B, 1, D]
+        # L = latent.shape[-2]
+        # assert L == 1 + self.grid_size()[0]*self.grid_size()[1], f'No class token...? {L} vs {self.grid_size()}'
+        return  latent[:, :1, :]
+
     def forward_decoder(self, x, ids_restore):
         # embed tokens
         x = self.decoder_embed(x)
@@ -197,8 +350,8 @@
 
     def forward_loss(self, imgs, pred, mask):
         """
-        imgs: [N, 3, H, W]
-        pred: [N, L, p*p*3]
+        imgs: [N, C, H, W]
+        pred: [N, L, ph*pw*C]
         mask: [N, L], 0 is keep, 1 is remove, 
         """
         target = self.patchify(imgs)
@@ -208,6 +361,9 @@
             target = (target - mean) / (var + 1.e-6)**.5
 
         loss = (pred - target) ** 2
+        # if torch.isnan(loss).any():
+        #     print('loss contains nan(s), which is replaced with 0...')
+        #     loss = torch.nan_to_num(loss)
         loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
 
         loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
@@ -219,6 +375,17 @@
         loss = self.forward_loss(imgs, pred, mask)
         return loss, pred, mask
 
+    def forward_viz(self, imgs, mask_ratio=0.75):
+        loss, pred, mask = self.forward(imgs, mask_ratio)
+        # recons_as_is = self.unpatchify(pred)
+        # overwrite visible patches with original image.
+        pred_org_on_mask = pred.clone()
+        visible = (mask == 0.)
+        pred_org_on_mask[visible] = self.patchify(imgs)[visible]
+        recons = self.unpatchify(pred_org_on_mask)
+        errormap = ((recons - imgs) ** 2).sqrt()
+        return loss, recons, errormap, mask.reshape(mask.shape[0], *self.grid_size())
+
 
 def mae_vit_base_patch16_dec512d8b(**kwargs):
     model = MaskedAutoencoderViT(
@@ -248,3 +415,950 @@
 mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
 mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b  # decoder: 512 dim, 8 blocks
 mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b  # decoder: 512 dim, 8 blocks
+
+
+def mae_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    model = MaskedAutoencoderViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def msm_mae_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    print(f'MSM-MAE **FORCED DEC DEPTH AS 4 (Your decoder_depth={decoder_depth} is ignored)**')
+    model = MaskedAutoencoderViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=384, decoder_depth=4, decoder_num_heads=6,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+# Masked Modeling Duo (M2D)
+def gather_features(set_of_features, rank=0, world_size=1):
+    # Simplified version from: https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/loss.py#L66
+    def gather_features_one(features):
+        gathered_features = [torch.zeros_like(features) for _ in range(world_size)]
+        dist.all_gather(gathered_features, features)
+
+        # ensure grads for local rank when all_* features don't have a gradient
+        gathered_features[rank] = features
+
+        all_features = torch.cat(gathered_features, dim=0)
+        return all_features
+
+    all_features = [gather_features_one(features) for features in set_of_features]
+
+    return all_features
+
+
+class ClipLoss(nn.Module):
+    # Simplified version from: https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/loss.py#L66
+    def __init__(
+            self,
+            rank=0,
+            world_size=1,
+    ):
+        super().__init__()
+        self.rank = rank
+        self.world_size = world_size
+        print(f' ClipLoss with rank={rank}, world_size={world_size}')
+
+    def get_ground_truth(self, device, num_logits) -> torch.Tensor:
+        # calculated ground-truth and cache if enabled
+        labels = torch.arange(num_logits, device=device, dtype=torch.long)
+        return labels
+
+    def get_logits(self, image_features, text_features, logit_scale):
+        image_features = F.normalize(image_features, dim=-1, p=2)
+        text_features = F.normalize(text_features, dim=-1, p=2)
+
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(
+                [image_features, text_features],
+                rank=self.rank, world_size=self.world_size)
+
+            logits_per_image = logit_scale * all_image_features @ all_text_features.T
+            logits_per_text = logits_per_image.T
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+
+        return logits_per_image, logits_per_text
+
+    def forward(self, image_features, text_features, logit_scale=None):
+        logit_scale = (1 / 0.07) if logit_scale is None else logit_scale  # the initial value used in CLIP
+        device = image_features.device
+        logits_per_image, logits_per_text = self.get_logits(image_features, text_features, logit_scale)
+        labels = self.get_ground_truth(device, logits_per_image.shape[0])
+        total_loss = (F.cross_entropy(logits_per_image, labels) + F.cross_entropy(logits_per_text, labels)) / 2
+        total_loss = total_loss
+
+        return total_loss
+
+
+def info_nce_mi(clip_loss_fn, z1, z2):
+    """ Estimate MI based on the InfoNCE loss, as in Oord et al., "Representation Learning with Contrastive Predictive Coding" """
+    info_nce_loss = clip_loss_fn(z1, z2)
+    effective_batch_size = z1.shape[0] * clip_loss_fn.world_size
+    mi_lower_bound = torch.log(torch.tensor(effective_batch_size, device=info_nce_loss.device)) - info_nce_loss
+    return mi_lower_bound
+
+
+def shared_dim_participation_ratio(z1, z2, eps=1e-12):
+    """ Compute the shared dimension participation ratio between two representations."""
+    # center
+    z1c = z1 - z1.mean(dim=0, keepdim=True)
+    z2c = z2 - z2.mean(dim=0, keepdim=True)
+
+    # cross-covariance
+    B = z1.shape[0]
+    C = z1c.T @ z2c / (B - 1)
+
+    # singular values
+    S = torch.linalg.svdvals(C.to(torch.float32))
+
+    # participation ratio
+    num = (S**2).sum()**2
+    den = (S**4).sum() + eps
+
+    return (num / den).to(z1.dtype)
+
+
+def set_requires_grad(model, val):
+    for p in model.parameters():
+        p.requires_grad = val
+
+
+def ema_model_weight(decay, old_model, new_model):
+    def ema(decay, old, new):
+        return old * decay + (1 - decay) * new
+
+    for new_params, old_params in zip(new_model.parameters(), old_model.parameters()):
+        old_weight, new_weight = old_params.data, new_params.data
+        old_params.data = ema(decay, old_weight, new_weight)
+
+
+class M2DViT(MaskedAutoencoderViT):
+    """ Masked Modeling Duo (M2D) implementation based on the MAE.
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., mlp_layer=Mlp, norm_layer=nn.LayerNorm, norm_pix_loss=False, norm_stats=None,
+                 loss_type='norm_mse', target_layers=None, world_rank=None, **kwargs):
+        super().__init__(img_size=img_size, patch_size=patch_size, in_chans=in_chans,
+                 embed_dim=embed_dim, depth=depth, num_heads=num_heads,
+                 decoder_embed_dim=decoder_embed_dim, decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
+                 mlp_ratio=mlp_ratio, mlp_layer=mlp_layer, norm_layer=norm_layer, norm_pix_loss=norm_pix_loss, norm_stats=norm_stats)
+        self.loss_type = loss_type
+        self.target_layers = target_layers
+        print(f'+ loss_type={loss_type}, target_layers={target_layers}')
+        if len(kwargs.keys()) > 0:
+            print(' CAUTION: You set unknown arguments ->', kwargs)
+        self.use_offline_target = False
+        world_size, rank = [1, 0] if world_rank is None else world_rank
+        self.clip_loss_fn = ClipLoss(rank=rank, world_size=world_size)
+
+        # --------------------------------------------------------------------------
+        # Target encoder specifics
+        self.target_blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, mlp_layer=mlp_layer)
+            for i in range(depth)])
+        self.target_norm = norm_layer(embed_dim)
+        set_requires_grad(self.target_blocks, False)
+        set_requires_grad(self.target_norm, False)
+        self.target_blocks.apply(self._init_weights)
+        self.target_norm.apply(self._init_weights)
+        # --------------------------------------------------------------------------
+
+        # --------------------------------------------------------------------------
+        # Decoder specifics
+        self.decoder_pred = nn.Linear(decoder_embed_dim, embed_dim, bias=True) # predict target embeddings
+        self.decoder_pred.apply(self._init_weights)
+        # --------------------------------------------------------------------------
+
+    def update_target_network(self, ema_decay):
+        ema_model_weight(ema_decay, self.target_blocks, self.blocks)
+        ema_model_weight(ema_decay, self.target_norm, self.norm)
+
+    def random_masking(self, x, mask_ratio, adjust_short=False, raw_input=None):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+
+        if isinstance(mask_ratio, (torch.Tensor, np.ndarray)):
+            # Prefixed mask. `mask` shall be 2x sized.
+            mask = mask_ratio.clone().detach()
+            #ids_shuffle = torch.where(mask.reshape(N, -1) == 0)[1].reshape(N, -1)
+            ids_shuffle = torch.argsort(mask.reshape(N, -1), dim=1)
+            ids_restore = torch.argsort(ids_shuffle, dim=1)
+            len_keep = (mask[0] == 0).sum() // 2
+        elif isinstance(mask_ratio, (list, tuple)):
+            # Prefixed ids_restore & len_keep.
+            ids_restore = mask_ratio[0]
+            ids_shuffle = torch.argsort(ids_restore, dim=1)
+            len_keep = mask_ratio[1]
+        elif mask_ratio == 0:
+            # No mask
+            mask = torch.zeros([N, L], device=x.device)
+            ids_restore = torch.tensor(list(range(L))).to(torch.int)
+            return x, None, mask, ids_restore
+        else:
+            # Random mask
+            HorF, WorT = self.grid_size()
+            if adjust_short and L < HorF * WorT:
+                # audio: shorten pos_embed for a short input
+                WorT = L // HorF
+            ids_shuffle, len_keep = self._random_mask_fn((N, HorF, WorT), mask_ratio, x.device,
+                                                         data=raw_input, patch_size=self.patch_size(),
+                                                         hint_ratio=self.hint_ratio(), model=self)
+            ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the visible patch indexes
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # keep the rest
+        ids_keep = ids_shuffle[:, len_keep:]
+        x_masked2 = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, x_masked2, mask, ids_restore
+
+    def forward_encoder(self, org_x, mask_ratio, return_layers=False, blocks=None, norm=None, adjust_short=False):
+        blocks, norm = blocks or self.blocks, norm or self.norm
+
+        # embed patches
+        x = self.patch_embed(org_x)
+
+        # add pos embed w/o cls token
+        pos_embed = self.pos_embed[:, 1:, :]
+        if adjust_short and x.shape[1] < pos_embed.shape[1]:
+            # audio: shorten pos_embed for a short input
+            dims = pos_embed.shape[-1]
+            fbins = self.grid_size()[0]
+            frames = x.shape[1] // fbins
+            pos_embed = pos_embed.reshape(1, fbins, -1, dims)[:, :, :frames, :].reshape(1, fbins*frames, dims)
+        x = x + pos_embed
+
+        # masking: length -> length * mask_ratio
+        x, x_targ, mask, ids_restore = self.random_masking(x, mask_ratio, adjust_short=adjust_short, raw_input=org_x)
+
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        layers = []
+        for blk in blocks:
+            x = blk(x)
+            if return_layers: layers.append(x)
+        x = norm(x)
+        if return_layers:
+            layers.pop() # replace the last feature with the normalized one.
+            layers.append(x)
+
+        if return_layers:
+            return torch.stack(layers), x_targ, mask, ids_restore
+        return x, x_targ, mask, ids_restore
+
+    def forward_decoder(self, x, ids_restore, keep_cls=False, also_pred_asis=False):
+        len_keep = x.shape[1] - 1 # tokens - cls
+
+        # embed tokens
+        x = self.decoder_embed(x)
+        D = x.shape[-1]
+
+        # append mask tokens to sequence
+        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
+        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token
+        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, D))  # unshuffle
+        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token
+
+        # add pos embed
+        x = x + self.decoder_pos_embed
+
+        # apply Transformer blocks
+        for blk in self.decoder_blocks:
+            x = blk(x)
+        x = self.decoder_norm(x)
+
+        # predictor projection
+        x = self.decoder_pred(x)
+
+        # remove cls token
+        y = self.drop_cls_token(x)
+        y_pred_asis = y
+        # re-shuffle, and keep prediction only
+        ids_shuffle = torch.argsort(ids_restore, dim=1)
+        y = torch.gather(y, dim=1, index=ids_shuffle.unsqueeze(-1).repeat(1, 1, y.shape[-1]))
+        y = y[:, len_keep:] # prediction only
+
+        # append cls if needed
+        if keep_cls:
+            y = torch.cat([x[:, :1, :], y], dim=1)
+        if also_pred_asis:
+            return y, y_pred_asis
+        return y
+
+    def forward_target_encoder(self, x_targ, drop_cls=True):
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x_targ.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x_targ), dim=1)
+
+        # apply Transformer blocks
+        xs = []
+        for l, blk in enumerate(self.target_blocks):
+            x = blk(x)
+            if self.target_layers and l in self.target_layers:
+                xs.append(x)
+        if xs:
+            x = torch.stack(xs).mean(0)
+        x = self.target_norm(x)
+
+        # remove cls token
+        if drop_cls:
+            x = self.drop_cls_token(x)
+
+        return x
+
+    def forward_loss(self, target, pred, norm_pix_loss, loss_type):
+        """
+        target: [N, targL, D]
+        pred: [N, targL, D]
+        """
+
+        if norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.e-6)**.5
+
+        if loss_type == 'mse':
+            loss = (pred - target) ** 2
+            loss = loss.mean(dim=-1)  # [N, L], mean loss per predicted patch embedding
+        elif loss_type == 'norm_mse':
+            target = torch.nn.functional.normalize(target, dim=-1, p=2)
+            pred = torch.nn.functional.normalize(pred, dim=-1, p=2)
+            loss = target * pred
+            loss = 2 - 2 * loss.sum(dim=-1)
+        else:
+            assert loss_type in ['WE NEED A KNOWN LOSS FN']
+
+        loss = loss.mean()
+
+        return loss
+
+    def forward(self, imgs, mask_ratio=0.7):
+        latent, x_targ, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        pred = self.forward_decoder(latent, ids_restore)  # [N, targL, D]
+        with torch.no_grad():
+            target = self.forward_target_encoder(x_targ)
+
+        loss = self.forward_loss(target, pred, self.norm_pix_loss, self.loss_type)
+        return loss, pred, (ids_restore, mask)
+
+    # def forward(self, imgs, mask_ratio=0.7):  # For monitoring MI between masked/visible patches.
+    #     latent, x_targ, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+    #     pred = self.forward_decoder(latent, ids_restore)  # [N, targL, D]
+    #     with torch.no_grad():
+    #         target = self.forward_target_encoder(x_targ)
+
+    #     loss = self.forward_loss(target, pred, self.norm_pix_loss, self.loss_type)
+    #     # calculate metrics for views
+    #     estimated_mi = info_nce_mi(self.clip_loss_fn, self.drop_cls_token(latent).mean(1), target.mean(1))
+    #     shared_pr = shared_dim_participation_ratio(self.drop_cls_token(latent).mean(1), target.mean(1)).mean()
+    #     return loss, pred, (ids_restore, mask, estimated_mi, shared_pr)
+
+    def forward_viz(self, imgs, mask_ratio=0.7):
+        # Visualize the input and mask.
+        loss, pred, (ids_restore, mask) = self.forward(imgs, mask_ratio)
+        recons, errormap = None, None
+        return recons, errormap, mask.reshape(mask.shape[0], *self.grid_size())
+
+
+def m2d_vit_base_patch16_dec512d8b(**kwargs):  # for image
+    model = M2DViT(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+m2d_vit_base_patch16 = m2d_vit_base_patch16_dec512d8b  # for image
+
+
+def m2d_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):  # for audio
+    model = M2DViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+# M2D variants
+
+class M2D_D2ViT(M2DViT):
+    """A data2vec-like M2D variant that feeds all patches to the target network."""
+
+    def random_masking(self, x, mask_ratio, adjust_short=False):
+        """Random masking that returns all patches as the x_target.
+
+        Returns:
+            x_masked: Maked patches
+            x_target: Target patches = all patches for Data2Vec
+            mask: Mask
+            ids_restore: indexes for restoration of masked patches
+        """
+        x_masked, _, mask, ids_restore = super().random_masking(x, mask_ratio, adjust_short=adjust_short)
+        return x_masked, x, mask, ids_restore
+
+    def forward(self, imgs, mask_ratio=0.75):
+        latent, x_targ, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        pred = self.forward_decoder(latent, ids_restore)  # [N, targL, D]
+        with torch.no_grad():
+            # x_targ holds all the input patches
+            target = self.forward_target_encoder(x_targ)
+            len_keep = latent.shape[1] - 1 # tokens - cls
+            ids_shuffle = torch.argsort(ids_restore, dim=1)
+            ids_keep = ids_shuffle[:, len_keep:]
+            # target to leave masked patch representations only 
+            target = torch.gather(target, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, target.shape[-1]))
+
+        loss = self.forward_loss(target, pred, norm_pix_loss=self.norm_pix_loss, loss_type=self.loss_type)
+
+        # calculate metrics for views
+        estimated_mi = info_nce_mi(self.clip_loss_fn, self.drop_cls_token(latent).mean(1), target.mean(1))
+        shared_pr = shared_dim_participation_ratio(self.drop_cls_token(latent).mean(1), target.mean(1)).mean()
+        return loss, pred, (ids_restore, mask, estimated_mi, shared_pr)
+
+
+def m2d_d2v_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):  # for audio
+    model = M2D_D2ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def m2d_d2v_vit_base_patch16_dec512d8b(**kwargs):  # for image
+    model = M2D_D2ViT(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+m2d_d2v_vit_base_patch16 = m2d_d2v_vit_base_patch16_dec512d8b  # for image
+
+
+class M2DEncoderViT(timm.models.vision_transformer.VisionTransformer):
+    """ Vision Transformer with support for global average pooling
+    """
+    def __init__(self, **kwargs):
+        norm_stats = kwargs['norm_stats'] if 'norm_stats' in kwargs else [-7.1, 4.2]
+        kwargs = {k:v for k, v in kwargs.items() if k not in ['decoder_depth', 'off_emb_dim', 'clip_args', 'norm_stats']}
+        super().__init__(**kwargs)
+        self.norm_stats = nn.Parameter(torch.tensor(norm_stats).clone().detach(), requires_grad=False)
+        # Our positional encoding is constant
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], self.grid_size(), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+        self.pos_embed.requires_grad_(False)
+        # Replace PatchEmbed to avoid unintended assertion failure. ex) AssertionError: Input image width (102) doesn't match model (608).
+        self.patch_embed = PatchEmbed(self.patch_embed.img_size, self.patch_embed.patch_size,
+                                      self.patch_embed.proj.in_channels, self.patch_embed.proj.out_channels)
+        # We do not use `head` for the M2D encoder only ViT
+        del self.head
+
+    def patch_size(self):
+        return self.patch_embed.proj.kernel_size
+
+    def grid_size(self):
+        # This fails with timm 0.4.5 -> return self.patch_embed.grid_size
+        # Workaround for avoid compatibility issue
+        img_size = np.array(self.patch_embed.img_size)
+        patch_size = np.array(self.patch_embed.patch_size)
+        grid_size = img_size // patch_size
+        return grid_size
+
+    def set_random_structured_mask(self):
+        print(' using random_structured_mask().')
+        self._random_mask_fn = random_structured_mask
+
+    def set_random_unstructured_mask(self):
+        print(' using random_unstructured_mask().')
+        self._random_mask_fn = random_unstructured_mask
+
+    def random_masking(self, x, mask_ratio, adjust_short=False):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.shape  # batch, length, dim
+
+        if isinstance(mask_ratio, (torch.Tensor, np.ndarray)):
+            # Prefixed mask. `mask` shall be 2x sized.
+            mask = mask_ratio.clone().detach()
+            #ids_shuffle = torch.where(mask.reshape(N, -1) == 0)[1].reshape(N, -1)
+            ids_shuffle = torch.argsort(mask.reshape(N, -1), dim=1)
+            ids_restore = torch.argsort(ids_shuffle, dim=1)
+            len_keep = (mask[0] == 0).sum() // 2
+        elif isinstance(mask_ratio, (list, tuple)):
+            # Prefixed ids_restore & len_keep.
+            ids_restore = mask_ratio[0]
+            ids_shuffle = torch.argsort(ids_restore, dim=1)
+            len_keep = mask_ratio[1]
+        elif mask_ratio == 0:
+            # No mask
+            mask = torch.zeros([N, L], device=x.device)
+            ids_restore = torch.tensor(list(range(L))).to(torch.int)
+            return x, None, mask, ids_restore
+        else:
+            # Random mask
+            HorF, WorT = self.grid_size()
+            if adjust_short and L < HorF * WorT:
+                # audio: shorten pos_embed for a short input
+                WorT = L // HorF
+            ids_shuffle, len_keep = self._random_mask_fn((N, HorF, WorT), mask_ratio, x.device)
+            ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the visible patch indexes
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # keep the rest
+        ids_keep = ids_shuffle[:, len_keep:]
+        x_masked2 = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, x_masked2, mask, ids_restore
+
+    def forward_encoder(self, x, mask_ratio, return_layers=False, blocks=None, norm=None, adjust_short=False):
+        blocks, norm = blocks or self.blocks, norm or self.norm
+
+        # embed patches
+        x = self.patch_embed(x)
+
+        # add pos embed w/o cls token
+        pos_embed = self.pos_embed[:, 1:, :]
+        if adjust_short and x.shape[1] < pos_embed.shape[1]:
+            # audio: shorten pos_embed for a short input
+            dims = pos_embed.shape[-1]
+            fbins = self.grid_size()[0]
+            frames = x.shape[1] // fbins
+            pos_embed = pos_embed.reshape(1, fbins, -1, dims)[:, :, :frames, :].reshape(1, fbins*frames, dims)
+        x = x + pos_embed
+
+        # masking: length -> length * mask_ratio
+        x, x_targ, mask, ids_restore = self.random_masking(x, mask_ratio, adjust_short=adjust_short)
+
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # apply Transformer blocks
+        layers = []
+        for blk in blocks:
+            x = blk(x)
+            if return_layers: layers.append(x)
+        x = norm(x)
+        if return_layers:
+            layers.pop() # replace the last feature with the normalized one.
+            layers.append(x)
+
+        if return_layers:
+            return torch.stack(layers), x_targ, mask, ids_restore
+        return x, x_targ, mask, ids_restore
+
+
+def m2d_vit_base_encoder_only(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):  # for audio, encoder only
+    model = M2DEncoderViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+msm_mae_vit_base_encoder_only = m2d_vit_base_encoder_only
+m2d_d2v_vit_base_encoder_only = m2d_vit_base_encoder_only
+
+
+def m2dsg_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):  # for audio
+    model = M2DViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_layer=SwiGLU,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+def m2dsg_vit_base_encoder_only(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):  # for audio, encoder only
+    model = M2DEncoderViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_layer=SwiGLU,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+# Masked Modeling Duo for X (M2D-X)
+
+class M2D_X_ViT(M2DViT):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., mlp_layer=Mlp, norm_layer=nn.LayerNorm, norm_pix_loss=False, norm_stats=None,
+                 loss_type='norm_mse', target_layers=None,
+                 loss_m2d=1.0, loss_off=0.0, off_emb_dim=3840, world_rank=None, **kwargs):
+        super().__init__(img_size=img_size, patch_size=patch_size, in_chans=in_chans,
+                 embed_dim=embed_dim, depth=depth, num_heads=num_heads,
+                 decoder_embed_dim=decoder_embed_dim, decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
+                 mlp_ratio=mlp_ratio, mlp_layer=mlp_layer, norm_layer=norm_layer, norm_pix_loss=norm_pix_loss, norm_stats=norm_stats,
+                 loss_type=loss_type, target_layers=target_layers, world_rank=world_rank)
+        self.loss_m2d = loss_m2d
+        self.loss_off = loss_off
+        print(f'+ loss_m2d={loss_m2d}, loss_off={loss_off}, off_emb_dim={off_emb_dim}')
+        self.use_offline_target = True
+
+        F = self.grid_size()[0]
+        self.offline_predictor = nn.Linear(F * embed_dim, off_emb_dim)
+        self.offline_predictor.apply(self._init_weights)
+
+        if len(kwargs) > 0:
+            print(f' **NOTE** ignored args: {kwargs}')
+
+    def forward_off_loss(self, target, pred):
+        _, T , _ = pred.shape
+        _, T2, _ = target.shape
+        assert T2 == T, f'Time frame resolution mismatch: offline GT {T2} != prediction {T}'
+        return self.forward_loss(target, pred, self.norm_pix_loss, self.loss_type)
+
+    def forward(self, imgs, target_offline, mask_ratio=0.7):
+        # online
+        latent, x_targ, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        pred_online = self.forward_decoder(latent, ids_restore)  # [N, targL, D]
+        with torch.no_grad():
+            target_online = self.forward_target_encoder(x_targ)
+        loss_online = self.forward_loss(target_online, pred_online, self.norm_pix_loss, self.loss_type)
+
+        # offline
+        z_online = torch.cat([self.drop_cls_token(latent), pred_online], dim=1)  # gather all the online tokes
+        z_online = torch.gather(z_online, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, latent.shape[-1]))  # unshuffle
+        B, FT, D = z_online.shape
+        F = self.grid_size()[0]
+        T = FT // F
+        z_online = torch.einsum('bftd->btfd', z_online.reshape(B, F, T, D)).reshape(B, T, F * D)
+        pred_offline = self.offline_predictor(z_online)
+
+        if target_offline is None:
+            loss_offline = torch.zeros_like(loss_online)
+        else:
+            loss_offline = self.forward_off_loss(target_offline, pred_offline)
+
+        loss = self.loss_m2d * loss_online + self.loss_off * loss_offline
+        return loss, pred_online, (ids_restore, mask, loss_online, loss_offline)
+
+    def forward_viz(self, imgs, mask_ratio=0.7):
+        latent, x_targ, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        return None, None, mask.reshape(mask.shape[0], *self.grid_size())  # only the mask to visualize
+
+
+def m2d_x_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs): # for audio
+    model = M2D_X_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+m2d_x_vit_base_encoder_only = m2d_vit_base_encoder_only
+
+
+class M2D_AS_ViT(M2D_X_ViT):
+    """M2D for AudioSet"""
+    def __init__(self, img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False, norm_stats=None,
+                 loss_type='norm_mse', target_layers=None,
+                 loss_m2d=1.0, loss_off=1.0, off_emb_dim=3840, **kwargs):
+        super().__init__(img_size=img_size, patch_size=patch_size, in_chans=in_chans,
+                 embed_dim=embed_dim, depth=depth, num_heads=num_heads,
+                 decoder_embed_dim=decoder_embed_dim, decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
+                 mlp_ratio=mlp_ratio, norm_layer=norm_layer, norm_pix_loss=norm_pix_loss, norm_stats=norm_stats,
+                 loss_type=loss_type, target_layers=target_layers,
+                 loss_m2d=loss_m2d, loss_off=loss_off, off_emb_dim=off_emb_dim, **kwargs)
+
+    def forward_off_loss(self, target, pred):
+        pred = pred.mean(1)  # (B, L, D) -> (B, D)
+        return F.binary_cross_entropy_with_logits(pred, target)
+
+
+def m2d_as_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs): # for audio
+    model = M2D_AS_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+m2d_as_vit_base_encoder_only = m2d_vit_base_encoder_only
+
+
+# Masked Modeling Duo for Speech (M2D-S)
+
+class M2D_S_ViT(M2DViT):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False, norm_stats=None,
+                 loss_type='norm_mse', target_layers=None,
+                 loss_m2d=0.7, loss_off=0.3, off_emb_dim=768):
+        super().__init__(img_size=img_size, patch_size=patch_size, in_chans=in_chans,
+                 embed_dim=embed_dim, depth=depth, num_heads=num_heads,
+                 decoder_embed_dim=decoder_embed_dim, decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
+                 mlp_ratio=mlp_ratio, norm_layer=norm_layer, norm_pix_loss=norm_pix_loss, norm_stats=norm_stats,
+                 loss_type=loss_type, target_layers=target_layers)
+        self.loss_m2d = loss_m2d
+        self.loss_off = loss_off
+        print(f'+ loss_m2d={loss_m2d}, loss_off={loss_off}')
+        self.use_offline_target = True
+
+        F = self.grid_size()[0]
+        self.offline_predictor = nn.Linear(F * embed_dim, off_emb_dim)
+
+    def forward(self, imgs, target_offline, mask_ratio=0.7):
+        # online
+        latent, x_targ, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        pred_online = self.forward_decoder(latent, ids_restore)  # [N, targL, D]
+        with torch.no_grad():
+            target_online = self.forward_target_encoder(x_targ)
+        loss_online = self.forward_loss(target_online, pred_online, self.norm_pix_loss, self.loss_type)
+
+        # offline
+        z_online = torch.cat([self.drop_cls_token(latent), pred_online], dim=1)  # gather all the online tokes
+        z_online = torch.gather(z_online, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, latent.shape[-1]))  # unshuffle
+        B, FT, D = z_online.shape
+        F = self.grid_size()[0]
+        T = FT // F
+        z_online = torch.einsum('bftd->btfd', z_online.reshape(B, F, T, D)).reshape(B, T, F * D)
+        pred_offline = self.offline_predictor(z_online)
+
+        B, T2, D = target_offline.shape
+        assert T2 == T, f'label length {T2} != patch length'
+
+        loss_offline = self.forward_loss(target_offline, pred_offline, self.norm_pix_loss, self.loss_type)
+
+        loss = self.loss_m2d * loss_online + self.loss_off * loss_offline
+        return loss, pred_online, (ids_restore, mask, loss_online, loss_offline)
+
+
+def m2d_s_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs): # for audio
+    model = M2D_S_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+m2d_s_vit_base_encoder_only = m2d_vit_base_encoder_only
+
+
+## LARGE MODELS
+
+def m2d_vit_large(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    model = M2DViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def m2d_x_vit_large(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    model = M2D_X_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def m2d_as_vit_large(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    model = M2D_AS_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def msm_mae_vit_large(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    print(f'MSM-MAE **FORCED DEC DEPTH AS 4 (Your decoder_depth={decoder_depth} is ignored)**')
+    model = MaskedAutoencoderViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=384, decoder_depth=4, decoder_num_heads=6,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def m2d_vit_large_encoder_only(patch_size=16, in_chans=1, **kwargs):  # for audio, encoder only
+    model = M2DEncoderViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+m2d_x_vit_large_encoder_only = m2d_vit_large_encoder_only
+m2d_as_vit_large_encoder_only = m2d_vit_large_encoder_only
+msm_mae_vit_large_encoder_only = m2d_vit_large_encoder_only
+
+
+# M2D-CLAP
+
+def get_MLP_projector(embed_dim, proj_hidden_dim, out_embed_dim):
+    projector = torch.nn.Sequential(
+        torch.nn.Linear(embed_dim, proj_hidden_dim),
+        torch.nn.ReLU(),
+        torch.nn.Linear(proj_hidden_dim, out_embed_dim),
+    )
+    return projector
+
+
+class AudioToSemantic(torch.nn.Module):
+    def __init__(self, embed_dim=768, sem_depth=1, sem_heads=1, sem_mlp_ratio=1):
+        # grid_size, pos_embed
+        super().__init__()
+        self.sem_token = torch.nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.sem_blocks = torch.nn.ModuleList([
+            timm.models.vision_transformer.Block(embed_dim, sem_heads, sem_mlp_ratio, qkv_bias=True, norm_layer=torch.nn.LayerNorm)
+            for i in range(sem_depth)])
+        self.norm = torch.nn.LayerNorm(embed_dim)
+        self.apply(self._init_weights)
+        self.dont_average = True
+
+    def _init_weights(self, m):
+        if isinstance(m, torch.nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, torch.nn.Linear) and m.bias is not None:
+                torch.nn.init.constant_(m.bias, 0)
+        elif isinstance(m, torch.nn.LayerNorm):
+            torch.nn.init.constant_(m.bias, 0)
+            torch.nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x):
+        # Append semantic token
+        sem_token = self.sem_token  # + self.pos_embed[:, :1, :]
+        sem_tokens = sem_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((sem_tokens, x), dim=1)
+
+        # Apply Transformer blocks
+        for blk in self.sem_blocks:
+            x = blk(x)
+        x = x[:, 0, :]  # Use semantic token only
+        x = self.norm(x)
+
+        return x
+
+
+class M2D_CLAP_ViT(M2D_X_ViT):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3,
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4., mlp_layer=Mlp, norm_layer=nn.LayerNorm, norm_pix_loss=False, norm_stats=None,
+                 loss_type='norm_mse', target_layers=None,
+                 loss_m2d=1.0, loss_off=0.01, off_emb_dim=768, sem_mode=0,
+                 world_rank=None, **kwargs):
+        super().__init__(img_size=img_size, patch_size=patch_size, in_chans=in_chans,
+                 embed_dim=embed_dim, depth=depth, num_heads=num_heads,
+                 decoder_embed_dim=decoder_embed_dim, decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
+                 mlp_ratio=mlp_ratio, mlp_layer=mlp_layer, norm_layer=norm_layer, norm_pix_loss=norm_pix_loss, norm_stats=norm_stats,
+                 loss_type=loss_type, target_layers=target_layers,
+                 loss_m2d=loss_m2d, loss_off=loss_off, off_emb_dim=off_emb_dim, world_rank=world_rank)
+        del self.offline_predictor
+
+        # CLIP
+        world_size, rank = [1, 0] if world_rank is None else world_rank
+        self.clip_loss_fn = ClipLoss(rank=rank, world_size=world_size)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        # Projectors
+        if sem_mode == 0:
+            self.audio_proj = get_MLP_projector(embed_dim, embed_dim, embed_dim)
+        elif sem_mode == 1:
+            self.audio_proj = AudioToSemantic(embed_dim=embed_dim, **{'sem_depth': 1, 'sem_heads': 1, 'sem_mlp_ratio': 1})
+        elif sem_mode == 2:
+            self.audio_proj = AudioToSemantic(embed_dim=embed_dim, **{'sem_depth': 2, 'sem_heads': 1, 'sem_mlp_ratio': 1})
+        else:
+            assert sem_mode in [0, 1, 2], f'Unknown sem_mode: {sem_mode}'
+        self.text_proj = nn.Identity() if off_emb_dim == embed_dim else nn.Linear(off_emb_dim, embed_dim)
+
+    def forward_off_loss(self, audio_embs, text_embs):
+        loss = self.clip_loss_fn(image_features=audio_embs, text_features=text_embs, logit_scale=self.logit_scale.exp())
+        return loss
+
+    def forward(self, imgs, text_embs, mask_ratio=0.7):
+        # online
+        latent, x_targ, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        pred_online = self.forward_decoder(latent, ids_restore)  # [N, targL, D]
+        with torch.no_grad():
+            target_online = self.forward_target_encoder(x_targ)
+        loss_online = self.forward_loss(target_online, pred_online, self.norm_pix_loss, self.loss_type)
+
+        # offline
+        audio_embs = self.drop_cls_token(latent)
+        if not hasattr(self.audio_proj, 'dont_average'):
+            audio_embs = audio_embs.mean(1)
+        audio_embs = self.audio_proj(audio_embs)
+
+        if text_embs is None:
+            loss_offline = torch.zeros_like(loss_online)
+        else:
+            text_embs = self.text_proj(text_embs)
+            loss_offline = self.forward_off_loss(text_embs, audio_embs)
+
+        loss = self.loss_m2d * loss_online + self.loss_off * loss_offline
+        return loss, pred_online, (ids_restore, mask, loss_online, loss_offline)
+
+    def clip_logit_scale(self):
+        # Thanks to https://github.com/openai/CLIP/issues/46#issuecomment-793434211
+        # logit scaling set as max 100
+        self.logit_scale.data = torch.clamp(self.logit_scale.data, 0, 4.6052)  # 4.6052 == log(100)
+
+
+def m2d_clap_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    model = M2D_CLAP_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def m2d_clap_vit_large(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):
+    model = M2D_CLAP_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+m2d_clap_vit_base_encoder_only = m2d_vit_base_encoder_only
+m2d_clap_vit_large_encoder_only = m2d_vit_large_encoder_only
+
+
+def m2dsg_clap_vit_base(patch_size=16, decoder_depth=8, in_chans=1, **kwargs):  # for audio
+    model = M2D_CLAP_ViT(
+        in_chans=in_chans, patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_layer=SwiGLU,
+        decoder_embed_dim=512, decoder_depth=decoder_depth, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+m2dsg_clap_vit_base_encoder_only = m2dsg_vit_base_encoder_only
--- _org/m2d/pos_embed.py	2023-07-07 10:22:30.217586685 +0900
+++ m2d/pos_embed.py	2023-07-07 10:22:30.229586731 +0900
@@ -23,12 +23,13 @@
     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)
+    gH, gW = grid_size
+    grid_h = np.arange(gH, dtype=np.float32)
+    grid_w = np.arange(gW, 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])
+    grid = grid.reshape([2, 1, gH, gW])
     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)
@@ -53,7 +54,7 @@
     out: (M, D)
     """
     assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float)
+    omega = np.arange(embed_dim // 2, dtype=float)
     omega /= embed_dim / 2.
     omega = 1. / 10000**omega  # (D/2,)
 
--- _org/m2d/timm_layers_pos_embed.py	2023-08-02 20:50:32.894880952 +0900
+++ m2d/timm_layers_pos_embed.py	2023-08-02 22:21:38.237502645 +0900
@@ -9,8 +9,6 @@
 import torch
 import torch.nn.functional as F
 
-from .helpers import to_2tuple
-
 _logger = logging.getLogger(__name__)
 
 
@@ -20,13 +18,13 @@
         old_size: Optional[List[int]] = None,
         num_prefix_tokens: int = 1,
         interpolation: str = 'bicubic',
-        antialias: bool = True,
         verbose: bool = False,
 ):
     # sort out sizes, assume square if old size not provided
     num_pos_tokens = posemb.shape[1]
     num_new_tokens = new_size[0] * new_size[1] + num_prefix_tokens
     if num_new_tokens == num_pos_tokens and new_size[0] == new_size[1]:
+        _logger.info(num_new_tokens, num_pos_tokens, new_size[0], new_size[1])
         return posemb
 
     if not old_size:
@@ -41,7 +39,7 @@
     # do the interpolation
     embed_dim = posemb.shape[-1]
     posemb = posemb.reshape(1, old_size[0], old_size[1], -1).permute(0, 3, 1, 2)
-    posemb = F.interpolate(posemb, size=new_size, mode=interpolation, antialias=antialias)
+    posemb = F.interpolate(posemb, size=new_size, mode=interpolation)  # removing "antialias=antialias" for backward compatibility.
     posemb = posemb.permute(0, 2, 3, 1).reshape(1, -1, embed_dim)
 
     # add back extra (class, etc) prefix tokens
--- _org/mae_train_audio.py	2023-07-07 10:22:30.225586716 +0900
+++ mae_train_audio.py	2026-02-20 15:37:38.549919216 +0900
@@ -15,69 +15,96 @@
 import os
 import time
 from pathlib import Path
+import subprocess
 
 import torch
 import torch.backends.cudnn as cudnn
 from torch.utils.tensorboard import SummaryWriter
-import torchvision.transforms as transforms
-import torchvision.datasets as datasets
+import matplotlib.pyplot as plt
 
 import timm
 
-assert timm.__version__ == "0.3.2"  # version check
 import timm.optim.optim_factory as optim_factory
 
 import util.misc as misc
 from util.misc import NativeScalerWithGradNormCount as NativeScaler
 
-import models_mae
+import m2d.models_mae
 
-from engine_pretrain import train_one_epoch
+from m2d.engine_pretrain_m2d import train_one_epoch_mae as train_one_epoch
+import audio_dataset
+import common
 
 
 def get_args_parser():
     parser = argparse.ArgumentParser('MAE pre-training', add_help=False)
-    parser.add_argument('--batch_size', default=64, type=int,
+    parser.add_argument('--batch_size', default=512, type=int,
                         help='Batch size per GPU (effective batch size is batch_size * accum_iter * # gpus')
-    parser.add_argument('--epochs', default=400, type=int)
+    parser.add_argument('--epochs', default=200, type=int)
     parser.add_argument('--accum_iter', default=1, type=int,
                         help='Accumulate gradient iterations (for increasing the effective batch size under memory constraints)')
+    parser.add_argument('--save_freq', default=50, type=int)
 
     # Model parameters
-    parser.add_argument('--model', default='mae_vit_large_patch16', type=str, metavar='MODEL',
+    parser.add_argument('--model', default='msm_mae_vit_base', type=str, metavar='MODEL',
                         help='Name of model to train')
+    parser.add_argument('--decoder_depth', type=int, default=8, metavar='DD',
+                        help='model decoder depth')
+                        
 
-    parser.add_argument('--input_size', default=224, type=int,
-                        help='images input size')
+    parser.add_argument('--input_size', default='80x208', type=str, help='images input size')
+    parser.add_argument('--patch_size', default='16x16', type=str, help='patch size')
 
     parser.add_argument('--mask_ratio', default=0.75, type=float,
                         help='Masking ratio (percentage of removed patches).')
+    parser.add_argument('--weighted_mask', action='store_true',
+                        help='Use weighted masking based on data contents.')
+    parser.add_argument('--var_weighted_mask', action='store_true',
+                        help='Use weighted masking (var) based on data contents.')
+    parser.add_argument('--det_weighted_mask', action='store_true',
+                        help='Use weighted masking -non probabilistic- based on data contents.')
+    parser.add_argument('--block_mask', action='store_true',
+                        help='Use block masking from data2vec 2.0.')
+    parser.add_argument('--obj_cent_mask', action='store_true',
+                        help='Use object centric masking based on data contents.')
+    parser.add_argument('--masking_hint_k', default=2, type=int)
+    parser.add_argument('--masking_start_epoch', default=50, type=int)
+    parser.add_argument('--abs_hint_ratio', default=-1.0, type=float)  # set value >= 0.0 to enable absolute hint ratio
 
-    parser.add_argument('--norm_pix_loss', action='store_true',
+    parser.add_argument('--no_norm_pix_loss', action='store_false', dest='norm_pix_loss',
                         help='Use (per-patch) normalized pixels as targets for computing loss')
-    parser.set_defaults(norm_pix_loss=False)
+    parser.set_defaults(norm_pix_loss=True)
+
+    parser.add_argument('--cont_mask', default=0, type=int,
+                        help='Use random 1-d (continuous) masking scheme. 0:off, 0<:# of frames to make them continuous.')
 
     # Optimizer parameters
     parser.add_argument('--weight_decay', type=float, default=0.05,
                         help='weight decay (default: 0.05)')
+    parser.add_argument('--clip_grad', type=float, default=3.0, metavar="NORM", ######
+                        help="Clip gradient norm (default: None, no clipping)")
 
     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',
+    parser.add_argument('--blr', type=float, default=3e-4, metavar='LR',
                         help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
     parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
                         help='lower lr bound for cyclic schedulers that hit 0')
 
-    parser.add_argument('--warmup_epochs', type=int, default=40, metavar='N',
+    parser.add_argument('--warmup_epochs', type=int, default=20, metavar='N',
                         help='epochs to warmup LR')
 
     # Dataset parameters
-    parser.add_argument('--data_path', default='/datasets01/imagenet_full_size/061417/', type=str,
+    parser.add_argument('--data_path', default='data', type=str,
                         help='dataset path')
+    parser.add_argument('--dataset', default='data/files_audioset.csv', type=str,
+                        help='dataset definition')
+    parser.add_argument('--norm_stats', default='None', type=str,
+                        help='dataset normalization stats')
 
-    parser.add_argument('--output_dir', default='./output_dir',
+    parser.add_argument('--output_dir', default='',
                         help='path where to save, empty for no saving')
-    parser.add_argument('--log_dir', default='./output_dir',
+    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')
@@ -104,6 +131,40 @@
     return parser
 
 
+def visualize_reconstruction(args, device, model, save_path):
+    ds, files = audio_dataset.build_viz_dataset(args)
+    if (ds is None) or (len(ds) == 0):
+        print(f'(Skipped visualization which require samples in {args.data_path}/vis_samples folder.)')
+        return
+    batch = torch.stack([ds[i] for i in range(len(ds))])
+    model.eval()
+    with torch.no_grad():
+        _, recons, _, masks = model.forward_viz(batch.to(device))
+    save_path.mkdir(parents=True, exist_ok=True)
+
+    grid_size = model.grid_size()
+    patch_size = model.patch_size()
+    masked = batch.clone().squeeze(1)  # [B, 1, H, W] -> [B, H, W]
+    for i, file in enumerate(files):
+        for y in range(grid_size[0]):
+            for x in range(grid_size[1]):
+                if not masks[i, y, x]: continue
+                y_start, y_end = y * patch_size[0], (y + 1) * patch_size[0]
+                x_start, x_end = x * patch_size[1], (x + 1) * patch_size[1]
+                masked[i, y_start:y_end, x_start:x_end] = 0.5 * torch.ones(patch_size, device=masked.device)
+
+    for i, file in enumerate(files):
+        # as .npy
+        np.save(f"{save_path}/recon_{Path(file).name}", recons[i].cpu().numpy())
+        # as .png
+        fig = plt.figure(figsize=[12, 8 if batch[0].shape[-1] < 310 else 6])
+        for j, img in enumerate([batch[i][0], recons[i][0], masked[i]]):
+            ax = fig.add_subplot(3, 1, j + 1)
+            ax.imshow(img.cpu().numpy(), origin='lower')
+        plt.margins(x=0, y=0)
+        fig.savefig(f'{save_path}/recon_{Path(file).stem}.png', bbox_inches = 'tight')
+
+
 def main(args):
     misc.init_distributed_mode(args)
 
@@ -119,13 +180,7 @@
 
     cudnn.benchmark = True
 
-    # simple augmentation
-    transform_train = transforms.Compose([
-            transforms.RandomResizedCrop(args.input_size, scale=(0.2, 1.0), interpolation=3),  # 3 is bicubic
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
-    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, 'train'), transform=transform_train)
+    dataset_train = audio_dataset.build_dataset(args)
     print(dataset_train)
 
     if True:  # args.distributed:
@@ -143,6 +198,8 @@
         log_writer = SummaryWriter(log_dir=args.log_dir)
     else:
         log_writer = None
+    common.PrintLogger(f'{args.log_dir}/console.txt')
+    print(args)
 
     data_loader_train = torch.utils.data.DataLoader(
         dataset_train, sampler=sampler_train,
@@ -153,12 +210,30 @@
     )
     
     # define the model
-    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
+    model = m2d.models_mae.__dict__[args.model](img_size=args.input_size, decoder_depth=args.decoder_depth,
+        norm_pix_loss=args.norm_pix_loss)
+    if args.weighted_mask:
+        assert (args.cont_mask == 0) and (not args.var_weighted_mask) and (not args.obj_cent_mask), "Masking options conflict."
+        model.set_random_weighted_mask_with_hint()
+    if args.var_weighted_mask:
+        assert (args.cont_mask == 0) and (not args.weighted_mask) and (not args.obj_cent_mask), "Masking options conflict."
+        model.set_random_weighted_mask_var_with_hint()
+    if args.obj_cent_mask:
+        assert (args.cont_mask == 0) and (not args.weighted_mask) and (not args.var_weighted_mask), "Masking options conflict."
+        if args.masking_start_epoch == 0:
+            print(f"⚠️  Warning: Too early to start object centric masking.")
+        model.set_object_centric_mask_with_hint()
+    if args.det_weighted_mask:
+        model.set_determ_weighted_mask_with_hint()
+    if args.block_mask:
+        model.set_data2vec2_block_masking()
+    if args.cont_mask > 0:
+        model.set_random_1d_mask(args.cont_mask)
 
     model.to(device)
 
     model_without_ddp = model
-    print("Model = %s" % str(model_without_ddp))
+    print("Model = %s" % common.short_model_desc(model_without_ddp))
 
     eff_batch_size = args.batch_size * args.accum_iter * misc.get_world_size()
     
@@ -176,7 +251,11 @@
         model_without_ddp = model.module
     
     # following timm: set wd as 0 for bias and norm layers
-    param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    try:
+        param_groups = optim_factory.add_weight_decay(model_without_ddp, args.weight_decay)
+    except:
+        print(' (for compatibility with timm) Switched add_weight_decay() to param_groups_weight_decay()')
+        param_groups = optim_factory.param_groups_weight_decay(model_without_ddp, args.weight_decay)
     optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
     print(optimizer)
     loss_scaler = NativeScaler()
@@ -186,18 +265,31 @@
     print(f"Start training for {args.epochs} epochs")
     start_time = time.time()
     for epoch in range(args.start_epoch, args.epochs):
+        epoch1 = epoch + 1
         if args.distributed:
             data_loader_train.sampler.set_epoch(epoch)
+        if args.weighted_mask or args.var_weighted_mask or args.det_weighted_mask or args.block_mask or args.obj_cent_mask:
+            if args.abs_hint_ratio < 0.0:
+                model_without_ddp.set_hint_ratio_by_epoch(cur_epoch=epoch, total_epochs=args.epochs, start_epoch=args.masking_start_epoch, k=args.masking_hint_k)
+            else:
+                model_without_ddp.set_hint_ratio(args.abs_hint_ratio)
         train_stats = train_one_epoch(
             model, data_loader_train,
             optimizer, device, epoch, loss_scaler,
             log_writer=log_writer,
             args=args
         )
-        if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
+        if args.output_dir and misc.is_main_process() and (epoch1 % args.save_freq == 0 or epoch1 == args.epochs):
             misc.save_model(
                 args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
-                loss_scaler=loss_scaler, epoch=epoch)
+                loss_scaler=loss_scaler, epoch=epoch1)
+            # visualize reconstructions
+            out_dir = Path(args.output_dir)/str(epoch)
+            visualize_reconstruction(args, device, model_without_ddp, out_dir)
+            # run the external evaluator
+            if epoch1 < args.epochs:
+                abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+                subprocess.Popen(['/bin/bash', './quick_eval.sh', abspath])
 
         log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                         'epoch': epoch,}
@@ -212,10 +304,47 @@
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))
     print('Training time {}'.format(total_time_str))
 
+    del model_without_ddp, model, data_loader_train, optimizer, loss_scaler
+    if misc.is_main_process():
+        abspath = Path(f'{args.output_dir}/checkpoint-{epoch1}.pth').absolute()
+        subprocess.call(['/bin/bash', './all_eval.sh', abspath])
+
+arg_conf_defaults = {
+   'dataset': ('data/files_audioset.csv', 'D', 'path'),
+    'ema_decay_init': (0.99995, 'ema', 'z'),
+    'ema_decay': (0.99999, 'ed', 'z'),
+    'decoder_depth': (8, 'dd', 'asis'),
+    'mask_ratio': (0.75, 'mr', 'z'),
+    'seed': (0, 's', 'asis'),
+    'norm_pix_loss':  (True, '~N', 'b'),
+    'loss_fn': ('norm_mse', 'L', 'head'),
+    'optim': ('adamw', 'O', 'asis'),
+    'blr': (3e-4, 'blr', 'z'),
+    'lr': (None, 'lr', 'z'),
+    'eff_batch_size': (2048, 'bs', 'asis'),
+    'accum_iter': (1, 'a', 'asis'),
+    'cont_mask': (0, 'C', 'asis'),
+    'weighted_mask': (False, 'Wm', 'b'),
+    'var_weighted_mask': (False, 'Vm', 'b'),
+    'det_weighted_mask': (False, 'Dm', 'b'),
+    'block_mask': (False, 'Bm', 'b'),
+    'obj_cent_mask': (False, 'Om', 'b'),
+    'masking_hint_k': [2, 'k', 'asis'],
+    'masking_start_epoch': [50, 'i', 'asis'],
+    'abs_hint_ratio': [-1.0, 'a', 'asis'],
+}
 
 if __name__ == '__main__':
     args = get_args_parser()
     args = args.parse_args()
-    if args.output_dir:
-        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    if not args.output_dir:
+        args.output_dir = f'{args.model}-{args.input_size}p{args.patch_size}'
+        args.output_dir += f'-{common.get_timestamp()[:6]}-{common.arg_conf_str(args, defaults=arg_conf_defaults)}'
+    if not args.log_dir:
+        args.log_dir = args.output_dir
+    args.input_size = [int(x) for x in args.input_size.split('x')]
+    args.patch_size = [int(x) for x in args.patch_size.split('x')]
+    args.norm_stats = eval(args.norm_stats) if args.norm_stats else None
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    print(args)
     main(args)