patch_m2d.diff

--- _org/train_audio.py	2023-07-07 10:22:30.225586716 +0900
+++ train_audio.py	2024-04-23 00:07:30.294216967 +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,105 @@
 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('--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 +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=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 +140,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 +148,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 +192,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
+        assert args.loss_off == 0.0, f'Missing --teacher while --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 +223,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 +233,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=off_emb_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, 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 +322,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_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'),
+    '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'),
+    '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	2024-05-13 08:17:23.051681712 +0900
@@ -1,3 +1,9 @@
+"""M2D-CLAP Pre-training Script
+
+M2D-CLAP: Masked Modeling Duo Meets CLAP for Learning General-purpose Audio-Language Representation
+https://TBD
+"""
+
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 
@@ -15,69 +21,204 @@
 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
+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 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
+
+
+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 = []
+        # 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):
+        emb_list = []
+        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)
+        # for debug
+        # if len(emb_list) > 1:
+        #     print(ytid, 'has multiple captions:', len(emb_list))
+        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
+        ytid = self.ds1.df.file_name.values[index].split('/')[-1][:11]
+        caption = self.get_random_caption(ytid)
+
+        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):
+    # get files and inflate the number of files (by repeating the list) if needed
+    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)
 
-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) + 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.')
 
-    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=.01, 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('--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('--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')
@@ -87,7 +228,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 +239,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 +247,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 +291,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 +309,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 +319,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, norm_stats=dataset_train.norm_stats,
+        off_emb_dim=dataset_train.caption_dim, clip_args=[global_rank, num_tasks])
+
+    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 +408,73 @@
                 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'),
+    '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_off': (0.01, '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/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	2024-06-04 18:11:30.006304208 +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_shceduler,
+                    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_shceduler(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,156 @@
             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_shceduler,
+                    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_shceduler(global_step)
+        if is_update_grad:
+            _model.update_target_network(ema_decay)
+            # Clip logit scale if a CLAP model own 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()))
 
     # 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	2024-06-03 09:47:54.848993793 +0900
@@ -9,14 +9,111 @@
 # 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
+import numpy as np
 
-from timm.models.vision_transformer import PatchEmbed, Block
+import timm
+from timm.models.vision_transformer import Block
 
-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
+
+
+def random_unstructured_mask(shape, mask_ratio, device):
+    B, F, T = shape # Batch, Freq bins, and Time frames; equivalent to Batch, Height, and Width for the image.
+    L = F * T
+    len_keep = int(L * (1 - mask_ratio))
+    noise = torch.rand(B, L, device=device)  # noise in [0, 1]
+    # sort noise for each sample
+    ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+    return ids_shuffle, len_keep
+
+
+def random_structured_mask(shape, mask_ratio, device):
+    """Random structured masking for training in audio tasks."""
+    B, F, T = shape
+
+    # We want true random freq/time masking but need to make the number of masks consistent among samples.
+    # We impose a constraint that the number of freq/time masks be consistent across samples while leaving it open where we mask.
+    NF = int(F * (mask_ratio + 1./F) * np.random.rand())
+    NF = min(F - 1, NF) # prevent to mask all freq. bins.
+    mask_ratio = max(mask_ratio + (.5/T) - (NF/F), 0.)
+    NT = int(T*mask_ratio)
+
+    # Make mask for each batch sample.
+    mask = torch.zeros((B, F, T), dtype=torch.int, device=device)
+    for b in range(B):
+        mask[b, torch.randperm(F)[:NF]] = 1
+    for b in range(B):
+        mask[b, :, torch.randperm(T)[:NT]] = 1
+
+    ids_shuffle = torch.argsort(mask.view(B, -1), descending=True)
+    len_keep = (mask[0] == 0).sum()
+    # print(len_keep, mask[:2])
+    return ids_shuffle, len_keep
+
+
+def make_one_1dmask(total_frames=800//2, mask_ratio=0.6, one_mask_frames=20):
+    M = int(total_frames * mask_ratio)
+    mask = np.zeros(total_frames, dtype=np.int8)
+    # Mask frames more than M, number of frames to mask.
+    while mask.sum() < M:
+        i = np.random.randint(low=0, high=total_frames)
+        mask[i:i+one_mask_frames] = 1
+    # Unmask frames from the tail to adjust total masked frames to M.
+    n_unmask = mask.sum() - M
+    if n_unmask > 0:
+        j_unmask = np.where(mask == 1)[0][-n_unmask:]
+        mask[j_unmask] = 0
+    return mask
+
+def random_1d_mask(shape, mask_ratio=0.6, device='cuda', one_mask_frames=20//2):
+    B, F, T = shape
+    mask = np.zeros((B, T), dtype=np.int8)
+    for i in range(B):
+        mask[i] = make_one_1dmask(total_frames=T, mask_ratio=mask_ratio, one_mask_frames=one_mask_frames)
+    mask = np.tile(mask, (1, F))
+    mask = torch.tensor(mask).to(device)
+
+    ids_shuffle = torch.argsort(mask.view(B, -1), dim=1)
+    len_keep = (mask[0] == 0).sum()
+    return ids_shuffle, len_keep
 
 
 class MaskedAutoencoderViT(nn.Module):
@@ -25,8 +122,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., 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 +137,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)
             for i in range(depth)])
         self.norm = norm_layer(embed_dim)
         # --------------------------------------------------------------------------
@@ -51,24 +151,54 @@
         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)
             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
+
+    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 +224,63 @@
 
     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 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
-        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)
+            ids_restore = torch.argsort(ids_shuffle, dim=1)
 
         # keep the first subset
         ids_keep = ids_shuffle[:, :len_keep]
@@ -147,15 +294,22 @@
 
         return x_masked, mask, ids_restore
 
-    def forward_encoder(self, x, mask_ratio):
+    def forward_encoder(self, x, mask_ratio, return_layers=False, adjust_short=False):
         # embed patches
         x = self.patch_embed(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)
 
         # append cls token
         cls_token = self.cls_token + self.pos_embed[:, :1, :]
@@ -163,12 +317,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 +370,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 +381,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 +395,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 +435,836 @@
 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 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., norm_layer=nn.LayerNorm, norm_pix_loss=False, norm_stats=None,
+                 loss_type='norm_mse', target_layers=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, 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
+
+        # --------------------------------------------------------------------------
+        # Target encoder specifics
+        self.target_blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_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):
+        """
+        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; TODO fix comment
+        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 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_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)
+        return loss, pred, ids_restore # mask
+
+
+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), 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 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; TODO fix comment
+        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
+
+
+# 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., 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, **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)
+        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
+from torch import distributed as dist
+
+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,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+    ):
+        super().__init__()
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def get_ground_truth(self, device, num_logits) -> torch.Tensor:
+        # calculated ground-truth and cache if enabled
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+        return labels
+
+    def get_logits(self, image_features, text_features, logit_scale):
+        image_features = torch.nn.functional.normalize(image_features, dim=-1, p=2)
+        text_features = torch.nn.functional.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):
+        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
+
+
+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., 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,
+                 clip_args=None):
+        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)
+        del self.offline_predictor
+
+        # CLIP
+        rank, world_size = [0, 1] if clip_args is None else clip_args
+        self.clip_loss_fn = ClipLoss(rank=rank, world_size=world_size)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        # Projectors
+        proj_hidden_dim = embed_dim
+        self.audio_proj = nn.Sequential(
+            nn.Linear(embed_dim, proj_hidden_dim),
+            nn.ReLU(),
+            nn.Linear(proj_hidden_dim, embed_dim),
+        )
+        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).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
--- _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	2024-03-25 19:07:49.243548273 +0900
@@ -15,40 +15,45 @@
 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).')
@@ -60,24 +65,30 @@
     # 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 +115,29 @@
     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)
+
+    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], masks[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 +153,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 +171,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 +183,13 @@
     )
     
     # 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)
 
     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 +207,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,6 +221,7 @@
     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)
         train_stats = train_one_epoch(
@@ -194,10 +230,17 @@
             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 +255,23 @@
     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])
+
 
 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)}'
+    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)