Initial commit

.gitignore

+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+.DS_Store

README.md

+# SAMix
+
+Official code for:  
+**"SAMix: Calibrated and Accurate Continual Learning via Sphere-Adaptive Mixup and Neural Collapse"**
+
+---
+
+## 🛠️ Prerequisites
+
+- **Python 3.8**  
+  Check installation:
+  ```bash
+  python3.8 --version
+  ```
+- **pip for Python 3.8**  
+  Check installation:
+  ```bash
+  python3.8 -m pip --version
+  ```
+- **wandb account**  
+  - This project uses **wandb** for convenient logging.  
+  - Create an account and log in before proceeding
+
+     ```bash
+         --wandb \
+         --wandb_project_name your_project_name \
+         --wandb_entity your_entity \
+    ```
+
+---
+
+## 📦 Environment Setup
+
+1. **Install virtualenv:**
+    ```bash
+    python3.8 -m pip install --upgrade pip
+    python3.8 -m pip install virtualenv
+    ```
+
+2. **Create and activate virtual environment:**
+    ```bash
+    # Create new environment
+    python3.8 -m virtualenv venv
+
+    # Activate environment
+    source venv/bin/activate  # On macOS/Linux
+    # or
+    .\venv\Scripts\activate  # On Windows
+    ```
+
+3. **Install required packages:**
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+4. **Install library for downloading the Tiny-ImageNet dataset:**
+    ```bash
+    pip uninstall googledrivedownloader
+    pip install assets/googledrivedownloader-0.4-py2.py3-none-any.whl
+    ```
+
+    **Note**: Due to the large size of the Tiny-ImageNet dataset, please download this dataset manually, upload it to your Google Drive, and update the script src/datasets/tiny_imagenet_dataset.py accordingly ("file_id='id_file_name'") to enable downloading and usage.
+
+> **Note**: To deactivate the virtual environment when you're done:
+> ```bash
+> deactivate
+> ```
+
+---
+
+## ⚙️ Configuration
+
+- **Set output folder:**  
+  Edit `config/base.yaml` and update `root_save_folder` to your desired path.
+
+- **wandb configuration:**  
+  In each bash script, set `wandb_project_name` and `wandb_entity` as needed.
+
+---
+
+## 🚀 Running Experiments
+
+### 4.1. CIFAR-10
+
+- **FC-NCCL with SAMix**
+    ```bash
+    MEM_SIZE=200 LOG_FOLDER_ZERO_MEM=not-use bash scripts/train/fc_nccl/fc_nccl_samix_cf10.sh
+    ```
+- **TA-NCCL with SAMix**
+    ```bash
+    MEM_SIZE=200 LOG_FOLDER_ZERO_MEM=not-use bash scripts/train/ta_nccl/ta_nccl_samix_cf10.sh
+    ```
+
+### 4.2. CIFAR-100
+
+- **FC-NCCL with SAMix**
+    ```bash
+    MEM_SIZE=200 LOG_FOLDER_ZERO_MEM=not-use bash scripts/train/fc_nccl/fc_nccl_samix_cf100.sh
+    ```
+- **TA-NCCL with SAMix**
+    ```bash
+    MEM_SIZE=200 LOG_FOLDER_ZERO_MEM=not-use bash scripts/train/ta_nccl/ta_nccl_samix_cf100.sh
+    ```
+
+### 4.3. Tiny-ImageNet
+
+- **FC-NCCL with SAMix**
+    ```bash
+    MEM_SIZE=200 LOG_FOLDER_ZERO_MEM=not-use bash scripts/train/fc_nccl/fc_nccl_samix_tiny_imagenet.sh
+    ```
+- **TA-NCCL with SAMix**
+    ```bash
+    MEM_SIZE=200 LOG_FOLDER_ZERO_MEM=not-use bash scripts/train/ta_nccl/ta_nccl_samix_tiny_imagenet.sh
+    ```
+
+> **Note 1:**  
+> Set `LOG_FOLDER_ZERO_MEM` to the log folder containing 200 auxiliary samples if `MEM_SIZE=0` (memory-free setting).
+
+> **Note 2:**  
+> **To run without SAMix:**
+> 1. Remove these arguments:
+>     ```bash
+>     --samix \
+>     --loss_normal_samples fnc2 \
+>     --weight_samix_loss 5.0 \
+>     ```
+> 2. Change `--loss nc_samix` to:
+>     ```bash
+>     --loss fnc2
+>     ```
+>     or
+>     ```bash
+>     --loss dr
+>     ```
+> 3. If using `--loss fnc2`, add `--focal_gamma` and `--main_mark`
+
+---
+
+## 🧪 Evaluation
+
+**Modify:** CHECKPOINT_FOLDER & LOG_FOLDER, and then run:
+
+### 5.1. CIFAR-10
+```bash
+CHECKPOINT_FOLDER=/path/check/point/folder LOG_FOLDER=/path/log/folder/ bash scripts/evaluate/eval_task_cf10.sh
+```
+
+### 5.2. CIFAR-100
+```bash
+CHECKPOINT_FOLDER=/path/check/point/folder LOG_FOLDER=/path/log/folder/ bash scripts/evaluate/eval_task_cf100.sh
+```
+
+### 5.3. Tiny-ImageNet
+```bash
+CHECKPOINT_FOLDER=/path/check/point/folder LOG_FOLDER=/path/log/folder/ bash scripts/evaluate/eval_task_tiny_imagenet.sh
+```
+
+---
+
+## 📊 Network Calibration
+
+**Modify bash file:** "scripts/calibration/run_calibration.sh", and then run:
+
+```bash
+    bash scripts/calibration/run_calibration.sh
+```
+
+This will compute and display:
+- Expected Calibration Error (ECE) each task
+- Overconfidence Error (OE) each task
+- Average Expected Calibration Error (AECE) 
+- Average Overconfidence Error (AOE)
+
+---
+
+## 📂 Project Structure
+
+```
+.
+├── assets/                             # Library for downloading Tiny-ImageNet dataset
+├── config/                             # Configuration files (e.g., base.yaml for experiment settings)
+├── scripts/                            # Training, evaluation, and calibration bash scripts
+├── src/                                # Main source code for SAMix and all modules
+│   ├── args/                           # Argument parsing and configuration utilities
+│   ├── backbones/                      # Backbone network architectures
+│   ├── buffers/                        # Buffer management for continual learning
+│   ├── datasets/                       # Dataset-related utilities and loaders
+│   ├── distillation/                   # Stability losses
+│   ├── plas_losses/                    # Plasticity losses
+│   ├── utils/                          # General utility functions and helper modules
+├── README.md                           # Project documentation
+├── main_compute_calibration.py         # Main file for computing calibration
+├── main_eval.py                        # Eval run file
+├── main_train.py                       # Train run file
+└── requirements.txt                    # Python dependencies
+```
+
+
+## 📁 Output Folder Structure
+
+```
+root_path: in 'root_save_folder' in config/base.yaml
+├── data
+├── experiments
+│   └── job_run_name-yyyy_mm_dd_HH_MM_SS
+│       ├── models
+│       │   ├── encoder_cosine[_warm]
+│       │   │   ├── task_{i}.pth
+│       │   │   ├── args_task_{i}.json
+│       │   ├── eval_linear_yyyy_mm_dd_HH_MM_SS
+│       │   │   │── accuracy_task_{i}.txt
+│       │   │   │── eval_args_task_{i}.json       
+│       │   │   │── linear_model_eval_linear_yyyy_mm_dd_HH_MM_SS_task_{i}.pth
+│       │────logs
+│       │   │── encoder_cosine[_warm]
+│       │   │   │── prototypes_dim={d}_k={n_cls}_seed={seed}.npy
+│       │   │   │── subset_indices_random_{i}.npy
+│       │   │   │── replay_indices_random_{i}.npy
+

config/base.yaml

+general:
+  save_path_config:
+      root_save_folder: "your/root/folder/"
+      data_folder: "{root_save_folder}/data"
+      experiments_folder: "{root_save_folder}/experiments"
+      wandb_folder: "{root_save_folder}/wandb"
+      run_exp_folder: "{experiments_folder}/{job_run_name}-{datetime}"
+  pattern_name:
+      model_task_name: "last_{policy}_{target_task}.pth"
+      replay_indices_task_name: "replay_indices_{policy}_{target_task}.npy"
+      subset_all_indices_task_name: "subset_indices_{policy}_{target_task}.npy"
+      set_prototypes_name: "prototypes_dim={d}_k={k}_seed={seed}.npy"
+      job_wandb_name_train: "TRAINING-{job_run_name}"
+      job_wandb_name_eval: "EVAL-{job_run_name}-task-{target_task}"
+      config_json_name_train: "train_args_task_{target_task}.json"
+      config_json_name_eval: "eval_args_task_{target_task}.json"
+      accuracy_single_task: "accuracy_task_{target_task}.txt"
+      eval_model_name: "linear_model_{eval_folder_name}_task_{target_task}.pth"
\ No newline at end of file

config/config_reader.py

+import yaml
+import os
+from src.utils.general_utils import print_x
+from yaml import Loader
+
+
+def read_config(file_path=None):
+    if not file_path:
+        current_working_dir = os.getcwd()
+        file_path = os.path.join(current_working_dir, "config/base.yaml")
+    print_x(f"{os.path.abspath(file_path)}")
+    with open(file_path, "r") as f:
+        cfg = yaml.load(f, Loader=Loader)
+        f.close()
+        return cfg

main_compute_calibration.py

+import torch
+from src.backbones.resnet_big import SupConResNet, LinearClassifier
+import torch.backends.cudnn as cudnn
+from src.utils.transforms_utils import get_transforms
+from torchvision import datasets
+from torch.utils.data import Dataset, Subset
+import numpy as np
+from src.utils.general_utils import print_x, get_feat_dim_projection
+import argparse
+from config.config_reader import read_config
+from src.datasets import DATA_INFORMATION
+
+device = torch.device("cpu")
+if torch.cuda.is_available():
+    torch.backends.cuda.max_split_size_mb = 128
+    device = torch.device("cuda")
+    print_x("CUDA available")
+else:
+    print_x("CPU available")
+
+
+def expected_calibration_and_overconfidence_error(predictions, labels, n_bins):
+    """
+    Compute Expected Calibration Error (ECE) and Overconfidence Error (OE).
+
+    Args:
+        predictions (torch.Tensor): Model outputs (logits or probabilities).
+        labels (torch.Tensor): Ground truth labels.
+        n_bins (int): Number of bins for confidence calibration.
+
+    Returns:
+        tuple: (ECE, OE)
+    """
+    probs = torch.softmax(predictions, dim=1)
+    confidences, predicted_classes = torch.max(probs, dim=1)
+    avg_confidence = confidences.mean()
+    accuracies = (predicted_classes == labels).float()
+    acc = accuracies.mean()
+    ece = 0.0
+    oe = 0.0
+    bin_boundaries = torch.linspace(0, 1, n_bins + 1)
+    bin_accs = torch.zeros(n_bins)
+    bin_confs = torch.zeros(n_bins)
+    bin_sizes = torch.zeros(n_bins)
+
+    for i in range(n_bins):
+        bin_lower, bin_upper = bin_boundaries[i], bin_boundaries[i + 1]
+        mask = (confidences > bin_lower) & (confidences <= bin_upper)
+        if mask.sum().item() > 0:
+            bin_sizes[i] = mask.sum().float()
+            bin_accs[i] = accuracies[mask].mean()
+            bin_confs[i] = confidences[mask].mean()
+            bin_weight = mask.float().mean()
+            ece += bin_weight * torch.abs(bin_accs[i] - bin_confs[i])
+            overconfidence = torch.clamp(bin_confs[i] - bin_accs[i], min=0)
+            oe += bin_weight * bin_confs[i] * overconfidence
+    return ece.item(), oe.item(), bin_accs.detach().numpy(), bin_confs.detach().numpy(), bin_sizes.detach().numpy(), avg_confidence, acc
+
+
+def load_classifier(path_classifier):
+    classifier = LinearClassifier(name="resnet18", num_classes=opt.n_cls)
+    ckpt_classifier = torch.load(path_classifier, map_location='cpu')
+    state_dict_classifier = ckpt_classifier['model']
+    classifier.load_state_dict(state_dict_classifier)
+    return classifier
+
+
+def load_backbone(path_checkpoint, predictor_hidden_dim):
+    feat_dim = get_feat_dim_projection(opt.dataset)
+    model = SupConResNet(name="resnet18", feat_dim=feat_dim, head='mlp', predictor_hidden_dim=predictor_hidden_dim)
+    ckpt = torch.load(path_checkpoint, map_location='cpu')
+    state_dict = ckpt['model']
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        k = k.replace("module.", "")
+        new_state_dict[k] = v
+    state_dict = new_state_dict
+    model = model.cpu()
+    cudnn.benchmark = True
+    model.load_state_dict(state_dict)
+    return model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Calibration computation for CL')
+    parser.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10', 'cifar100', 'tiny-imagenet'],
+                        help='dataset name (default: cifar10)')
+    parser.add_argument('--target_task', type=int, default=4,
+                        help='target task number (default: 4)')
+    parser.add_argument('--n_bins', type=int, default=10,
+                        help='number of bins to calculate calibration')
+    parser.add_argument('--checkpoint_path', type=str, required=True,
+                        help='path to model checkpoint')
+    parser.add_argument('--classifier_path', type=str, required=True,
+                        help='path to linear classifier checkpoint')
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    fixed_config = read_config()
+    save_path_config = fixed_config['general']['save_path_config']
+    root_path = save_path_config['root_save_folder']
+    data_folder = save_path_config['data_folder'].replace("{root_save_folder}", root_path)
+    opt = parse_args()
+    opt.n_cls, opt.cls_per_task, opt.size = DATA_INFORMATION[opt.dataset]
+
+    # data
+    train_transform, val_transform = get_transforms(opt)
+    all_target_classes = list(range(0, (opt.target_task + 1) * opt.cls_per_task))
+    subset_indices = []
+    _val_dataset = datasets.CIFAR100(root=data_folder, train=False, download=True, transform=val_transform)
+    for tc in all_target_classes:
+        subset_indices += np.where(np.array(_val_dataset.targets) == tc)[0].tolist()
+    val_dataset = Subset(_val_dataset, subset_indices)
+    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=len(val_dataset), shuffle=False)
+    data, labels = next(iter(val_loader))
+
+    images = data.float().to(device)
+    labels = labels.to(device)
+
+    num_tasks = opt.n_cls // opt.cls_per_task
+    backbone = load_backbone(path_checkpoint=opt.checkpoint_path,
+                             predictor_hidden_dim=get_feat_dim_projection(opt.dataset))
+    linear_classifier = load_classifier(path_classifier=opt.classifier_path)
+    if torch.cuda.is_available():
+        backbone = backbone.cuda()
+        linear_classifier = linear_classifier.cuda()
+    features = backbone.encoder(images)
+
+    sum_ece = 0.0
+    sum_oe = 0.0
+    for task_idx in range(num_tasks):
+        lower_bound = task_idx * opt.cls_per_task
+        upper_bound = (task_idx + 1) * opt.cls_per_task
+        mask = (labels >= lower_bound) & (labels < upper_bound)
+        filtered_features = features[mask]
+        filtered_labels = labels[mask]
+        output_logits = linear_classifier(filtered_features)
+        ECE, OE, accuracies, confidences, bin_sizes, avg_confidence, acc = expected_calibration_and_overconfidence_error(
+            predictions=output_logits, labels=filtered_labels, n_bins=opt.n_bins)
+        sum_ece += ECE
+        sum_oe += OE
+        print_x(f"Task-{task_idx + 1}: classes {lower_bound} to {upper_bound - 1}")
+        print_x(
+            f"ECE = {ECE:.4f}, OE = {OE:.8f}, avg_confidence = {avg_confidence:.4f}, acc = {acc:.4f}, \nbin_sizes = {bin_sizes},\naccuracies = {accuracies}, \nconfidences = {confidences}")
+        print_x(f"-----------------------")
+
+    print_x(f"Average across tasks: AECE = {(sum_ece / num_tasks):.4f}, AOE = {(sum_oe / num_tasks):.8f}")

main_eval.py

+'''
+    This is experimental code for evaluating a single task
+'''
+
+from __future__ import print_function
+import sys
+import time
+import numpy as np
+import copy
+import torch
+import torch.backends.cudnn as cudnn
+from src.datasets.evaluate_dataloader import set_loader
+from src.utils.general_utils import AverageMeter
+from src.utils.general_utils import adjust_learning_rate, warmup_learning_rate
+from src.utils.general_utils import set_optimizer
+from src.utils.metrics import TaskAccuracy
+from src.backbones.resnet_big import SupConResNet, LinearClassifier
+from config.config_reader import read_config
+from src.utils.general_utils import print_x, save_config_file, write_eval_acc, \
+    get_feat_dim_projection
+from src.args.setup_evaluating import parse_all_evaluating_args, parse_general_evaluating_args
+import wandb
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def set_model(opt):
+    feat_dim = get_feat_dim_projection(opt.dataset)
+    if hasattr(opt, "predictor_hidden_dim") and opt.predictor_hidden_dim:
+        model = SupConResNet(name=opt.model, feat_dim=feat_dim, head='mlp', predictor_hidden_dim=opt.predictor_hidden_dim)
+    else:
+        model = SupConResNet(name=opt.model, feat_dim=feat_dim, head='mlp')
+    criterion = torch.nn.CrossEntropyLoss()
+    classifier = LinearClassifier(name=opt.model, num_classes=opt.n_cls)
+    ckpt = torch.load(opt.ckpt, map_location='cpu')
+    state_dict = ckpt['model']
+    if torch.cuda.is_available():
+        if torch.cuda.device_count() > 1:
+            model.encoder = torch.nn.DataParallel(model.encoder)
+        else:
+            new_state_dict = {}
+            for k, v in state_dict.items():
+                k = k.replace("module.", "")
+                new_state_dict[k] = v
+            state_dict = new_state_dict
+        model = model.cuda()
+        classifier = classifier.cuda()
+        criterion = criterion.cuda()
+        cudnn.benchmark = True
+        model.load_state_dict(state_dict)
+    return model, classifier, criterion
+
+
+def train(train_loader, model, classifier, criterion, optimizer, epoch, opt):
+    """one epoch training"""
+    model.eval()
+    classifier.train()
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    losses = AverageMeter()
+    end = time.time()
+    acc = 0.0
+    cnt = 0.0
+    for idx, (images, labels) in enumerate(train_loader):
+        data_time.update(time.time() - end)
+
+        if torch.cuda.is_available():
+            images = images.cuda(non_blocking=True)
+            labels = labels.cuda(non_blocking=True)
+        bsz = labels.shape[0]
+
+        # warm-up learning rate
+        warmup_learning_rate(opt, epoch, idx, len(train_loader), optimizer)
+
+        # compute loss
+        with torch.no_grad():
+            features = model.encoder(images)
+        output = classifier(features.detach())
+        loss = criterion(output, labels)
+
+        # update metric
+        losses.update(loss.item(), bsz)
+        acc += (output.argmax(1) == labels).float().sum().item()
+        cnt += bsz
+
+        # SGD
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if opt.wandb:
+            wandb.log({"linear train/loss": losses.val,
+                       "linear train/acc_top1": acc / cnt * 100.})
+        if (idx + 1) % opt.print_freq == 0:
+            print_x('Train: Epoch:[{0}] batch [{1}/{2}]\t'
+                    'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                    'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
+                    'loss {loss.val:.3f} ({loss.avg:.3f})\t'
+                    'Acc@1 {top1:.3f}'.format(
+                epoch, idx + 1, len(train_loader), batch_time=batch_time,
+                data_time=data_time, loss=losses, top1=acc / cnt * 100.))
+            sys.stdout.flush()
+
+    return losses.avg, acc / cnt * 100.
+
+
+def validate(val_loader, model, classifier, criterion, epoch, opt):
+    """validation each epoch"""
+    model.eval()
+    classifier.eval()
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    corr = [0.] * (opt.target_task + 1) * opt.cls_per_task
+    cnt = [0.] * (opt.target_task + 1) * opt.cls_per_task
+    correct_task = 0.0
+
+    with torch.no_grad():
+        end = time.time()
+        for idx, (images, labels) in enumerate(val_loader):
+            images = images.float().to(device)
+            labels = labels.to(device)
+            batch_size = labels.shape[0]
+
+            # forward
+            output = classifier(model.encoder(images))
+            loss = criterion(output, labels)
+
+            # update metric
+            losses.update(loss.item(), batch_size)
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+            cls_list = np.unique(labels.cpu())
+            correct_all = (output.argmax(1) == labels)
+
+            for tc in cls_list:
+                mask = labels == tc
+                # task-il acc
+                check_tc = output[mask,
+                           (tc // opt.cls_per_task) * opt.cls_per_task: ((
+                                                                                 tc // opt.cls_per_task) + 1) * opt.cls_per_task]
+                correct_task += (check_tc.argmax(1) == (tc % opt.cls_per_task)).float().sum()
+
+                # class-il acc
+                corr[tc] += correct_all[mask].float().sum().item()
+                cnt[tc] += mask.float().sum().item()
+
+            if idx % opt.print_freq == 0:
+                print_x('Test: Epoch: [{0}] batch [{1}/{2}]\t'
+                        'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                        'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
+                        'Acc@1 {top1:.3f} {task_il:.3f}'.format(
+                    epoch, idx, len(val_loader), batch_time=batch_time,
+                    loss=losses, top1=np.sum(corr) / np.sum(cnt) * 100., task_il=correct_task / np.sum(cnt) * 100.))
+    if opt.wandb:
+        wandb.log({"val/loss": losses.val,
+                   "val/acc_top1": np.sum(corr) / np.sum(cnt) * 100.,
+                   "val/acc_task_il": correct_task / np.sum(cnt) * 100.})
+    print_x(
+        ' * Epoch {0} - Acc@1={top1:.3f} task_il={task_il:.3f}'.format(epoch,
+                                                                       top1=np.sum(corr) / np.sum(cnt) * 100.,
+                                                                       task_il=correct_task / np.sum(cnt) * 100.))
+    return losses.avg, corr, cnt, correct_task / np.sum(cnt) * 100.
+
+
+def evaluate_single_task(opt):
+    overall_acc_task = TaskAccuracy(opt.target_task, opt.epochs)
+    best_model = None
+
+    print_x(f"Args = {opt}")
+    save_config_file(opt, opt.eval_config_path)
+    if opt.target_task is not None:
+        if opt.target_task == 0:
+            replay_indices = np.array([])
+        else:
+            replay_indices = np.load(opt.logpt)
+        print_x(f"len(replay_indices) = {len(replay_indices)}")
+    # build data loader
+    train_loader, val_loader, cls_num_list = set_loader(opt, replay_indices)
+
+    # build model and criterion
+    model, classifier, criterion = set_model(opt)
+
+    # build optimizer
+    optimizer = set_optimizer(opt, classifier)
+
+    # training routine
+    for epoch in range(1, opt.epochs + 1):
+        print_x(f"Start eval epoch: {epoch}/{opt.epochs}")
+        adjust_learning_rate(opt, optimizer, epoch)
+
+        # train for one epoch
+        print_x("-------------Train classifier------")
+        print_x('Train: [EPOCH][idx in batch/number of batch]\t'
+                'BT batch_time.val (batch_time.avg)\t'
+                'DT data_time.val (data_time.avg)\t'
+                'loss loss.val (loss.avg)\t'
+                'Acc@1 top1_acc')
+        time1 = time.time()
+        loss, acc = train(train_loader, model, classifier, criterion, optimizer, epoch, opt)
+        print_x('Train epoch: {}, total time {:.2f}, loss: {:.2f}, accuracy: {:.2f}, lr: {:.3f}'.format(
+            epoch, time.time() - time1, loss, acc, optimizer.param_groups[0]['lr']))
+
+        # eval for one epoch
+        print_x("-------------Test classifier------")
+        print_x('Test: [batch_idx/total_batches]\t'
+                'Time batch_time.val (batch_time.avg)\t'
+                'Loss loss.val (loss.avg)\t'
+                'Acc@1 top1_acc task_il')
+        loss, val_corr, val_cnt, task_acc = validate(val_loader, model, classifier, criterion, epoch, opt)
+        val_acc = np.sum(val_corr) / np.sum(val_cnt) * 100.
+        if val_acc > overall_acc_task.best_acc[1]:
+            overall_acc_task.best_acc = (epoch, val_acc)
+            overall_acc_task.best_task_acc = (epoch, task_acc)
+            best_model = copy.deepcopy(classifier)
+            val_acc_stats = {}
+            for cls, (cr, c) in enumerate(zip(val_corr, val_cnt)):
+                if c > 0:
+                    val_acc_stats[str(cls)] = cr / c * 100.
+            overall_acc_task.dict_acc_classes = val_acc_stats
+
+        # wandb log - per epoch
+        if opt.wandb:
+            wandb.log({f"overall/train_acc": acc,
+                       f"overall/val_acc": val_acc,
+                       f"overall/task_acc": task_acc,
+                       f"overall/epoch": epoch})
+
+    # end of task
+    overall_acc_task.last_acc = val_acc
+    overall_acc_task.task_acc = task_acc
+    # wandb log
+    if opt.wandb:
+        wandb.run.summary[f"Task-{opt.target_task}: best_accuracy"] = overall_acc_task.best_acc[1]
+        wandb.run.summary[f"Task-{opt.target_task}: best_task_il_accuracy"] = overall_acc_task.best_task_acc[1]
+        wandb.run.summary[f"Task-{opt.target_task}: best_accuracy at epoch"] = overall_acc_task.best_acc[0]
+        wandb.run.summary[f"Task-{opt.target_task}: last_accuracy"] = val_acc
+        wandb.run.summary[f"Task-{opt.target_task}: last_task_il_accuracy"] = task_acc
+
+    # report
+    print_x(f"===> Report acc result into path: {opt.accuracy_single_task_path}")
+    write_eval_acc(opt.accuracy_single_task_path, overall_acc_task.best_acc[1], val_acc,
+                   overall_acc_task.dict_acc_classes, overall_acc_task.best_task_acc[1], task_acc)
+
+    # save model
+    print_x('==> Saving model into path:...' + opt.eval_model_path)
+    torch.save({
+        'opt': opt,
+        'model': best_model.state_dict(),
+        'optimizer': optimizer.state_dict()
+    }, opt.eval_model_path)
+    return overall_acc_task
+
+
+if __name__ == '__main__':
+    fixed_config = read_config()
+    opt = parse_general_evaluating_args()
+    opt = parse_all_evaluating_args(fixed_config, opt)
+    seed = opt.seed
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    if opt.wandb:
+        wandb.init(project=opt.wandb_project_name, name=opt.wandb_run_name_eval, entity=opt.wandb_entity,
+                   dir=opt.wandb_dir, config=vars(opt))
+
+    overall_acc_task = evaluate_single_task(opt=opt)

main_train.py

+from __future__ import print_function
+
+import os
+import copy
+import time
+import wandb
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+from src.utils.general_utils import (AverageMeter, adjust_learning_rate, warmup_learning_rate, print_x, set_optimizer,
+                                     save_model, load_model, set_replay_samples,
+                                     save_config_file)
+from src.backbones.resnet_big import SupConResNet
+from src.datasets.pretrain_dataloader import set_loader_with_replay
+from src.plas_losses import LOSSES, LIST_LOSSES_PROTOTYPES
+from config.config_reader import read_config
+from src.args.setup_training import parse_all_training_args
+from src.distillation.stab_loss import compute_distillation_loss
+from src.utils.mix_views_generation import generate_samix
+from src.utils.prototypes import generate
+
+device = torch.device("cpu")
+if torch.cuda.is_available():
+    torch.backends.cuda.max_split_size_mb = 128
+    device = torch.device("cuda")
+    print_x("CUDA available")
+else:
+    print_x("CPU available")
+
+
+def get_prototypes(criterion, opt):
+    n_proto_per_cls = opt.n_proto_per_cls if hasattr(opt, "n_proto_per_cls") else 1
+    current_prototypes = criterion.new_prototypes if hasattr(criterion, "new_prototypes") else criterion.points
+    indices_prototypes_target_classes = list(
+        range(0, (opt.target_task + 1) * opt.cls_per_task * n_proto_per_cls))
+    return current_prototypes[indices_prototypes_target_classes]
+
+
+def train(train_loader, model, frozen_model, criterion, optimizer, epoch, opt):
+    """one epoch training
+    """
+    model.train()
+    # Time
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    losses = AverageMeter()
+    distill = AverageMeter()
+    end = time.time()
+    for idx, (images, labels) in enumerate(train_loader):
+        # images: [2 x bsz, C, H, W]
+        data_time.update(time.time() - end)
+        images = torch.cat([images[0], images[1]], dim=0)
+        _views, lamb_ratio, mixed_indices = None, None, None
+        if hasattr(opt, "samix") and opt.samix:
+            _views, lamb_ratio, mixed_indices = generate_samix(data=images)
+            if torch.cuda.is_available():
+                _views = _views.cuda(non_blocking=True)
+                mixed_indices = mixed_indices.cuda(non_blocking=True)
+        if torch.cuda.is_available():
+            images = images.cuda(non_blocking=True)
+            labels = labels.cuda(non_blocking=True)
+        bsz = labels.shape[0]
+        # warm-up lr
+        warmup_learning_rate(opt, epoch, idx, len(train_loader), optimizer)
+        distill_predictor_output = None
+        if (opt.target_task > 0 and opt.distillation != "no-distillation" and hasattr(opt,
+                                                                                      "predictor_hidden_dim") and opt.predictor_hidden_dim and opt.predictor_hidden_dim != 0):
+            features, _encoded, distill_predictor_output = model(images, return_feat=True, distill=True)
+        else:
+            features, _encoded = model(images, return_feat=True)
+        row_original_size = features.size(0)
+        # samix
+        samix_features = None  # [2 x bsz, d]
+        samix_distill_predictor_output = None
+        if _views is not None:
+            if opt.target_task > 0 and opt.distillation != "no-distillation" and hasattr(opt,
+                                                                                         "predictor_hidden_dim") and opt.predictor_hidden_dim and opt.predictor_hidden_dim != 0:
+                samix_features, samix_encoded, samix_distill_predictor_output = model(_views, return_feat=True,
+                                                                                      distill=True)
+            else:
+                samix_features, samix_encoded = model(_views, return_feat=True)
+        f1, f2 = torch.split(features, [bsz, bsz], dim=0)
+        features = torch.cat([f1.unsqueeze(1), f2.unsqueeze(1)], dim=1)  # [bsz, 2, d]
+        target_labels = list(range(opt.target_task * opt.cls_per_task, (opt.target_task + 1) * opt.cls_per_task))
+        loss = criterion(features=features, labels=labels, target_labels=target_labels,
+                         samix_features=samix_features, lamb_ratio=lamb_ratio, mixed_indices=mixed_indices)
+        if opt.wandb:
+            wandb.log({"general_training/plasticity_loss": loss})
+        if opt.target_task > 0 and opt.distillation != "no-distillation":
+            if distill_predictor_output is not None:
+                features1_prev_task = distill_predictor_output
+            else:
+                features1_prev_task = features.transpose(0, 1).reshape(row_original_size, -1)
+            alpha_balance_distillation = None  # increase
+            if opt.distillation == 'hsd':
+                if hasattr(opt, "main_mark") and hasattr(opt, "second_mark"):
+                    if opt.main_mark is not None and opt.second_mark is not None:
+                        assert 0 < opt.main_mark < opt.epochs - 1 and opt.main_mark < opt.second_mark <= opt.epochs
+                        if epoch < opt.main_mark:
+                            alpha_balance_distillation = 0
+                        elif opt.main_mark <= epoch < opt.second_mark:
+                            alpha_balance_distillation = (epoch - opt.main_mark) / opt.epochs
+                        else:
+                            alpha_balance_distillation = (opt.second_mark - opt.main_mark) / opt.epochs
+                    else:
+                        alpha_balance_distillation = (epoch - 1) / opt.epochs
+                else:
+                    alpha_balance_distillation = (epoch - 1) / opt.epochs
+                assert alpha_balance_distillation >= 0
+            if _views is not None:
+                if samix_distill_predictor_output is not None:
+                    features1_prev_task = torch.cat((features1_prev_task, samix_distill_predictor_output), dim=0)
+                else:
+                    features1_prev_task = torch.cat((features1_prev_task, samix_features), dim=0)
+                images = torch.cat((images, _views), dim=0)
+            if opt.filter_distillation_point:
+                _distill = compute_distillation_loss(opt, features1_prev_task, frozen_model, images,
+                                                     set_prototypes=get_prototypes(criterion, opt),
+                                                     alpha_balance_distillation=alpha_balance_distillation)
+            else:  # no filter proto
+                _distill = compute_distillation_loss(opt, features1_prev_task, frozen_model, images,
+                                                     set_prototypes=None)
+            loss += opt.distill_power * _distill
+            distill.update(_distill.item(), bsz)
+        losses.update(loss.item(), bsz)
+        if opt.wandb:
+            wandb.log({"general_training/overall_loss": losses.val})
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if (idx + 1) % opt.print_freq == 0 or idx + 1 == len(train_loader):
+            print_x('Train: Epoch: [{0}] batch [{1}/{2}]\t'
+                    'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                    'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
+                    'loss {loss.val:.3f} ({loss.avg:.3f} {distill.avg:.3f})\t'
+                    'lr {lr: .3f}'.format(
+                epoch, idx + 1, len(train_loader), batch_time=batch_time,
+                data_time=data_time, loss=losses, distill=distill, lr=optimizer.param_groups[0]['lr']))
+    return losses.avg, frozen_model, distill.avg
+
+
+def set_model(opt, set_prototypes=None):
+    if hasattr(opt, "predictor_hidden_dim") and opt.predictor_hidden_dim:
+        model = SupConResNet(name=opt.model, feat_dim=opt.dim,
+                             head='mlp',
+                             predictor_hidden_dim=opt.predictor_hidden_dim)
+    else:
+        model = SupConResNet(name=opt.model, feat_dim=opt.dim,
+                             head='mlp')
+    LossClass = LOSSES[opt.loss]
+    criterion = LossClass(**opt.__dict__)
+    if torch.cuda.is_available():
+        if torch.cuda.device_count() > 1:
+            model.encoder = torch.nn.DataParallel(model.encoder)
+        model = model.cuda()
+        criterion = criterion.cuda()
+        cudnn.benchmark = True
+    if set_prototypes is not None and opt.loss in LIST_LOSSES_PROTOTYPES:
+        criterion.set_set_prototypes(set_prototypes)
+    return model, criterion
+
+
+def main():
+    fixed_config = read_config()
+    pattern_name = fixed_config['general']['pattern_name']
+    pattern_model_task_name = pattern_name['model_task_name']
+    pattern_replay_indices_task_name = pattern_name['replay_indices_task_name']
+    pattern_subset_all_indices_task_name = pattern_name['subset_all_indices_task_name']
+    pattern_set_prototypes_name = pattern_name['set_prototypes_name']
+    config_json_name_train_name = pattern_name['config_json_name_train']
+    opt = parse_all_training_args(fixed_config)
+    seed = opt.seed
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    print_x(opt)
+    opt.path_save_prototypes = None
+    set_prototypes = None
+    if opt.loss in LIST_LOSSES_PROTOTYPES:
+        path_save_prototypes = os.path.join(opt.log_folder, pattern_set_prototypes_name
+                                            .format(d=opt.dim, k=opt.n_cls, seed=seed))
+        opt.path_save_prototypes = path_save_prototypes
+        set_prototypes = generate(d=opt.dim, k=opt.n_cls, path_save_prototype=path_save_prototypes, seed=seed)
+        set_prototypes = torch.from_numpy(set_prototypes)
+    if opt.wandb:
+        wandb.init(project=opt.wandb_project_name, name=opt.wandb_run_name_train, entity=opt.wandb_entity,
+                   dir=opt.wandb_dir, config=vars(opt))
+
+    # model, optimizer
+    model, criterion = set_model(opt, set_prototypes)
+    frozen_model, _ = set_model(opt)
+    frozen_model.eval()
+    optimizer = set_optimizer(opt, model)
+
+    # resume running from a specific checkpoint
+    replay_indices = None
+    if opt.resume_target_task is not None:
+        print_x(f"Resume running after task {opt.resume_target_task}...")
+        assert (opt.path_resume_model_folder is not None) or (
+                opt.path_model_resume_task is not None and opt.path_replay_indices is not None)
+        if opt.path_resume_model_folder is not None:
+            print_x(f"Resume training, reuse old log folder")
+            load_file = os.path.join(opt.save_folder, pattern_model_task_name.format(policy=opt.replay_policy,
+                                                                                     target_task=opt.resume_target_task))
+            model, optimizer = load_model(model, optimizer, load_file)
+            path_resume_replay = os.path.join(opt.log_folder,
+                                              pattern_replay_indices_task_name.format(policy=opt.replay_policy,
+                                                                                      target_task=opt.resume_target_task))
+            replay_indices = np.load(path_resume_replay).tolist()
+        elif opt.path_model_resume_task is not None and opt.path_replay_indices is not None:
+            print_x(
+                f"Resume training, create new log folder")
+            model, optimizer = load_model(model, optimizer, opt.path_model_resume_task)
+            replay_indices = np.load(opt.path_replay_indices).tolist()
+
+    # separate tasks
+    original_epochs = opt.epochs
+    if opt.end_task is not None:  # if opt.end_task < T
+        if opt.resume_target_task is not None:
+            assert opt.end_task > opt.resume_target_task
+        opt.end_task = min(opt.end_task + 1, opt.n_cls // opt.cls_per_task)
+    else:  # end_task = T
+        opt.end_task = opt.n_cls // opt.cls_per_task
+
+    # Train each task
+    for target_task in range(0 if opt.resume_target_task is None else opt.resume_target_task + 1, opt.end_task):
+        print_x('>>> Start Training current task {}'.format(target_task))
+        opt.target_task = target_task
+        frozen_model = copy.deepcopy(model)
+
+        path_config_json = os.path.join(
+            opt.save_folder, config_json_name_train_name.replace("{target_task}", str(target_task)))
+        print_x(f"Saving args task: {target_task} into path: {path_config_json}")
+        save_config_file(opt, path_config_json)
+
+        replay_indices = set_replay_samples(opt, model, prev_indices=replay_indices)
+        np.save(os.path.join(opt.log_folder, pattern_replay_indices_task_name
+                             .format(policy=opt.replay_policy, target_task=target_task)), np.array(replay_indices))
+        train_loader, subset_indices = set_loader_with_replay(opt, replay_indices)
+
+        np.save(os.path.join(opt.log_folder, pattern_subset_all_indices_task_name
+                             .format(policy=opt.replay_policy, target_task=target_task)), np.array(subset_indices))
+
+        if target_task == 0 and opt.start_epoch is not None:
+            opt.epochs = opt.start_epoch
+        else:
+            opt.epochs = original_epochs
+
+        for epoch in range(1, opt.epochs + 1):
+            adjust_learning_rate(opt, optimizer, epoch)
+            time1 = time.time()
+            loss, frozen_model, distill = train(train_loader, model, frozen_model, criterion, optimizer, epoch, opt)
+            time2 = time.time()
+            print_x('epoch {}, total time {:.2f}'.format(epoch, time2 - time1))
+            if opt.wandb:
+                wandb.log({f"task_{target_task}/loss": loss,
+                           f"task_{target_task}/distill_loss": distill,
+                           f"task_{target_task}/learning_rate": optimizer.param_groups[0]['lr'],
+                           f"task_{target_task}/epoch": epoch})
+        save_file = os.path.join(
+            opt.save_folder,
+            pattern_model_task_name.format(policy=opt.replay_policy, target_task=target_task))
+        save_model(model, optimizer, opt, opt.epochs, save_file)
+    print_x(f"Finish training...")
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

+absl-py==2.1.0
+astunparse==1.6.3
+cachetools==5.3.3
+certifi==2024.2.2
+charset-normalizer==3.3.2
+click==8.1.7
+dataclasses==0.6
+dm-tree==0.1.8
+docker-pycreds==0.4.0
+flatbuffers==1.12
+fsspec==2024.10.0
+gast==0.4.0
+gitdb==4.0.11
+GitPython==3.1.43
+google-auth==2.29.0
+google-auth-oauthlib==0.4.6
+google-pasta==0.2.0
+googledrivedownloader @ file:///home/dang/code/continual_learning_framework/assets/googledrivedownloader-0.4-py2.py3-none-any.whl#sha256=278dc5fa6ad7027786d648977169a20ab3ba0b0985fd33965443e8585874ed87
+grpcio==1.64.0
+h5py==3.11.0
+huggingface-hub==0.26.1
+idna==3.7
+importlib_metadata==7.1.0
+joblib==1.4.2
+keras==2.9.0
+Keras-Preprocessing==1.1.2
+kornia==0.7.2
+kornia_rs==0.1.3
+libclang==18.1.1
+Markdown==3.6
+MarkupSafe==2.1.5
+mkl-fft @ file:///croot/mkl_fft_1695058164594/work
+mkl-random @ file:///croot/mkl_random_1695059800811/work
+mkl-service==2.4.0
+numpy==1.24.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.20.5
+nvidia-nvjitlink-cu12==12.3.101
+nvidia-nvtx-cu12==12.1.105
+oauthlib==3.2.2
+opencv-python==4.9.0.80
+opt-einsum==3.3.0
+packaging==24.1
+pandas==2.0.3
+Pillow==9.3.0
+platformdirs==4.2.2
+protobuf==3.19.6
+psutil==5.9.8
+pyaml==24.4.0
+pyasn1==0.6.0
+pyasn1_modules==0.4.0
+python-dateutil==2.9.0.post0
+pytz==2024.1
+PyYAML==6.0.1
+regex==2024.9.11
+requests==2.32.3
+requests-oauthlib==2.0.0
+rsa==4.9
+safetensors==0.4.5
+scikit-learn==1.3.2
+scipy==1.10.1
+sentry-sdk==2.2.1
+setproctitle==1.3.3
+six==1.16.0
+smmap==5.0.1
+termcolor==2.4.0
+threadpoolctl==3.1.0
+tiktoken==0.7.0
+tokenizers==0.20.1
+torch==1.12.0+cu102
+torchaudio==0.12.1+cu102
+torchvision==0.13.1+cu102
+transformers==4.45.2
+triton==2.3.0
+typing_extensions @ file:///croot/typing_extensions_1715268824938/work
+tzdata==2024.1
+umap-learn==0.1.5
+urllib3==2.2.1
+wandb==0.17.0
+Werkzeug==3.0.3
+wrapt==1.16.0
+zipp==3.18.2

scripts/calibration/run_calibration.sh

+#!/bin/bash
+# Run calibration computation
+python main_compute_calibration.py \
+    --dataset cifar10 \
+    --target_task 4 \
+    --n_bins 10 \
+    --checkpoint_path \path\checkpoint_name.pth \
+    --classifier_path \path\classifier.pth
\ No newline at end of file

scripts/evaluate/eval_task_cf10.sh

+python3 -W ignore main_eval.py \
+  --seed 1234 \
+  --target_task 4 \
+  --dataset cifar10 \
+  --learning_rate 0.1 \
+  --batch_size 256 \
+  --epochs 100 \
+  --num_workers 16 \
+  --model resnet18 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name LINEAR-eval-cifar10  \
+  --ckpt ${CHECKPOINT_FOLDER} \
+  --logpt ${LOG_FOLDER}
\ No newline at end of file

scripts/evaluate/eval_task_cf100.sh

+python3 -W ignore main_eval.py \
+  --seed 1234 \
+  --target_task 4 \
+  --dataset cifar100 \
+  --learning_rate 0.1 \
+  --batch_size 256 \
+  --epochs 100 \
+  --num_workers 16 \
+  --model resnet18 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name LINEAR-eval-cifar100  \
+  --ckpt ${CHECKPOINT_FOLDER} \
+  --logpt ${LOG_FOLDER}
\ No newline at end of file

scripts/evaluate/eval_task_tiny_imagenet.sh

+python3 -W ignore main_eval.py \
+  --seed 1234 \
+  --target_task 9 \
+  --dataset tiny-imagenet \
+  --learning_rate 0.1 \
+  --batch_size 256 \
+  --epochs 100 \
+  --num_workers 16 \
+  --model resnet18 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name LINEAR-eval-tiny-imagenet  \
+  --ckpt ${CHECKPOINT_FOLDER} \
+  --logpt ${LOG_FOLDER}
\ No newline at end of file

scripts/train/fc_nccl/fc_nccl_samix_cf10.sh

+python3 -W ignore main_train.py \
+  --seed 1234 \
+  --batch_size 512 \
+  --model resnet18 \
+  --dataset cifar10 \
+  --mem_size $MEM_SIZE \
+  --epochs 100 \
+  --start_epoch 500 \
+  --learning_rate 0.5 \
+  --temperature 0.5 \
+  --current_temp 0.2 \
+  --past_temp 0.01  \
+  --current_temp_sprd 0.2  \
+  --past_temp_sprd 0.1  \
+  --cosine \
+  --num_workers 16 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name fcnccl-samix-cf10-m-${MEM_SIZE} \
+  --distillation hsd \
+  --log_folder_zero_mem ${LOG_FOLDER_ZERO_MEM} \
+  --loss nc_samix \
+  --loss_normal_samples fnc2 \
+  --weight_samix_loss 5.0 \
+  --filter_distillation_point \
+  --samix \
+  --predictor_hidden_dim 128 \
+  --main_mark 30 \
+  --focal_gamma 1
\ No newline at end of file

scripts/train/fc_nccl/fc_nccl_samix_cf100.sh

+python3 -W ignore main_train.py \
+  --seed 1234 \
+  --batch_size 512 \
+  --model resnet18 \
+  --dataset cifar100 \
+  --mem_size $MEM_SIZE \
+  --epochs 100 \
+  --start_epoch 500 \
+  --learning_rate 0.5 \
+  --temperature 0.5 \
+  --current_temp 0.2 \
+  --past_temp 0.01  \
+  --current_temp_sprd 0.2  \
+  --past_temp_sprd 0.1  \
+  --cosine \
+  --num_workers 16 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name fcnccl-samix-cf100-m-${MEM_SIZE} \
+  --distillation hsd \
+  --log_folder_zero_mem ${LOG_FOLDER_ZERO_MEM} \
+  --loss nc_samix \
+  --loss_normal_samples fnc2 \
+  --weight_samix_loss 5.0 \
+  --filter_distillation_point \
+  --samix \
+  --predictor_hidden_dim 128 \
+  --main_mark 30 \
+  --focal_gamma 4
\ No newline at end of file

scripts/train/fc_nccl/fc_nccl_samix_tiny_imagenet.sh

+python3 -W ignore main_train.py \
+  --seed 1234 \
+  --batch_size 512 \
+  --model resnet18 \
+  --dataset tiny-imagenet \
+  --mem_size $MEM_SIZE \
+  --epochs 50 \
+  --start_epoch 500 \
+  --learning_rate 0.1 \
+  --temperature 0.5 \
+  --current_temp 0.1 \
+  --past_temp 0.1  \
+  --current_temp_sprd 0.2  \
+  --past_temp_sprd 0.1  \
+  --cosine \
+  --num_workers 16 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name fcnccl-samix-tiny-imagenet-m-${MEM_SIZE} \
+  --distillation hsd \
+  --log_folder_zero_mem ${LOG_FOLDER_ZERO_MEM} \
+  --loss nc_samix \
+  --loss_normal_samples fnc2 \
+  --weight_samix_loss 5.0 \
+  --filter_distillation_point \
+  --samix \
+  --predictor_hidden_dim 256 \
+  --main_mark 20 \
+  --focal_gamma 4
\ No newline at end of file

scripts/train/ta_nccl/ta_nccl_samix_cf10.sh

+python3 -W ignore main_train.py \
+  --seed 1234 \
+  --batch_size 512 \
+  --model resnet18 \
+  --dataset cifar10 \
+  --mem_size $MEM_SIZE \
+  --epochs 100 \
+  --start_epoch 500 \
+  --learning_rate 0.5 \
+  --temperature 0.5 \
+  --current_temp 0.2 \
+  --past_temp 0.01  \
+  --current_temp_sprd 0.2  \
+  --past_temp_sprd 0.1  \
+  --cosine \
+  --num_workers 16 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name tanccl-samix-cf00-m-${MEM_SIZE} \
+  --distillation hsd \
+  --log_folder_zero_mem ${LOG_FOLDER_ZERO_MEM} \
+  --loss nc_samix \
+  --loss_normal_samples dr \
+  --weight_samix_loss 5.0 \
+  --filter_distillation_point \
+  --samix \
+  --predictor_hidden_dim 128

scripts/train/ta_nccl/ta_nccl_samix_cf100.sh

+python3 -W ignore main_train.py \
+  --seed 1234 \
+  --batch_size 512 \
+  --model resnet18 \
+  --dataset cifar100 \
+  --mem_size $MEM_SIZE \
+  --epochs 100 \
+  --start_epoch 500 \
+  --learning_rate 0.5 \
+  --temperature 0.5 \
+  --current_temp 0.2 \
+  --past_temp 0.01  \
+  --current_temp_sprd 0.2  \
+  --past_temp_sprd 0.1  \
+  --cosine \
+  --num_workers 16 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name tanccl-samix-cf100-m-${MEM_SIZE} \
+  --distillation hsd \
+  --log_folder_zero_mem ${LOG_FOLDER_ZERO_MEM} \
+  --loss nc_samix \
+  --loss_normal_samples dr \
+  --weight_samix_loss 5.0 \
+  --filter_distillation_point \
+  --samix \
+  --predictor_hidden_dim 128

scripts/train/ta_nccl/ta_nccl_samix_tiny_imagenet.sh

+python3 -W ignore main_train.py \
+  --seed 1234 \
+  --batch_size 512 \
+  --model resnet18 \
+  --dataset tiny-imagenet \
+  --mem_size $MEM_SIZE \
+  --epochs 50 \
+  --start_epoch 500 \
+  --learning_rate 0.1 \
+  --temperature 0.5 \
+  --current_temp 0.1 \
+  --past_temp 0.1  \
+  --current_temp_sprd 0.2  \
+  --past_temp_sprd 0.1  \
+  --cosine \
+  --num_workers 16 \
+  --wandb \
+  --wandb_project_name your_project_name \
+  --wandb_entity your_entity \
+  --job_run_name tanccl-samix-tiny-imagenet-m-${MEM_SIZE} \
+  --distillation hsd \
+  --log_folder_zero_mem ${LOG_FOLDER_ZERO_MEM} \
+  --loss nc_samix \
+  --loss_normal_samples dr \
+  --weight_samix_loss 5.0 \
+  --filter_distillation_point \
+  --samix \
+  --predictor_hidden_dim 256

src/args/continual.py

+from argparse import ArgumentParser
+
+
+def continual_args(parser: ArgumentParser):
+    """Adds continual learning arguments to a parser.
+
+    Args:
+        parser (ArgumentParser): parser to add dataset args to.
+    """
+    parser.add_argument('--target_task', type=int, default=0)
+    parser.add_argument('--end_task', type=int, default=None)
+    parser.add_argument('--cls_per_task', type=int, default=2, help='Use for class incremental learning')
+    SPLIT_STRATEGIES = ["class", "data", "domain"]
+    parser.add_argument("--split_strategy", choices=SPLIT_STRATEGIES, type=str, default='class')
+    # distillation args
+    parser.add_argument("--distiller", type=str, default=None)
+    parser.add_argument('--distill_power', type=float, default=1.0)
+    parser.add_argument('--predictor_hidden_dim', type=int, default=None)

src/args/dataset.py

+from argparse import ArgumentParser
+from pathlib import Path
+
+
+def dataset_args(parser: ArgumentParser):
+    """Adds dataset-related arguments to a parser.
+
+    Args:
+        parser (ArgumentParser): parser to add dataset args to.
+    """
+
+    SUPPORTED_DATASETS = [
+        "cifar10",
+        "cifar100",
+        "tiny-imagenet",
+        "custom",
+    ]
+    parser.add_argument('--dataset', type=str, default='cifar10',
+                        choices=SUPPORTED_DATASETS, required=True, help='dataset')
+    parser.add_argument('--data_folder', type=str, default=None, help='path to custom dataset')
+    parser.add_argument("--train_dir", type=Path, default=None, help='train_dir use for custom dataset')
+    parser.add_argument("--val_dir", type=Path, default=None, help='val_dir use for custom dataset')

src/args/evaluate.py

+from argparse import ArgumentParser
+from src.args.utils import general_args
+
+
+def evaluate_args(parser: ArgumentParser):
+    """Adds all arguments to a parser.
+    Args:
+        parser (ArgumentParser): parser to add dataset args to.
+    """
+    general_args(parser)
+    parser.add_argument('--target_task', type=int, default=0, help='Use all classes if None else learned tasks so far')
+
+    # saving models frequency settings
+    parser.add_argument('--save_freq', type=int, default=50, help='save frequency')
+    parser.add_argument('--epochs', type=int, default=100, help='number of training epochs')
+
+    # optimization
+    parser.add_argument('--learning_rate', type=float, default=0.1, help='learning rate')
+    parser.add_argument('--lr_decay_epochs', type=str, default='60,75,90', help='where to decay lr, can be a list')
+    parser.add_argument('--lr_decay_rate', type=float, default=0.2, help='decay rate for learning rate')
+    parser.add_argument('--weight_decay', type=float, default=0, help='weight decay')
+
+    # load model path - ckpt & log folder
+    parser.add_argument('--ckpt', type=str, default='', help='path to pre-trained model')
+    parser.add_argument('--logpt', type=str, default='', help='path to pre-trained model')
+
+    # If using all train dataset for evaluating
+    parser.add_argument('--no_eval_mem', action='store_true', default=False)

src/args/setup_evaluating.py

+from argparse import ArgumentParser, Namespace
+from src.args.evaluate import evaluate_args
+from src.args.dataset import dataset_args
+from src.utils.general_utils import print_x, create_folders
+import os
+import math
+from src.datasets import DATA_INFORMATION
+from datetime import datetime
+
+
+def parse_general_evaluating_args():
+    parser = ArgumentParser('argument for evaluating')
+    dataset_args(parser)
+    evaluate_args(parser)
+    opt = parser.parse_args()
+    return opt
+
+
+def parse_all_evaluating_args(fixed_config, opt) -> Namespace:
+    opt.data_folder = None
+
+    # config
+    save_path_config = fixed_config['general']['save_path_config']
+    pattern_name = fixed_config['general']['pattern_name']
+    root_path = save_path_config['root_save_folder']
+    if opt.data_folder is None:
+        opt.data_folder = save_path_config['data_folder'].replace("{root_save_folder}", root_path)
+
+    # eval folder
+    current_time = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
+    eval_folder_name = f"eval_linear_{current_time}"
+
+    opt.eval_folder_path = os.path.join(opt.ckpt, '..', eval_folder_name)
+    create_folders(opt.eval_folder_path)
+
+    opt.eval_model_path = os.path.join(opt.eval_folder_path,
+                                       pattern_name['eval_model_name'].replace(
+                                           "{eval_folder_name}", eval_folder_name).replace(
+                                           "{target_task}", str(opt.target_task)))
+
+    opt.eval_config_path = os.path.join(opt.eval_folder_path,
+                                        pattern_name['config_json_name_eval'].replace(
+                                            "{target_task}", str(opt.target_task)))
+
+    opt.accuracy_single_task_path = os.path.join(opt.eval_folder_path,
+                                                 pattern_name['accuracy_single_task'].replace(
+                                                     "{target_task}", str(opt.target_task)))
+
+    iterations = opt.lr_decay_epochs.split(',')
+    opt.lr_decay_epochs = list([])
+    for it in iterations:
+        opt.lr_decay_epochs.append(int(it))
+
+    # warm-up for large-batch training
+    if opt.warm:
+        opt.warmup_from = 0.01
+        opt.warm_epochs = 10
+        if opt.cosine:
+            eta_min = opt.learning_rate * (opt.lr_decay_rate ** 3)
+            opt.warmup_to = eta_min + (opt.learning_rate - eta_min) * (
+                    1 + math.cos(math.pi * opt.warm_epochs / opt.epochs)) / 2
+        else:
+            opt.warmup_to = opt.learning_rate
+
+    opt.n_cls, opt.cls_per_task, opt.size = DATA_INFORMATION[opt.dataset]
+    if opt.num_tasks is not None:
+        print_x(f"opt.num_task is not None => process: opt.num_tasks = {opt.num_tasks}")
+        assert opt.n_cls % opt.num_tasks == 0
+        opt.cls_per_task = int(opt.n_cls / opt.num_tasks)
+
+    opt.predictor_hidden_dim = 256 if opt.dataset == 'tiny-imagenet' else 128
+
+    opt.origin_ckpt = opt.ckpt
+
+    opt.ckpt = os.path.join(opt.ckpt, f'last_{opt.replay_policy}_{opt.target_task}.pth')
+    opt.logpt = os.path.join(opt.logpt, f'replay_indices_{opt.replay_policy}_{opt.target_task}.npy')
+
+    # wandb dir
+    if opt.wandb:
+        opt.wandb_dir = root_path
+        opt.wandb_run_name_eval = pattern_name['job_wandb_name_eval'] \
+            .replace('{job_run_name}', opt.job_run_name) \
+            .replace('{target_task}', str(opt.target_task))
+    return opt

src/args/setup_training.py

+import argparse
+from src.args.train import utils_training_args
+from src.args.dataset import dataset_args
+from src.args.continual import continual_args
+import os
+import math
+from src.utils.general_utils import print_x, create_folders, get_feat_dim_projection
+from src.backbones.resnet_big import model_dict
+from datetime import datetime
+from src.plas_losses import LOSSES
+from src.datasets import DATA_INFORMATION
+
+
+def parse_all_training_args(fixed_config) -> argparse.Namespace:
+    parser = argparse.ArgumentParser('argument for training')
+
+    # add shared arguments
+    dataset_args(parser)
+    utils_training_args(parser)
+    continual_args(parser)
+
+    temp_args, _ = parser.parse_known_args()
+    LOSSES[temp_args.loss].add_specific_args(parser)
+
+    opt = parser.parse_args()
+
+    assert opt.model in model_dict
+
+    opt.save_freq = opt.epochs // 2
+
+    opt.n_cls, opt.cls_per_task, opt.size = DATA_INFORMATION[opt.dataset]
+    if opt.num_tasks is not None:
+        assert opt.n_cls % opt.num_tasks == 0
+        opt.cls_per_task = int(opt.n_cls / opt.num_tasks)
+
+    if opt.dataset == 'custom':
+        assert opt.data_folder is not None and opt.mean is not None and opt.std is not None
+
+    opt.dim = get_feat_dim_projection(dataset_name=opt.dataset)
+
+    # Set the path according to the environment
+    save_path_config = fixed_config['general']['save_path_config']
+    pattern_name = fixed_config['general']['pattern_name']
+
+    root_path = save_path_config['root_save_folder']
+    if opt.data_folder is None:
+        opt.data_folder = save_path_config['data_folder'].replace("{root_save_folder}", root_path)
+
+    experiments_folder = save_path_config['experiments_folder'].replace("{root_save_folder}", root_path)
+
+    current_time = datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
+
+    job_run_name = opt.job_run_name
+    run_exp_folder = save_path_config['run_exp_folder'].replace(
+        "{experiments_folder}", experiments_folder).replace(
+        "{job_run_name}", job_run_name).replace(
+        "{datetime}", current_time)
+    if opt.resume_target_task is not None and opt.path_resume_model_folder is not None:
+        print_x(f"Resume task: {opt.resume_target_task}, path_resume_model_folder = {opt.path_resume_model_folder}")
+        run_exp_folder = opt.path_resume_model_folder
+
+    opt.model_path = '{}/models'.format(run_exp_folder, opt.dataset)
+    opt.log_path = '{}/logs'.format(run_exp_folder)
+
+    iterations = opt.lr_decay_epochs.split(',')
+    opt.lr_decay_epochs = list([])
+    for it in iterations:
+        opt.lr_decay_epochs.append(int(it))
+
+    opt.encoder_folder = 'encoder'
+    if opt.cosine:
+        opt.encoder_folder = 'encoder_cosine'
+
+    if opt.loss == 'nc_samix' or opt.loss == 'fnc2':
+        if opt.dataset == 'tiny-imagenet':
+            opt.second_mark = 40
+        else:
+            opt.second_mark = 80
+
+    # warm-up
+    if opt.batch_size > 256:
+        opt.warm = True
+    if opt.warm:
+        opt.encoder_folder = '{}_warm'.format(opt.encoder_folder)
+        opt.warmup_from = 0.01
+        opt.warm_epochs = 10
+        if opt.cosine:
+            eta_min = opt.learning_rate * (opt.lr_decay_rate ** 3)
+            opt.warmup_to = eta_min + (opt.learning_rate - eta_min) * (
+                    1 + math.cos(math.pi * opt.warm_epochs / opt.epochs)) / 2
+        else:
+            opt.warmup_to = opt.learning_rate
+
+    opt.save_folder = os.path.join(opt.model_path, opt.encoder_folder)
+    create_folders(opt.save_folder)
+
+    opt.log_folder = os.path.join(opt.log_path, opt.encoder_folder)
+    create_folders(opt.log_folder)
+
+    # wandb dir
+    if opt.wandb:
+        opt.wandb_dir = root_path
+        opt.wandb_run_name_train = pattern_name['job_wandb_name_train'] \
+            .replace('{job_run_name}', job_run_name)
+    return opt

src/args/train.py

+from argparse import ArgumentParser
+from src.plas_losses import LOSSES
+from src.args.utils import general_args
+from src.distillation import LIST_DISTILLATION_LOSSES
+
+
+def utils_training_args(parser: ArgumentParser):
+    """Adds all arguments to a parser.
+    Args:
+        parser (ArgumentParser): parser to add dataset args to.
+    """
+    SUPPORTED_METHOD_LOSSES = LOSSES.keys()
+
+    # general args
+    general_args(parser)
+
+    # saving models frequency settings
+    parser.add_argument('--save_freq', type=int, default=1000,
+                        help='save frequency')
+
+    parser.add_argument('--epochs', type=int, default=1000,
+                        help='number of training epochs')
+    parser.add_argument('--start_epoch', type=int, default=None)
+
+    # optimization
+    parser.add_argument('--learning_rate', type=float, default=0.05,
+                        help='learning rate')
+    parser.add_argument('--lr_decay_epochs', type=str, default='700,800,900',
+                        help='where to decay lr, can be a list')
+    parser.add_argument('--lr_decay_rate', type=float, default=0.1,
+                        help='decay rate for learning rate')
+    parser.add_argument('--weight_decay', type=float, default=1e-4,
+                        help='weight decay')
+
+    # for customized dataset
+    parser.add_argument('--mean', type=str, help='mean of dataset in path in form of str tuple')
+    parser.add_argument('--std', type=str, help='std of dataset in path in form of str tuple')
+
+    # temperature
+    parser.add_argument('--current_temp', type=float, default=0.2,
+                        help='temperature for loss function')
+    parser.add_argument('--past_temp', type=float, default=0.01,
+                        help='temperature for loss function')
+    parser.add_argument('--temperature', type=float, default=0.07,
+                        help='temperature for loss function')
+
+    # memory
+    parser.add_argument('--mem_size', type=int, default=200)
+
+    # loss
+    parser.add_argument('--loss', type=str, default='nc_samix', choices=SUPPORTED_METHOD_LOSSES)
+    parser.add_argument('--distillation', type=str, default='hsd', choices=LIST_DISTILLATION_LOSSES)
+
+    # for resume training
+    parser.add_argument('--resume_target_task', type=int, default=None)
+    parser.add_argument('--path_resume_model_folder', type=str, default=None)
+    parser.add_argument('--path_model_resume_task', type=str, default=None)
+    parser.add_argument('--path_replay_indices', type=str, default=None)
+
+    # filter protos from task 0 to curren task for distillation
+    parser.add_argument('--filter_distillation_point', action='store_true', default=False)
+
+    # In case which mem = 0, set as 'not-use' with mem > 0
+    parser.add_argument('--log_folder_zero_mem', type=str, default='')
+
+    parser.add_argument('--eval_batch_size', type=int, default=256)
+    parser.add_argument('--eval_epochs', type=int, default=100)
+    parser.add_argument('--eval_learning_rate', type=float, default=1.0)
+    parser.add_argument('--eval_num_workers', type=int, default=16)
+
+    # hsd loss
+    parser.add_argument('--main_mark', type=int, default=20)
+
+    # samix
+    parser.add_argument('--samix', action='store_true', default=False)

src/args/utils.py

+from argparse import ArgumentParser
+from src.backbones.resnet_big import model_dict
+
+
+def general_args(parser: ArgumentParser):
+    # training settings
+    SUPPORTED_BACKBONES = model_dict.keys()
+    parser.add_argument('--seed', type=int, default=5, help='seed number')
+    parser.add_argument('--batch_size', type=int, default=256, help='batch_size')
+    parser.add_argument('--num_workers', type=int, default=16, help='num of workers to use')
+    parser.add_argument('--print_freq', type=int, default=10, help='print frequency')
+    # model backbone
+    parser.add_argument('--model', type=str, default='resnet18', choices=SUPPORTED_BACKBONES)
+    parser.add_argument('--size', type=int, default=32, help='parameter for RandomResizedCrop')
+    # other settings
+    parser.add_argument('--cosine', action='store_true', help='using cosine annealing')
+    parser.add_argument('--warm', action='store_true', help='warm-up for large batch training')
+    parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
+    # wandb arguments
+    parser.add_argument('--wandb', action='store_true', default=False,
+                        help='Enable wandb logging')
+    parser.add_argument('--wandb_project_name', type=str, required=False)
+    parser.add_argument('--wandb_entity', type=str, required=False)
+    parser.add_argument('--job_run_name', type=str, required=True)
+    parser.add_argument('--replay_policy', type=str,
+                        default='random', help='Use to find the model name to eval', required=False)
+    parser.add_argument('--num_tasks', type=int, default=None, help='Numer of tasks to run')

src/backbones/mlp.py

+import torch
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+
+    def __init__(self):
+        super(MLP, self).__init__()
+
+        self.conv1 = nn.Conv2d(1, 20, 5, 1)
+        self.conv2 = nn.Conv2d(20, 50, 5, 1)
+        self.fc1 = nn.Linear(4 * 4 * 50, 500)
+
+        self.encoder = nn.Sequential(
+            self.conv1,
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+            self.conv2,
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+            nn.Flatten(),
+            self.fc1,
+            nn.ReLU()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x)
+
+
+class SupConMLP(nn.Module):
+    def __init__(self, feat_dim=500):
+        super(SupConMLP, self).__init__()
+        self.encoder = MLP()
+        self.head = nn.Sequential(
+            nn.Linear(500, 500),
+            nn.ReLU(inplace=True),
+            nn.Linear(500, 500)
+        )
+
+    def forward(self, x, return_feat=False):
+        encoded = self.encoder(x)
+        feat = self.head(encoded)
+        if return_feat:
+            return feat, encoded
+        else:
+            return feat
+
+
+class LinearClassifier(nn.Module):
+    """Linear classifier"""
+
+    def __init__(self, num_classes=10):
+        super(LinearClassifier, self).__init__()
+        self.fc = nn.Linear(500, num_classes)
+
+    def forward(self, features):
+        return self.fc(features)

src/backbones/resnet_big.py

+"""Code from ResNet in PyTorch
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1, is_last=False):
+        super(BasicBlock, self).__init__()
+        self.is_last = is_last
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion * planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        preact = out
+        out = F.relu(out)
+        if self.is_last:
+            return out, preact
+        else:
+            return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1, is_last=False):
+        super(Bottleneck, self).__init__()
+        self.is_last = is_last
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, self.expansion * planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion * planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion * planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion * planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        preact = out
+        out = F.relu(out)
+        if self.is_last:
+            return out, preact
+        else:
+            return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, in_channel=3, zero_init_residual=False):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.conv1 = nn.Conv2d(in_channel, 64, kernel_size=3, stride=1, padding=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1] * (num_blocks - 1)
+        layers = []
+        for i in range(num_blocks):
+            stride = strides[i]
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x, layer=100):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = self.avgpool(out)
+        out = torch.flatten(out, 1)
+        return out
+
+
+def resnet18(**kwargs):
+    return ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+
+
+def resnet34(**kwargs):
+    return ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+
+
+def resnet50(**kwargs):
+    return ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+
+
+def resnet101(**kwargs):
+    return ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+
+
+model_dict = {
+    'resnet18': [resnet18, 512],
+    'resnet34': [resnet34, 512],
+    'resnet50': [resnet50, 2048],
+    'resnet101': [resnet101, 2048]
+}
+
+
+class LinearBatchNorm(nn.Module):
+    """Implements BatchNorm1d by BatchNorm2d, for SyncBN purpose"""
+
+    def __init__(self, dim, affine=True):
+        super(LinearBatchNorm, self).__init__()
+        self.dim = dim
+        self.bn = nn.BatchNorm2d(dim, affine=affine)
+
+    def forward(self, x):
+        x = x.view(-1, self.dim, 1, 1)
+        x = self.bn(x)
+        x = x.view(-1, self.dim)
+        return x
+
+
+class MLPLayer(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.fc = nn.Linear(in_dim, out_dim)
+        self.bn = nn.BatchNorm1d(out_dim)
+        self.act = nn.GELU()
+
+    def forward(self, x):
+        x = self.fc(x)
+        x = self.bn(x)
+        x = self.act(x)
+        return x
+
+
+class L2Norm(nn.Module):
+    def __init__(self, dim=1):
+        super().__init__()
+        self.dim = dim
+
+    def extra_repr(self):
+        return "dim={}".format(self.dim)
+
+    def forward(self, x):
+        return F.normalize(x, dim=self.dim)
+
+
+class SupConResNet(nn.Module):
+    """backbone + projection head"""
+
+    def __init__(self, name='resnet50', head='mlp', feat_dim=128, predictor_hidden_dim=None):
+        """
+        :param name:
+        :param head:
+        :param feat_dim:
+        :param predictor_hidden_dim:
+        :param dynamic_prototype_n_class: Number of classes (all prototypes created from scratch)
+        """
+        super(SupConResNet, self).__init__()
+        model_fun, dim_in = model_dict[name]
+        self.encoder = model_fun()
+        if head == 'linear':
+            self.head = nn.Linear(dim_in, feat_dim)
+        elif head == 'mlp':
+            self.head = nn.Sequential(
+                nn.Linear(dim_in, dim_in),
+                nn.BatchNorm1d(dim_in),
+                nn.ReLU(inplace=True),
+                nn.Linear(dim_in, feat_dim)
+            )
+        else:
+            raise NotImplementedError(
+                'head not supported: {}'.format(head))
+
+        self.distill_predictor = None
+        if predictor_hidden_dim is not None and predictor_hidden_dim != 0:
+            self.distill_predictor = nn.Sequential(
+                nn.Linear(feat_dim, predictor_hidden_dim),
+                nn.BatchNorm1d(predictor_hidden_dim),
+                nn.ReLU(),
+                nn.Linear(predictor_hidden_dim, feat_dim),
+            )
+
+    def reinit_head(self):
+        for layers in self.head.children():
+            if hasattr(layers, 'reset_parameters'):
+                layers.reset_parameters()
+
+    def forward(self, x, return_feat=False, norm=True, distill=False):
+        """
+        @param x: [2 x bsz, C, W, H]
+        @param return_feat:
+        @param norm:
+        @param distill:
+        @return:
+        """
+        encoded = self.encoder(x)
+        if norm:
+            feat = F.normalize(self.head(encoded), dim=1)
+        else:
+            feat = self.head(encoded)
+        output_return = ()
+        if return_feat:
+            output_return += (feat, encoded)
+        else:
+            output_return += (feat,)
+
+        if distill and self.distill_predictor:
+            distill_predictor_output = F.normalize(self.distill_predictor(feat), dim=1)
+            output_return += (distill_predictor_output,)
+
+        if len(output_return) == 1:
+            return output_return[0]
+        return output_return
+
+
+class SupCEResNet(nn.Module):
+    """encoder + classifier"""
+
+    def __init__(self, name='resnet50', num_classes=10):
+        super(SupCEResNet, self).__init__()
+        model_fun, dim_in = model_dict[name]
+        self.encoder = model_fun()
+        self.fc = nn.Linear(dim_in, num_classes)
+
+    def forward(self, x):
+        return self.fc(self.encoder(x))
+
+
+class LinearClassifier(nn.Module):
+    """Linear classifier"""
+
+    def __init__(self, name='resnet50', num_classes=10, two_layers=False):
+        super(LinearClassifier, self).__init__()
+        _, feat_dim = model_dict[name]
+        if two_layers:
+            self.fc = nn.Sequential(
+                nn.Linear(feat_dim, feat_dim),
+                nn.ReLU(),
+                nn.Linear(feat_dim, num_classes)
+            )
+        else:
+            self.fc = nn.Linear(feat_dim, num_classes)
+
+    def forward(self, features):
+        return self.fc(features)

src/buffers/buffer_setup.py

+from torch.utils.data import Dataset
+from torchvision import transforms, datasets
+import numpy as np
+from src.datasets.tiny_imagenet_dataset import TinyImagenet
+import random
+import math
+import torch
+
+
+def set_replay_samples(opt, model, prev_indices=None):
+    is_training = model.training
+    model.eval()
+
+    class IdxDataset(Dataset):
+        def __init__(self, dataset, indices):
+            self.dataset = dataset
+            self.indices = indices
+
+        def __len__(self):
+            return len(self.dataset)
+
+        def __getitem__(self, idx):
+            return self.indices[idx], self.dataset[idx]
+
+    val_transform = transforms.Compose([
+        transforms.ToTensor(),
+    ])
+
+    if opt.dataset == 'cifar10':
+        subset_indices = []
+        val_dataset = datasets.CIFAR10(root=opt.data_folder,
+                                       transform=val_transform,
+                                       download=True)
+        val_targets = np.array(val_dataset.targets)
+    elif opt.dataset == 'cifar100':
+        subset_indices = []
+        val_dataset = datasets.CIFAR100(root=opt.data_folder,
+                                        transform=val_transform,
+                                        download=True)
+        val_targets = np.array(val_dataset.targets)
+
+    elif opt.dataset == 'tiny-imagenet':
+        subset_indices = []
+        val_dataset = TinyImagenet(root=opt.data_folder,
+                                   transform=val_transform,
+                                   download=True)
+        val_targets = val_dataset.targets
+
+    else:
+        raise ValueError('dataset not supported: {}'.format(opt.dataset))
+
+    if prev_indices is None:
+        prev_indices = []
+        observed_classes = list(range(0, opt.target_task * opt.cls_per_task))
+    else:
+        shrink_size = ((opt.target_task - 1) * opt.mem_size / opt.target_task)
+        if len(prev_indices) > 0:
+            unique_cls = np.unique(val_targets[prev_indices])
+            _prev_indices = prev_indices
+            prev_indices = []
+            for c in unique_cls:
+                mask = val_targets[_prev_indices] == c
+                size_for_c = shrink_size / len(unique_cls)
+                p = size_for_c - (shrink_size // len(unique_cls))
+                if random.random() < p:
+                    size_for_c = math.ceil(size_for_c)
+                else:
+                    size_for_c = math.floor(size_for_c)
+                prev_indices += torch.tensor(_prev_indices)[mask][torch.randperm(mask.sum())[:size_for_c]].tolist()
+            print(np.unique(val_targets[prev_indices], return_counts=True))
+        observed_classes = list(
+            range(max(opt.target_task - 1, 0) * opt.cls_per_task, (opt.target_task) * opt.cls_per_task))
+
+    if len(observed_classes) == 0:
+        return prev_indices
+    observed_indices = []
+    for tc in observed_classes:
+        observed_indices += np.where(val_targets == tc)[0].tolist()
+    val_observed_targets = val_targets[observed_indices]
+    val_unique_cls = np.unique(val_observed_targets)
+
+    selected_observed_indices = []
+    for c_idx, c in enumerate(val_unique_cls):
+        size_for_c_float = (
+                (opt.mem_size - len(prev_indices) - len(selected_observed_indices)) / (len(val_unique_cls) - c_idx))
+        p = size_for_c_float - ((opt.mem_size - len(prev_indices) - len(selected_observed_indices)) // (
+                len(val_unique_cls) - c_idx))
+        if random.random() < p:
+            size_for_c = math.ceil(size_for_c_float)
+        else:
+            size_for_c = math.floor(size_for_c_float)
+        mask = val_targets[observed_indices] == c
+        selected_observed_indices += torch.tensor(observed_indices)[mask][
+            torch.randperm(mask.sum())[:size_for_c]].tolist()
+    print(np.unique(val_targets[selected_observed_indices], return_counts=True))
+    model.is_training = is_training
+    return prev_indices + selected_observed_indices

src/datasets/__init__.py

+from enum import Enum
+
+
+class Dataset(Enum):
+    CIFAR_10 = "cifar10"
+    TINY_IMAGENET = "tiny-imagenet"
+    CIFAR_100 = "cifar100"
+    CUSTOM = "custom"
+
+
+DATA_INFORMATION = {
+    "cifar10": [10, 2, 32],  # n_cls, cls_per_task, image_size
+    "tiny-imagenet": [200, 20, 64],
+    "cifar100": [100, 20, 32],
+    "custom": None
+}

src/datasets/evaluate_dataloader.py

+import numpy as np
+import torch
+from torch.utils.data import Dataset, Subset, WeightedRandomSampler
+from src.datasets import DATA_INFORMATION
+from src.utils.transforms_utils import get_transforms
+from src.utils.general_utils import get_train_datasets, get_val_datasets
+
+
+def set_loader(opt, replay_indices):
+    replay_indices = replay_indices.tolist()
+    train_transform, val_transform = get_transforms(opt)
+    all_target_classes = list(range(0, (opt.target_task + 1) * opt.cls_per_task))
+    _train_dataset = get_train_datasets(opt, train_transform)
+    if opt.dataset in DATA_INFORMATION:
+        subset_indices = []
+        _train_targets = np.array(_train_dataset.targets)
+        start_cls = opt.target_task * opt.cls_per_task
+        if opt.no_eval_mem:
+            start_cls = 0
+            replay_indices = []
+
+        for tc in range(start_cls, (opt.target_task + 1) * opt.cls_per_task):
+            subset_indices += np.where(np.array(_train_dataset.targets) == tc)[0].tolist()
+        subset_indices += replay_indices
+        ut, uc = np.unique(_train_targets[subset_indices], return_counts=True)
+
+        weights = np.array([0.] * len(subset_indices))
+        for t, c in zip(ut, uc):
+            weights[_train_targets[subset_indices] == t] = 1. / c
+
+        train_dataset = Subset(_train_dataset, subset_indices)
+        subset_indices = []
+        _val_dataset = get_val_datasets(opt, val_transform)
+        for tc in all_target_classes:
+            subset_indices += np.where(np.array(_val_dataset.targets) == tc)[0].tolist()
+        val_dataset = Subset(_val_dataset, subset_indices)
+    else:
+        raise ValueError('Dataset not supported: {}'.format(opt.dataset))
+    train_sampler = WeightedRandomSampler(torch.Tensor(weights), len(weights))
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=opt.batch_size, shuffle=(train_sampler is None),
+        num_workers=opt.num_workers, pin_memory=True, sampler=train_sampler)
+    val_loader = torch.utils.data.DataLoader(
+        val_dataset, batch_size=256, shuffle=False,
+        num_workers=8, pin_memory=True)
+    return train_loader, val_loader, uc

src/datasets/is_subset.py

+from torch.utils.data import Subset
+
+
+class IS_Subset(Subset):
+    """
+    Defines dataset with importance sampling weight.
+    """
+
+    def __init__(self, dataset, indices, IS_weight) -> None:
+        super().__init__(dataset, indices)
+        self.weight = IS_weight
+
+    def __getitem__(self, idx):
+        if isinstance(idx, list):
+            index = [self.indices[i] for i in idx]
+            weight = [self.weight[i] for i in idx]
+        else:
+            index = self.indices[idx]
+            weight = self.weight[idx]
+
+        return super().__getitem__(idx) + (weight, index)
+
+    def __len__(self):
+        return super().__len__()

src/datasets/pretrain_dataloader.py

+import numpy as np
+import torch
+from torch.utils.data import Subset
+from src.utils.general_utils import TwoCropTransform, get_train_datasets
+from src.utils.transforms_utils import get_transforms
+
+
+def set_loader_with_replay(opt, replay_indices):
+    train_transform, _ = get_transforms(opt)
+    all_target_classes = list(range(opt.target_task * opt.cls_per_task, (opt.target_task + 1) * opt.cls_per_task))
+    subset_indices = []
+    _train_dataset = get_train_datasets(opt, TwoCropTransform(train_transform))
+    for tc in all_target_classes:
+        subset_indices += np.where(np.array(_train_dataset.targets) == tc)[0].tolist()
+    subset_indices += replay_indices
+    train_dataset = Subset(_train_dataset, subset_indices)
+    train_sampler = None
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=opt.batch_size, shuffle=(train_sampler is None),
+        num_workers=opt.num_workers, pin_memory=True, sampler=train_sampler)
+    return train_loader, subset_indices

src/datasets/tiny_imagenet_dataset.py

+import os
+import numpy as np
+import torchvision.transforms as transforms
+from torch.utils.data import Dataset
+from PIL import Image
+
+
+class TinyImagenet(Dataset):
+    """
+    Defines Tiny Imagenet as for the others pytorch datasets.
+    """
+
+    def __init__(self, root: str, train: bool = True, transform: transforms = None,
+                 target_transform: transforms = None, download: bool = False) -> None:
+        self.not_aug_transform = transforms.Compose([transforms.ToTensor()])
+        self.root = root
+        self.train = train
+        self.transform = transform
+        self.target_transform = target_transform
+        self.download = download
+        self.extracted_folder = os.path.join(root, 'tiny-imagenet')
+        if not os.path.isdir(self.extracted_folder):
+            os.mkdir(self.extracted_folder)
+        if download:
+            print(f"[TINY IMAGENET] Download Tiny Imagenet dataset...")
+            if os.path.isdir(self.extracted_folder) and len(os.listdir(self.extracted_folder)) > 0:
+                print('Download not needed, files already on disk.')
+            else:
+                from google_drive_downloader import GoogleDriveDownloader as gdd
+                print('Downloading dataset')
+                dest_path = os.path.join(self.extracted_folder, 'tiny-imagenet-processed.zip')
+                gdd.download_file_from_google_drive(
+                    file_id='id_file_name',  # replace with actual file ID - Due to anonymity requirements, please set this manually using your own Google Drive file ID
+                    dest_path=dest_path,
+                    unzip=True)
+                # unzip again
+                print('Finish downloading dataset')
+        self.data = []
+        for num in range(20):
+            self.data.append(np.load(os.path.join(self.extracted_folder,
+                                                  'processed/x_%s_%02d.npy' % (
+                                                      'train' if self.train else 'val', num + 1))))
+        self.data = np.concatenate(np.array(self.data))
+        self.targets = []
+        for num in range(20):
+            self.targets.append(np.load(os.path.join(self.extracted_folder,
+                                                     'processed/y_%s_%02d.npy' % (
+                                                         'train' if self.train else 'val', num + 1))))
+        self.targets = np.concatenate(np.array(self.targets))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        img, target = self.data[index], self.targets[index]
+        img = Image.fromarray(np.uint8(255 * img))
+        original_img = img.copy()
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+        if hasattr(self, 'logits'):
+            return img, target, original_img, self.logits[index]
+        return img, target

src/distillation/__init__.py

+LIST_DISTILLATION_LOSSES = [
+    "ird",
+    "sprd",
+    "hsd",
+    "no-distillation"
+]

src/distillation/stab_loss.py

+import torch
+
+
+def compute_ird(features, temp):
+    features1_sim = torch.div(torch.matmul(features, features.T), temp)
+    logits_mask = torch.scatter(
+        torch.ones_like(features1_sim), 1,
+        torch.arange(features1_sim.size(0)).view(-1, 1).cuda(non_blocking=True), 0)
+
+    logits_max1, _ = torch.max(features1_sim * logits_mask, dim=1, keepdim=True)
+    features1_sim = features1_sim - logits_max1.detach()
+    row_size = features1_sim.size(0)
+
+    logits1 = torch.exp(features1_sim[logits_mask.bool()].view(row_size, -1)) / torch.exp(
+        features1_sim[logits_mask.bool()].view(row_size, -1)).sum(dim=1, keepdim=True)
+    return logits1
+
+
+def compute_sprd(features, prototypes, temp):
+    features1_sim = torch.div(torch.matmul(features, prototypes.T), temp)
+    logits_max1, _ = torch.max(features1_sim, dim=1, keepdim=True)
+    features1_sim = features1_sim - logits_max1.detach()
+    return torch.exp(features1_sim) / torch.exp(features1_sim).sum(dim=1, keepdim=True)
+
+
+def compute_distillation_loss(opt, features1_prev_task, frozen_model, images, set_prototypes=None,
+                              alpha_balance_distillation=None):
+    if opt.distillation == 'ird':
+        logits1 = compute_ird(features1_prev_task, opt.current_temp)
+        with torch.no_grad():
+            features2_prev_task = frozen_model(images)
+            logits2 = compute_ird(features2_prev_task, opt.past_temp)
+        return (-logits2 * torch.log(logits1)).sum(1).mean()
+
+    elif opt.distillation == 'sprd':
+        assert set_prototypes is not None
+        logits1 = compute_sprd(features1_prev_task, set_prototypes, opt.current_temp)
+        with torch.no_grad():
+            features2_prev_task = frozen_model(images)
+            logits2 = compute_sprd(features2_prev_task, set_prototypes, opt.past_temp)
+        return (-logits2 * torch.log(logits1)).sum(1).mean()
+
+    elif opt.distillation == 'hsd':
+        assert set_prototypes is not None
+        assert alpha_balance_distillation is not None
+        current_temp_sprd = opt.current_temp
+        past_temp_sprd = opt.past_temp
+        if hasattr(opt, "current_temp_sprd") and hasattr(opt, "past_temp_sprd"):
+            if opt.current_temp_sprd is not None and opt.past_temp_sprd is not None:
+                current_temp_sprd = opt.current_temp_sprd
+                past_temp_sprd = opt.past_temp_sprd
+        logits1 = compute_ird(features1_prev_task, opt.current_temp)
+        logits1_proto = compute_sprd(features1_prev_task, set_prototypes, current_temp_sprd)
+        with torch.no_grad():
+            features2_prev_task = frozen_model(images)
+            logits2 = compute_ird(features2_prev_task, opt.past_temp)
+            logits2_proto = compute_sprd(features2_prev_task, set_prototypes, past_temp_sprd)
+        return (1 - alpha_balance_distillation) * (-logits2 * torch.log(logits1)).sum(
+            1).mean() + alpha_balance_distillation * (-logits2_proto * torch.log(logits1_proto)).sum(1).mean()

src/plas_losses/__init__.py

+from src.plas_losses.main_loss_fnc2 import FNC2
+from src.plas_losses.main_loss_dot_regression import DotRegression
+from src.plas_losses.loss_supcon import SupConLoss
+from src.plas_losses.main_loss_nc_samix import NCSAMix
+
+__all__ = [
+    "FNC2",
+    "SupConLoss",
+    "NCSAMix",
+    "DotRegression"
+]
+
+LIST_LOSSES_PROTOTYPES = [
+    "fnc2",
+    "dot_regression",
+    "nc_samix"
+]
+
+LOSSES = {
+    "fnc2": FNC2,
+    "supcon": SupConLoss,
+    "nc_samix": NCSAMix,
+    "dr": DotRegression
+}

src/plas_losses/base_loss.py

+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+
+
+class BaseLoss(ABC):
+    @staticmethod
+    @abstractmethod
+    def add_specific_args(parser: ArgumentParser):
+        """Add model specific args here"""
+        pass

src/plas_losses/base_prototype_loss.py

+from abc import ABC, abstractmethod
+
+
+class BasePrototypeBasedLoss(ABC):
+
+    def __init__(self):
+        self.points = None
+
+    @abstractmethod
+    def set_set_prototypes(self, set_prototypes):
+        print(f"Shape of fixed protos = {set_prototypes.shape}")
+        self.points = set_prototypes

src/plas_losses/loss_supcon.py

+from __future__ import print_function
+import torch
+import torch.nn as nn
+from src.utils.general_utils import print_x
+from src.plas_losses.base_loss import BaseLoss
+
+import argparse
+
+'''
+The original source code from Co2L[ICCV'21] paper
+'''
+
+
+class SupConLoss(BaseLoss, nn.Module):
+    def __init__(self, temperature=0.07, contrast_mode='all', base_temperature=0.07, **kwargs):
+        super(SupConLoss, self).__init__()
+        print_x(f"**kwargs = {kwargs}")
+        self.temperature = temperature
+        self.contrast_mode = contrast_mode
+        self.base_temperature = base_temperature
+
+    @staticmethod
+    def add_specific_args(parser: argparse.ArgumentParser):
+        pass
+
+    def forward(self, features, labels=None, mask=None, target_labels=None, reduction='mean', **kwargs):
+        assert target_labels is not None and len(
+            target_labels) > 0, "Target labels should be given as a list of integer"
+        device = (torch.device('cuda')
+                  if features.is_cuda
+                  else torch.device('cpu'))
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, d-dimensions],'
+                             'at least 3 dimensions are required')
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)
+        batch_size = features.shape[0]  # [bsz, 2, d]
+        if labels is not None and mask is not None:
+            raise ValueError('Cannot define both `labels` and `mask`')
+        elif labels is None and mask is None:
+            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
+        elif labels is not None:
+            labels = labels.contiguous().view(-1, 1)
+            if labels.shape[0] != batch_size:
+                raise ValueError('Num of labels does not match num of features')
+            mask = torch.eq(labels, labels.T).float().to(device)
+        else:
+            mask = mask.float().to(device)
+        contrast_count = features.shape[1]
+        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
+        if self.contrast_mode == 'one':
+            anchor_feature = features[:, 0]
+            anchor_count = 1
+        elif self.contrast_mode == 'all':
+            anchor_feature = contrast_feature
+            anchor_count = contrast_count
+        else:
+            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))
+        anchor_dot_contrast = torch.div(
+            torch.matmul(anchor_feature, contrast_feature.T),
+            self.temperature)
+        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
+        logits = anchor_dot_contrast - logits_max.detach()
+
+        mask = mask.repeat(anchor_count, contrast_count)
+        logits_mask = torch.scatter(torch.ones_like(mask), 1,
+                                    torch.arange(batch_size * anchor_count).view(-1, 1).to(device), 0)
+        mask = mask * logits_mask
+        exp_logits = torch.exp(logits) * logits_mask
+        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
+        pos_per_sample = mask.sum(1)
+        pos_per_sample[pos_per_sample < 1e-6] = 1.0
+        mean_log_prob_pos = (mask * log_prob).sum(1) / pos_per_sample
+        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+
+        curr_class_mask = torch.zeros_like(labels)
+        for tc in target_labels:
+            curr_class_mask += (labels == tc)
+        curr_class_mask = curr_class_mask.view(-1).to(device)
+        loss = curr_class_mask * loss.view(anchor_count, batch_size)
+
+        if reduction == 'mean':
+            loss = loss.mean()
+        elif reduction == 'none':
+            loss = loss.mean(0)
+        else:
+            raise ValueError('loss reduction not supported: {}'.
+                             format(reduction))
+        return loss

src/plas_losses/main_loss_dot_regression.py

+# -*- coding:  utf-8 -*-
+from __future__ import print_function
+import torch
+import torch.nn as nn
+from src.plas_losses.base_prototype_loss import BasePrototypeBasedLoss
+from src.plas_losses.base_loss import BaseLoss
+from src.utils.general_utils import print_x
+import argparse
+
+
+class DotRegression(BaseLoss, BasePrototypeBasedLoss, nn.Module):
+    def set_set_prototypes(self, set_prototypes):
+        super().set_set_prototypes(set_prototypes)
+
+    def __init__(self,
+                 dim: int = None,
+                 n_cls: int = None,
+                 **kwargs):
+        nn.Module.__init__(self)
+        BasePrototypeBasedLoss.__init__(self)
+        BaseLoss.__init__(self)
+        print_x(f"**kwargs = {kwargs}")
+        self.d = dim
+        self.k = n_cls
+        print_x(self)
+
+    def __str__(self):
+        lines = [f"Fixed prototypes: (dim_d={self.d}, n_cls={self.k})"]
+        return "[Dot-Regression]" + "\n" + "\n".join(["\t" + line for line in lines])
+
+    @staticmethod
+    def add_specific_args(parser: argparse.ArgumentParser):
+        print_x(f"Adding specific args....")
+        parser.add_argument('--current_temp_sprd', type=float, default=None,
+                            help='temperature for current - S-PRD distillation')
+        parser.add_argument('--past_temp_sprd', type=float, default=None,
+                            help='temperature for previous - S-PRD distillation')
+
+    def forward(self, features, labels=None, target_labels=None, **kwargs):
+        """
+        @param features: [bsz, 2, d]
+        @param labels: [bsz,]
+        @param target_labels: [n_cls_per_task,]
+        @param kwargs:
+        @param return:
+       """
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, d-dimensions],'
+                             'at least 3 dimensions are required')
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)
+        anchor_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
+        batch_size, anchor_count = features.shape[:2]
+
+        self.points = self.points.to(features)
+        repeated_labels = labels.repeat(anchor_count)
+        corresponding_prototypes = self.points[repeated_labels]
+        dot_product = torch.sum(anchor_feature * corresponding_prototypes, dim=1)
+
+        curr_class_mask = torch.zeros_like(repeated_labels)
+        for tc in target_labels:
+            curr_class_mask += (repeated_labels == tc)
+        curr_class_mask = curr_class_mask.view(-1).to(features)
+        dot_product = curr_class_mask * dot_product
+        return torch.mean(0.5 * ((dot_product - 1) ** 2))

src/plas_losses/main_loss_fnc2.py

+# -*- coding:  utf-8 -*-
+from __future__ import print_function
+import torch
+import torch.nn as nn
+from src.utils.general_utils import print_x
+from src.plas_losses.base_loss import BaseLoss
+from src.plas_losses.base_prototype_loss import BasePrototypeBasedLoss
+import argparse
+
+
+class FNC2(BaseLoss, BasePrototypeBasedLoss, nn.Module):
+    def set_set_prototypes(self, set_prototypes):
+        super().set_set_prototypes(set_prototypes)
+
+    def __init__(self,
+                 temperature: float = 0.07,
+                 base_temperature: float = 0.07,
+                 focal_gamma: int = 0,
+                 **kwargs):
+        super(FNC2, self).__init__()
+        print_x(f"**kwargs = {kwargs}")
+        self.temperature = temperature
+        self.base_temperature = base_temperature
+        self.focal_gamma = focal_gamma
+        print(self)
+
+    def __str__(self):
+        lines = [f"temperature: {self.temperature}", f"base_temperature: {self.base_temperature}"]
+        return "[FNC2]" + "\n" + "\n".join(["\t" + line for line in lines])
+
+    def get_mask(self, labels, mask, batch_size: int, device):
+        # labels|  mask | mask_output
+        #   0   |   0   | torch.eye(batch_size)
+        #   0   |   1   | mask
+        #   1   |   0   | torch.eq(labels, labels.T)
+        #   1   |   1   | "Cannot define both `labels` and `mask`"
+        if labels is not None and mask is not None:
+            raise ValueError("Cannot define both `labels` and `mask`")
+        elif labels is None and mask is None:
+            return torch.eye(batch_size, dtype=torch.float32).to(device)
+        elif labels is not None:
+            labels = labels.contiguous().view(-1, 1)
+            if labels.shape[0] != batch_size:
+                raise ValueError("Num of labels does not match num of features")
+            return torch.eq(labels, labels.T).float().to(device)
+        else:
+            return mask.float().to(device)
+
+    @staticmethod
+    def add_specific_args(parser: argparse.ArgumentParser):
+        print_x(f"Adding specific args for FNC2....")
+        parser.add_argument('--focal_gamma', type=int, default=0)
+        parser.add_argument('--current_temp_sprd', type=float, default=None,
+                            help='temperature for current - S-PRD distillation')
+        parser.add_argument('--past_temp_sprd', type=float, default=None,
+                            help='temperature for previous - S-PRD distillation')
+
+    def compute_fnc2_loss(self, anchor_feature, anchor_count, labels, batch_size, device, target_labels,
+                          corresponding_prototypes):
+        """
+        @param anchor_feature: [bsz x 2, d]
+        @param anchor_count: 2
+        @param labels:
+        @param batch_size:
+        @param device:
+        @param target_labels:
+        @param corresponding_prototypes:
+        @return:
+        """
+        mask = self.get_mask(labels, None, batch_size, device)
+        anchor_dot_contrast = torch.div(torch.matmul(anchor_feature, anchor_feature.T), self.temperature)
+        index_previous_proto = min(target_labels)
+        previous_proto = self.points[:index_previous_proto]
+        anchor_contrast_previous_proto = None
+        if previous_proto.numel() != 0:
+            anchor_contrast_previous_proto = torch.div(
+                torch.matmul(anchor_feature, previous_proto.T), self.temperature)
+
+        points_dot_contrast = torch.div(
+            torch.mul(anchor_feature, corresponding_prototypes).sum(1), self.temperature
+        )
+        I = torch.arange(corresponding_prototypes.size(0))
+        anchor_dot_contrast[I, I] = points_dot_contrast
+
+        if anchor_contrast_previous_proto is not None:
+            logits_max, _ = torch.max(torch.cat((anchor_dot_contrast, anchor_contrast_previous_proto), dim=1), dim=1,
+                                      keepdim=True)
+            previous_logits_sample_proto = anchor_contrast_previous_proto - logits_max.detach()
+        else:
+            logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
+        logits = anchor_dot_contrast - logits_max.detach()
+        mask = mask.repeat(anchor_count, anchor_count)
+        logits_mask = torch.scatter(torch.ones_like(mask), 1,
+                                    torch.arange(batch_size * anchor_count).view(-1, 1).to(device), 0)
+        exp_logits = torch.exp(logits) * logits_mask
+        if anchor_contrast_previous_proto is not None:
+            exp_previous_proto = torch.exp(previous_logits_sample_proto)
+            log_prob = logits - torch.log(
+                exp_logits.sum(1, keepdim=True) + exp_previous_proto.sum(1, keepdim=True))
+        else:
+            log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
+        weight = (1 - torch.exp(log_prob)) ** self.focal_gamma
+        pos_per_sample = mask.sum(1)
+        pos_per_sample[pos_per_sample < 1e-6] = 1.0
+        mean_log_prob_pos = (weight * mask * log_prob).sum(1) / pos_per_sample
+        mean_log_prob_pos = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+        print_x(f"Shape of mean_log_prob_pos in FNC2 loss = {mean_log_prob_pos.shape}")
+        return mean_log_prob_pos  # [2 x bsz,]
+
+    def forward(self, features, labels=None, target_labels=None, **kwargs):
+        """
+        @param features: [bsz, 2, d]
+        @param labels: [bsz,]
+        @param target_labels: [n_clas_per_task]
+        @param kwargs:
+        @return:
+        """
+        assert target_labels is not None and len(
+            target_labels) > 0, "Target labels should be given as a list of integer"
+        device = (torch.device('cuda') if features.is_cuda else torch.device('cpu'))
+
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, d-dimensions],'
+                             'at least 3 dimensions are required')
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)
+
+        anchor_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
+        batch_size, anchor_count = features.shape[:2]
+        self.points = self.points.to(features)
+        repeated_labels = labels.repeat(anchor_count)
+        corresponding_prototypes = self.points[repeated_labels]
+        mean_log_prob_pos = self.compute_fnc2_loss(anchor_feature=anchor_feature, anchor_count=anchor_count,
+                                                   labels=labels, batch_size=batch_size, device=device,
+                                                   target_labels=target_labels,
+                                                   corresponding_prototypes=corresponding_prototypes)
+        curr_class_mask = torch.zeros_like(labels)
+        for tc in target_labels:
+            curr_class_mask += (labels == tc)
+        curr_class_mask = curr_class_mask.view(-1).to(device)
+        loss = curr_class_mask * mean_log_prob_pos.view(anchor_count, batch_size)
+        return loss.mean()

src/plas_losses/main_loss_nc_samix.py

+# -*- coding:  utf-8 -*-
+from __future__ import print_function
+import torch
+import torch.nn as nn
+from src.plas_losses.base_prototype_loss import BasePrototypeBasedLoss
+from src.utils.general_utils import print_x
+from src.plas_losses.base_loss import BaseLoss
+import argparse
+
+
+class NCSAMix(BaseLoss, BasePrototypeBasedLoss, nn.Module):
+    def set_set_prototypes(self, set_prototypes):
+        super().set_set_prototypes(set_prototypes)
+
+    def __init__(self,
+                 loss_normal_samples: str = None,
+                 temperature: float = 0.07,
+                 base_temperature: float = 0.07,
+                 focal_gamma: int = 0,
+                 weight_samix_loss: float = 1.0,
+                 **kwargs):
+        nn.Module.__init__(self)
+        BasePrototypeBasedLoss.__init__(self)
+        BaseLoss.__init__(self)
+        print_x(f"**kwargs = {kwargs}")
+        self.loss_normal_samples = loss_normal_samples
+        self.temperature = temperature
+        self.base_temperature = base_temperature
+        self.focal_gamma = focal_gamma
+        self.weight_samix_loss = weight_samix_loss
+        self.print_one_previous = True
+        self.print_one_future = True
+        print(self)
+
+    def __str__(self):
+        lines = [f"loss_normal_samples = {self.loss_normal_samples}",
+                 f"weight_samix_loss = {self.weight_samix_loss}",
+                 f"temperature: {self.temperature}",
+                 f"base_temperature: {self.base_temperature}",
+                 f"focal_gamma: {self.focal_gamma}"]
+        return "[NCSAMix]" + "\n" + "\n".join(["\t" + line for line in lines])
+
+    @staticmethod
+    def add_specific_args(parser: argparse.ArgumentParser):
+        print_x(f"Adding specific args....")
+        parser.add_argument('--focal_gamma', type=int, default=None)
+        parser.add_argument('--current_temp_sprd', type=float, default=None,
+                            help='temperature for current - S-PRD distillation')
+        parser.add_argument('--past_temp_sprd', type=float, default=None,
+                            help='temperature for previous - S-PRD distillation')
+        parser.add_argument('--loss_normal_samples', type=str, default='fnc2', choices=['dr', 'fnc2'])
+        parser.add_argument('--weight_samix_loss', type=float, default=1.0)
+
+    def get_mask(self, labels, mask, batch_size: int, device):
+        # labels|  mask | mask_output
+        #   0   |   0   | torch.eye(batch_size)
+        #   0   |   1   | mask
+        #   1   |   0   | torch.eq(labels, labels.T)
+        #   1   |   1   | "Cannot define both `labels` and `mask`"
+        if labels is not None and mask is not None:
+            raise ValueError("Cannot define both `labels` and `mask`")
+        elif labels is None and mask is None:
+            return torch.eye(batch_size, dtype=torch.float32).to(device)
+        elif labels is not None:
+            labels = labels.contiguous().view(-1, 1)
+            if labels.shape[0] != batch_size:
+                raise ValueError("Num of labels does not match num of features")
+            return torch.eq(labels, labels.T).float().to(device)
+        else:
+            return mask.float().to(device)
+
+    def compute_fnc2_loss(self, anchor_feature, anchor_count, labels, batch_size, device, target_labels,
+                          corresponding_prototypes):
+        """
+        @param anchor_feature: [bsz x 2, d]
+        @param anchor_count: 2
+        @param labels:
+        @param batch_size:
+        @param device:
+        @param target_labels:
+        @param corresponding_prototypes:
+        @return:
+        """
+        mask = self.get_mask(labels, None, batch_size, device)
+        anchor_dot_contrast = torch.div(torch.matmul(anchor_feature, anchor_feature.T), self.temperature)
+        index_previous_proto = min(target_labels)
+        previous_proto = self.points[:index_previous_proto]
+        anchor_contrast_previous_proto = None
+        if previous_proto.numel() != 0:
+            anchor_contrast_previous_proto = torch.div(
+                torch.matmul(anchor_feature, previous_proto.T), self.temperature)
+
+        points_dot_contrast = torch.div(
+            torch.mul(anchor_feature, corresponding_prototypes).sum(1), self.temperature
+        )
+        I = torch.arange(corresponding_prototypes.size(0))
+        anchor_dot_contrast[I, I] = points_dot_contrast
+
+        if anchor_contrast_previous_proto is not None:
+            logits_max, _ = torch.max(torch.cat((anchor_dot_contrast, anchor_contrast_previous_proto), dim=1), dim=1,
+                                      keepdim=True)
+            previous_logits_sample_proto = anchor_contrast_previous_proto - logits_max.detach()
+        else:
+            logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
+        logits = anchor_dot_contrast - logits_max.detach()
+        mask = mask.repeat(anchor_count, anchor_count)
+        logits_mask = torch.scatter(torch.ones_like(mask), 1,
+                                    torch.arange(batch_size * anchor_count).view(-1, 1).to(device), 0)
+        exp_logits = torch.exp(logits) * logits_mask
+        if anchor_contrast_previous_proto is not None:
+            exp_previous_proto = torch.exp(previous_logits_sample_proto)
+            log_prob = logits - torch.log(
+                exp_logits.sum(1, keepdim=True) + exp_previous_proto.sum(1, keepdim=True))
+        else:
+            log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
+        weight = (1 - torch.exp(log_prob)) ** self.focal_gamma
+        pos_per_sample = mask.sum(1)
+        pos_per_sample[pos_per_sample < 1e-6] = 1.0
+        mean_log_prob_pos = (weight * mask * log_prob).sum(1) / pos_per_sample
+        mean_log_prob_pos = - (self.temperature / self.base_temperature) * mean_log_prob_pos
+        return mean_log_prob_pos  # [2 x bsz,]
+
+    def generate_samix_prototypes_slerp(self, repeated_labels, mixed_indices, corresponding_prototypes, lamb_ratio):
+        perm_prototypes = self.points[repeated_labels[mixed_indices]]
+        cos_theta = torch.sum(corresponding_prototypes * perm_prototypes, dim=1, keepdim=True)
+        theta = torch.acos(torch.clamp(cos_theta, -1.0, 1.0))
+        sin_theta = torch.sin(theta)
+        slerp_prototypes = torch.empty_like(corresponding_prototypes)
+        small_angle_threshold = 1e-6
+        small_angle_mask = (theta < small_angle_threshold).squeeze(-1)
+        slerp_prototypes[small_angle_mask] = corresponding_prototypes[small_angle_mask]
+        non_zero_mask = ~small_angle_mask
+        if non_zero_mask.any():
+            theta_non_zero = theta[non_zero_mask]
+            sin_theta_non_zero = sin_theta[non_zero_mask]
+            p1_weight = torch.sin(lamb_ratio * theta_non_zero) / sin_theta_non_zero
+            p2_weight = torch.sin((1 - lamb_ratio) * theta_non_zero) / sin_theta_non_zero
+            p1_weight = p1_weight.expand(-1, corresponding_prototypes.size(1))
+            p2_weight = p2_weight.expand(-1, perm_prototypes.size(1))
+            slerp_prototypes[non_zero_mask] = (p1_weight * corresponding_prototypes[non_zero_mask]
+                                               + p2_weight * perm_prototypes[non_zero_mask])
+        return slerp_prototypes
+
+    def forward(self, features, labels=None, target_labels=None, samix_features=None, lamb_ratio=None,
+                mixed_indices=None, **kwargs):
+        """
+        @param features: [bsz, 2, d]
+        @param labels: [bsz,]
+        @param target_labels: [n_cls_per_task,]
+        @param samix_features: [2 x bsz, d]
+        @param lamb_ratio: samixed_views = data * lamb_ratio + (1 - lamb_ratio) * data_perm
+        @param mixed_indices: [2 x bsz,]
+        @param kwargs:
+        @return:
+        """
+        assert target_labels is not None and len(
+            target_labels) > 0, "Target labels should be given as a list of integer"
+        assert self.loss_normal_samples in ['dr', 'fnc2']
+        device = (torch.device('cuda') if features.is_cuda else torch.device('cpu'))
+        if len(features.shape) < 3:
+            raise ValueError('`features` needs to be [bsz, n_views, d-dimensions],'
+                             'at least 3 dimensions are required')
+        if len(features.shape) > 3:
+            features = features.view(features.shape[0], features.shape[1], -1)
+
+        anchor_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
+        batch_size, anchor_count = features.shape[:2]
+        self.points = self.points.to(features)
+        repeated_labels = labels.repeat(anchor_count)
+        corresponding_prototypes = self.points[repeated_labels]
+
+        # For normal features
+        if self.loss_normal_samples == 'dr':
+            loss_normal_features = torch.sum(anchor_feature * corresponding_prototypes, dim=1)
+            loss_normal_features = 0.5 * ((loss_normal_features - 1) ** 2)
+        elif self.loss_normal_samples == 'fnc2':
+            assert self.focal_gamma is not None
+            loss_normal_features = self.compute_fnc2_loss(anchor_feature=anchor_feature, anchor_count=anchor_count,
+                                                          labels=labels, batch_size=batch_size, device=device,
+                                                          target_labels=target_labels,
+                                                          corresponding_prototypes=corresponding_prototypes)
+        else:
+            raise ValueError(f"args --loss_normal_samples is not valid")
+
+        # Remove old samples as anchors
+        curr_class_mask = torch.zeros_like(repeated_labels)
+        for tc in target_labels:
+            curr_class_mask += (repeated_labels == tc)
+        curr_class_mask = curr_class_mask.view(-1).to(features)
+        loss_normal_features = curr_class_mask * loss_normal_features
+        loss_normal_features = torch.mean(loss_normal_features)
+
+        # For samix features
+        loss_samix_features = None
+        if samix_features is not None and lamb_ratio is not None and mixed_indices is not None:
+            mixed_prototypes = self.generate_samix_prototypes_slerp(repeated_labels, mixed_indices,
+                                                                    corresponding_prototypes, lamb_ratio)
+            # dr - samix samples
+            loss_samix_features = torch.sum(samix_features * mixed_prototypes, dim=1)  # [2 x bsz,]
+            loss_samix_features = 0.5 * ((loss_samix_features - 1) ** 2)
+            loss_samix_features = torch.mean(loss_samix_features)
+        return loss_normal_features + self.weight_samix_loss * loss_samix_features

src/utils/general_utils.py

+from __future__ import print_function
+
+import os
+import inspect
+import math
+import numpy as np
+import torch
+import random
+import torch.optim as optim
+from torch.utils.data import Dataset
+from src.datasets.tiny_imagenet_dataset import TinyImagenet
+from torchvision import transforms, datasets
+from src.datasets import Dataset
+import json
+
+
+def print_x(s):
+    caller_frame = inspect.stack()[1]
+    caller_file_path = caller_frame.filename
+    caller_file_name = os.path.basename(caller_file_path)
+    print(f"[{caller_file_name}] {s}")
+
+
+def get_train_datasets(opt, data_transform):
+    if opt.dataset == 'cifar10':
+        return datasets.CIFAR10(root=opt.data_folder,
+                                transform=data_transform,
+                                download=True)
+    elif opt.dataset == 'tiny-imagenet':
+        return TinyImagenet(root=opt.data_folder,
+                            transform=data_transform,
+                            download=True)
+    elif opt.dataset == 'cifar100':
+        return datasets.CIFAR100(root=opt.data_folder,
+                                 transform=data_transform,
+                                 download=True)
+    elif opt.dataset == 'custom':
+        return datasets.ImageFolder(root=opt.data_folder,
+                                    transform=data_transform)
+    else:
+        raise ValueError('Dataset not supported: {}'.format(opt.dataset))
+
+
+def get_val_datasets(opt, val_transform):
+    if opt.dataset == 'cifar10':
+        return datasets.CIFAR10(root=opt.data_folder,
+                                train=False,
+                                transform=val_transform)
+    elif opt.dataset == 'tiny-imagenet':
+        return TinyImagenet(root=opt.data_folder,
+                            transform=val_transform,
+                            train=False)
+    elif opt.dataset == 'cifar100':
+        return datasets.CIFAR100(root=opt.data_folder,
+                                 train=False,
+                                 transform=val_transform)
+    elif opt.dataset == 'custom':
+        return datasets.ImageFolder(root=opt.data_folder,
+                                    transform=val_transform)
+    else:
+        raise ValueError('Dataset not supported: {}'.format(opt.dataset))
+
+
+class TwoCropTransform:
+    """Create two crops of the same image"""
+
+    def __init__(self, transform):
+        self.transform = transform
+
+    def __call__(self, x):
+        return [self.transform(x), self.transform(x)]
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def adjust_learning_rate(args, optimizer, epoch):
+    lr = args.learning_rate
+    if args.cosine:
+        eta_min = lr * (args.lr_decay_rate ** 3)
+        lr = eta_min + (lr - eta_min) * (
+                1 + math.cos(math.pi * epoch / args.epochs)) / 2
+    else:
+        steps = np.sum(epoch > np.asarray(args.lr_decay_epochs))
+        if steps > 0:
+            lr = lr * (args.lr_decay_rate ** steps)
+
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def warmup_learning_rate(args, epoch, batch_id, total_batches, optimizer):
+    if args.warm and epoch <= args.warm_epochs:
+        p = (batch_id + (epoch - 1) * total_batches) / \
+            (args.warm_epochs * total_batches)
+        lr = args.warmup_from + p * (args.warmup_to - args.warmup_from)
+
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+
+def set_optimizer(opt, model):
+    optimizer = optim.SGD(model.parameters(),
+                          lr=opt.learning_rate,
+                          momentum=opt.momentum,
+                          weight_decay=opt.weight_decay)
+    return optimizer
+
+
+def save_model(model, optimizer, opt, epoch, save_file):
+    print_x('==> Saving...' + save_file)
+    state = {
+        'opt': opt,
+        'model': model.state_dict(),
+        'optimizer': optimizer.state_dict(),
+        'epoch': epoch,
+    }
+    torch.save(state, save_file)
+    del state
+
+
+def load_model(model, optimizer, save_file):
+    print_x('==> Loading...' + save_file)
+    loaded = torch.load(save_file)
+    model.load_state_dict(loaded['model'])
+    optimizer.load_state_dict(loaded['optimizer'])
+    del loaded
+    return model, optimizer
+
+
+def set_replay_samples(opt, model, prev_indices=None):
+    is_training = model.training
+    model.eval()
+    val_transform = transforms.Compose([
+        transforms.ToTensor(),
+    ])
+    val_dataset = get_train_datasets(opt, val_transform)
+    val_targets = val_dataset.targets
+    val_targets = np.array(val_targets)
+
+    if prev_indices is None:
+        prev_indices = []
+        observed_classes = list(range(0, opt.target_task * opt.cls_per_task))
+    else:
+        shrink_size = ((opt.target_task - 1) * opt.mem_size / opt.target_task)
+        if len(prev_indices) > 0:
+            unique_cls = np.unique(val_targets[prev_indices])
+            _prev_indices = prev_indices
+            prev_indices = []
+            for c in unique_cls:
+                mask = val_targets[_prev_indices] == c
+                size_for_c = shrink_size / len(unique_cls)
+                p = size_for_c - (shrink_size // len(unique_cls))
+                if random.random() < p:
+                    size_for_c = math.ceil(size_for_c)
+                else:
+                    size_for_c = math.floor(size_for_c)
+                prev_indices += torch.tensor(_prev_indices)[mask][torch.randperm(mask.sum())[:size_for_c]].tolist()
+        observed_classes = list(
+            range(max(opt.target_task - 1, 0) * opt.cls_per_task, (opt.target_task) * opt.cls_per_task))
+    if len(observed_classes) == 0:
+        return prev_indices
+
+    observed_indices = []
+    for tc in observed_classes:
+        observed_indices += np.where(val_targets == tc)[0].tolist()
+    val_observed_targets = val_targets[observed_indices]
+    val_unique_cls = np.unique(val_observed_targets)
+    selected_observed_indices = []
+    for c_idx, c in enumerate(val_unique_cls):
+        size_for_c_float = (
+                (opt.mem_size - len(prev_indices) - len(selected_observed_indices)) / (len(val_unique_cls) - c_idx))
+        p = size_for_c_float - ((opt.mem_size - len(prev_indices) - len(selected_observed_indices)) // (
+                len(val_unique_cls) - c_idx))
+        if random.random() < p:
+            size_for_c = math.ceil(size_for_c_float)
+        else:
+            size_for_c = math.floor(size_for_c_float)
+        mask = val_targets[observed_indices] == c
+        selected_observed_indices += torch.tensor(observed_indices)[mask][
+            torch.randperm(mask.sum())[:size_for_c]].tolist()
+    print_x(np.unique(val_targets[selected_observed_indices], return_counts=True))
+
+    model.is_training = is_training
+    output = prev_indices + selected_observed_indices
+    return output
+
+
+def save_config_file(opt, save_path):
+    dict_opt = vars(opt)
+    with open(save_path, 'w') as f:
+        json.dump(dict_opt, f, indent=4)
+        f.close()
+
+
+def create_folders(dirs):
+    if not os.path.isdir(dirs):
+        os.makedirs(dirs)
+
+
+def get_feat_dim_projection(dataset_name):
+    if dataset_name == Dataset.TINY_IMAGENET.value:
+        return 256
+    return 128
+
+
+def write_eval_acc(accuracy_single_task_path, best_acc, val_acc,
+                   dict_acc_classes, best_task_acc, val_task_acc):
+    with open(accuracy_single_task_path, 'w') as f:
+        out = 'Best class-il accuracy: {:.2f}, task-il accuracy: {:.2f}\n'.format(best_acc, best_task_acc)
+        out += '{:.2f} {:.2f}'.format(val_acc, val_task_acc)
+        print_x(out)
+        out += '\n'
+        for k, v in dict_acc_classes.items():
+            print_x(v)
+            out += f'class: {k} - acc: {v:.2f}\n'
+        f.write(out)
+
+
+def compute_forgetting(final_dict_acc_tasks, maximum_acc):
+    list_subtract = [maximum_acc[i][1] - final_dict_acc_tasks[i] for i in final_dict_acc_tasks.keys()][:-1]
+    return sum(list_subtract) / len(list_subtract)

src/utils/metrics.py

+class TaskAccuracy:
+
+    def __init__(self, task_index, num_epochs):
+        self.task_index = task_index
+        self.num_epochs = num_epochs
+
+        # class-il
+        self.last_acc = 0.0
+        self.best_acc = (0, 0.0)
+        self.dict_acc_classes = {}
+
+        # task-il
+        self.task_acc = 0.0
+        self.best_task_acc = (0, 0.0)

src/utils/mix_views_generation.py

+import numpy as np
+import torch
+
+
+def generate_samix(data, alpha=25.0):
+    lamb = np.random.beta(alpha, alpha)
+    indices = get_perm(data.size(0))
+    data_perm = data[indices]
+    mixed_views = data * lamb + (1 - lamb) * data_perm
+
+    return mixed_views, lamb, indices
+
+
+def get_perm(l):
+    perm = torch.randperm(l)
+    while torch.all(torch.eq(perm, torch.arange(l))):
+        perm = torch.randperm(l)
+    return perm

src/utils/prototypes.py

+import random
+import math
+from typing import Optional
+import numpy as np
+
+
+def pedcc_generation(
+        n: int, k: int = None, seed: Optional[int] = None
+) -> np.ndarray:
+    def pedcc_frame(n: int, k: int = None) -> np.ndarray:
+        assert 0 < k <= n + 1
+        zero = [0] * (n - k + 1)
+        u0 = [-1][:0] + zero + [-1][0:]
+        u1 = [1][:0] + zero + [1][0:]
+        u = np.stack((u0, u1)).tolist()
+        for i in range(k - 2):
+            c = np.insert(u[len(u) - 1], 0, 0)
+            for j in range(len(u)):
+                p = np.append(u[j], 0).tolist()
+                s = len(u) + 1
+                u[j] = math.sqrt(s * (s - 2)) / (s - 1) * np.array(p) - 1 / (
+                        s - 1
+                ) * np.array(c)
+            u.append(c)
+        return np.array(u)
+
+    U = pedcc_frame(n=n, k=k)
+    r = np.random.RandomState(seed)
+    while True:
+        try:
+            noise = r.rand(n, n)  # [0, 1)
+            V, _ = np.linalg.qr(noise)
+            break
+        except np.linalg.LinAlgError:
+            continue
+    points = np.dot(U, V)
+    return points
+
+
+def generate(
+        d: int,
+        k: int,
+        path_save_prototype: str,
+        seed: Optional[int] = None,
+) -> np.ndarray:
+    """Generate k evenly distributed R^n points in a unit (n-1)-hypersphere
+    Args:
+        d (int): dimension of the Euclidean space
+        k (int): number of points to generate
+        method (str): method to generate the points. Defaults to "simplex".
+        path_save_prototype (str): filepath to save the points.
+        seed (int, optional): seed for the random number generator. Defaults to None.
+
+    Returns:
+        np.ndarray: k evenly distributed points in a unit (n-1)-hypersphere
+    >>> generate(2, 3, method="simplex")
+    array([[ 0.        ,  0.        ],
+           [ 0.70710678,  0.70710678],
+    """
+    assert k < d + 2
+    if seed is None or not (isinstance(seed, int) and 0 <= seed < 2 ** 32):
+        seed = random.randrange(2 ** 32)
+    points = pedcc_generation(d, k, seed=seed)
+    if path_save_prototype:
+        np.save(path_save_prototype, points)
+    return points

src/utils/transforms_utils.py

+from __future__ import print_function
+from torchvision import transforms
+
+
+def get_transforms(opt):
+    if opt.dataset == 'cifar10':
+        mean = (0.4914, 0.4822, 0.4465)
+        std = (0.2023, 0.1994, 0.2010)
+    elif opt.dataset == 'cifar100':
+        mean = (0.5071, 0.4867, 0.4408)
+        std = (0.2675, 0.2565, 0.2761)
+    elif opt.dataset == 'tiny-imagenet':
+        mean = (0.4802, 0.4480, 0.3975)
+        std = (0.2770, 0.2691, 0.2821)
+    else:
+        raise ValueError('dataset not supported: {}'.format(opt.dataset))
+    normalize = transforms.Normalize(mean=mean, std=std)
+    train_transform = transforms.Compose([
+        transforms.Resize(size=(opt.size, opt.size)),
+        transforms.RandomResizedCrop(size=opt.size, scale=(0.1 if opt.dataset == 'tiny-imagenet' else 0.2, 1.)),
+        transforms.RandomHorizontalFlip(),
+        transforms.RandomApply([
+            transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
+        ], p=0.8),
+        transforms.RandomGrayscale(p=0.2),
+        transforms.RandomApply([transforms.GaussianBlur(kernel_size=opt.size // 20 * 2 + 1, sigma=(0.1, 2.0))],
+                               p=0.5 if opt.size > 32 else 0.0),
+        transforms.ToTensor(),
+        normalize,
+    ])
+    val_transform = transforms.Compose([
+        transforms.Resize(size=(opt.size, opt.size)),
+        transforms.ToTensor(),
+        normalize,
+    ])
+    return train_transform, val_transform