main.py

# -*- coding: utf-8 -*-

from dataclasses import dataclass, field
from typing import Any, Dict, List, Tuple

import numpy as np
import hydra
import math
import wandb
import copy
from omegaconf import DictConfig
import torch
import torch.nn as nn
import torch.optim as optim
from absl import logging
from timm.utils.model_ema import ModelEmaV2
from torch.utils.data import DataLoader
from torchmetrics import Accuracy
from info_nce import InfoNCE
from tqdm import tqdm
from torch.distributions.multivariate_normal import MultivariateNormal

from libml.core import App, Conf, DataModule, Module, ModuleConf
from libml.data import build_imagenet_transform
from libml.model.pet import Adapter, Conv2dAdapter, KVLoRA, Prefix, PromptPool, KVLoRAPool, KVVeRAPool, Prompt
from libml.model.utils import freeze, unfreeze, get_knn_acc, get_cifar100_mapping, BarlowTwins, EWC, SI, orth_loss, target2onehot, optimise_ridge_parameter


from libml.model.backbone.vit_adam import MultiBranchCosineIncrementalNet


@dataclass
class MyModuleConf(ModuleConf):
    adapt_blocks: List[int] = field(default_factory=lambda: [0, 1, 2, 3, 4])
    pet_cls: str = "LAE-Adapter"  # enum: Adapter, LoRA, Prefix
    pet_kwargs: Dict[str, Any] = field(default_factory=lambda: {})
    query_pet_cls: str = None
    query_pet_kwargs: Dict[str, Any] = field(default_factory=lambda: {})
    num_emas: int = 1
    ema_decay: float = 0.9999
    num_freeze_epochs: int = 3
    eval_only_emas: bool = False
    debug: bool = False
    prune: int = 0
    mask_logits: bool = True
    regularizer: str = False
    reg_lambda: float = 0.1


@dataclass
class MyConf(Conf):
    module: MyModuleConf = MyModuleConf()


class MyDataModule(DataModule):
    def train_transform(self):
        return build_imagenet_transform(train=True, norm=False)

    def val_transform(self):
        return build_imagenet_transform(train=False, norm=False)


class MyModule(Module):
    pets: nn.Module  # online pet modules
    pets_emas: nn.ModuleList  # offline pet modules
    original_vit: nn.Module
    train_acc: Accuracy
    loss_fn: nn.Module

    def setup_model(self):
        """
        Setup model and pets
        :return:
        """
        super().setup_model()

        if self.cfg.pet_cls.startswith('L2P') or self.cfg.pet_cls.startswith('DP') \
                or self.cfg.pet_cls == 'LLS' or self.cfg.pet_cls == 'CODA' or self.cfg.pet_cls.startswith('ADAM') or self.cfg.pet_cls == 'HIDE' or self.cfg.pet_cls == 'OA':
            if self.cfg.model.freeze:
                if self.query_function is not None and self.query_function != 'model':
                    # all parameters are frozen for original vit model
                    for p in self.query_function.parameters():
                        p.requires_grad = False

                # freeze args.freeze[blocks, patch_embed, cls_token] parameters
                if self.cfg.pet_cls == 'CODA':
                    freeze_comp = tuple([f'feat.{c}' for c in self.cfg.model.freeze])
                else:
                    freeze_comp = tuple(self.cfg.model.freeze)
                for n, p in self.model.backbone.named_parameters():
                    if n.startswith(freeze_comp):
                        p.requires_grad = False
        elif self.cfg.pet_cls.startswith('LAE') or self.cfg.pet_cls.endswith('LoRA'):
            if getattr(self, "pets_emas", None) is None:
                freeze(self.model.backbone)
                self.pets_emas = nn.ModuleList([])
                self.pets = self.create_pets_vit(self.cfg.pet_cls, self.cfg.pet_kwargs)
                logging.info(f"==> pets:\n{self.pets}")
            elif len(self.pets_emas) < self.cfg.num_emas:
                idx = len(self.pets_emas)
                ema = ModelEmaV2(self.pets, decay=self.cfg.ema_decay)
                self.pets_emas.append(ema)
            self.attach_pets(self.pets)
            if getattr(self, 'query_function', None) is not None and \
                    self.cfg.model.query.endswith('lae') and \
                    self.cfg.model.query_kwargs.adapt_query:
                self.query_pets = self.create_pets_vit(self.cfg.query_pet_cls, self.cfg.query_pet_kwargs)
                self.attach_pets(self.query_pets, to='query')
        elif self.cfg.pet_cls == 'None':
            freeze(self.model.backbone)
        if self.cfg.model.head == 'DynamicNMCHead' and not self.cfg.pet_cls.startswith('ADAM'):
            freeze(self.model.head)
        n_parameters = sum(p.numel() for p in self.model.backbone.parameters() if p.requires_grad)
        logging.info(f'number of backbone params: {n_parameters}')
        n_parameters = sum(p.numel() for p in self.model.head.parameters() if p.requires_grad)
        logging.info(f'number of head params: {n_parameters}')
        self.train_acc = Accuracy().to(self.device)
        self.loss_fn = nn.CrossEntropyLoss()
        if self.cfg.model.adapt_query:
            if self.cfg.model.query_train_kwargs.loss == 'info_nce':
                self.query_loss_fn = InfoNCE(negative_mode='paired')
                self.query_transforms = build_imagenet_transform(train=True, query=True)
            elif self.cfg.model.query_train_kwargs.loss == 'barlow_twins':
                self.query_loss_fn = BarlowTwins()
                self.query_transforms = build_imagenet_transform(train=True, query=True)
            elif self.cfg.model.query_train_kwargs.loss == 'cross_entropy':
                self.query_loss_fn = nn.CrossEntropyLoss()
                self.query_transforms = None
            else:
                raise NotImplementedError

        if self.cfg.regularization == 'ewc' and self.current_task > 0:
            self.regularizer = EWC(self.model, self.train_loader)
        if self.cfg.regularization == 'si':
            if self.current_task == 0:
                self.regularizer = SI(self.model, self.cfg.reg_lambda)
            else:
                self.regularizer.update_omega()

    #@property
    def create_pets_vit(self, pet_cls, pet_kwargs):
        """
        Create PET modules for Vision Transformer
        :param pet_cls: PET class name
        :param pet_kwargs: PET class arguments
        :return: nn.ModuleList of PET modules
        """
        assert not (pet_cls.startswith('L2P') or pet_cls.startswith('DP')), "Not implemented PET"

        n = len(self.cfg.adapt_blocks)
        embed_dim = self.model.backbone.embed_dim

        kwargs = dict(**pet_kwargs)
        if pet_cls.endswith("Adapter"):
            kwargs["embed_dim"] = embed_dim
            kwargs.pop('embed_key')
            return nn.ModuleList([Adapter(**kwargs) for _ in range(n)])

        if pet_cls.startswith("Prompt") or pet_cls.endswith("Prompt"):
            assert n == 1, "Only one prompt is supported"
            kwargs.pop('embed_key')
            return nn.ModuleList([Prompt(**kwargs)])

        if pet_cls.startswith("Pool"):
            kwargs["in_features"] = embed_dim
            kwargs["out_features"] = embed_dim
            kwargs.pop('pull_constraint')
            kwargs.pop('pull_constraint_coeff')
            kwargs.pop('push_constraint')
            kwargs.pop('push_constraint_coeff')
            if isinstance(pet_kwargs.pool_size, int):
                pool_size = [pet_kwargs.pool_size] * n
            else:
                pool_size = pet_kwargs.pool_size
            assert len(pool_size) == n
            pets = nn.ModuleList()
            for l in range(n):
                kwargs["pool_size"] = pool_size[l]
                if pet_cls.endswith("LoRA"):
                    pets.append(KVLoRAPool(**kwargs))
                elif pet_cls.endswith("VeRA"):
                    A = torch.zeros((kwargs['rank'], kwargs['in_features'])).to(self.device)
                    B = torch.zeros((kwargs['out_features'], kwargs['rank'])).to(self.device)
                    nn.init.kaiming_uniform_(A, a=math.sqrt(5))
                    nn.init.kaiming_uniform_(B, a=math.sqrt(5))
                    kwargs["A"] = A
                    kwargs["B"] = B
                    pets.append(KVVeRAPool(**kwargs))
                else:
                    raise NotImplementedError
            return pets

        if pet_cls.endswith("LoRA"):
            kwargs["in_features"] = embed_dim
            kwargs["out_features"] = embed_dim
            kwargs.pop('embed_key')
            return nn.ModuleList([KVLoRA(**kwargs) for _ in range(n)])

        kwargs["dim"] = embed_dim
        kwargs.pop('embed_key')
        return nn.ModuleList([Prefix(**kwargs) for i in range(n)])

    def create_pets(self):
        """
        Create PET modules
        :return: nn.ModuleList of PET modules
        """
        return self.create_pets_vit

    def attach_pets_vit(self, pets: nn.ModuleList, to):
        """
        Attach PET modules to Vision Transformer
        :param pets: PET modules
        :param to: 'model' or 'query'
        :return:
        """
        assert not (self.cfg.pet_cls.startswith('L2P') or self.cfg.pet_cls.startswith('DP')), "Not implemented PET"

        if self.cfg.pet_cls.endswith("Adapter"):
            for i, b in enumerate(self.cfg.adapt_blocks):
                if to == 'model':
                    self.model.backbone.blocks[b].attach_adapter(attn=pets[i])
                elif to == 'query':
                    self.query_function.backbone.blocks[b].attach_adapter(attn=pets[i])
                else:
                    raise NotImplementedError
            return

        if self.cfg.pet_cls.endswith("LoRA") or self.cfg.pet_cls.endswith("VeRA"):
            for i, b in enumerate(self.cfg.adapt_blocks):
                if to == 'model':
                    self.model.backbone.blocks[b].attn.attach_adapter(qkv=pets[i])
                elif to == 'query':
                    self.query_function.backbone.blocks[b].attach_adapter(qkv=pets[i])
                else:
                    raise NotImplementedError
            return

        if self.cfg.pet_cls.startswith("Prompt") or self.cfg.pet_cls.endswith("Prompt"):
            for pet in self.pets:
                if to == 'model':
                    self.model.backbone.attach_prompt(pet)
                elif to == 'query':
                    self.query_function.backbone.attach_prompt(pet)
                else:
                    raise NotImplementedError

        for i, b in enumerate(self.cfg.adapt_blocks):
            if to == 'model':
                self.model.backbone.blocks[b].attn.attach_prefix(pets[i])
            elif to == 'query':
                self.query_function.backbone.blocks[b].attach_prefix(pets[i])
            else:
                raise NotImplementedError

    def attach_pets(self, pets: nn.ModuleList, to='model'):
        """
        Attach PET modules to Vision Transformer
        :param pets: PET modules
        :param to: 'model' or 'query'
        :return:
        """
        return self.attach_pets_vit(pets, to=to)

    def filter_state_dict(self, state_dict):
        """
        Filter state_dict to remove PET modules
        :param state_dict: Model state_dict
        :return: Filtered state_dict
        """
        ps = ("model_wrap.", "model.backbone.")
        return {k: v for k, v in state_dict.items() if not k.startswith(ps)}

    def forward(self, input):
        """
        Forward pass
        :param input: input data
        :return: model output
        """
        return super().forward(input)

    def pre_train_epoch(self):
        """
        Pre-train epoch
        :return:
        """
        self.train_acc.reset()

        if hasattr(self.cfg.model.backbone_kwargs, 'load_avg_prompt') and self.cfg.model.backbone_kwargs.load_avg_prompt:
            self.model.backbone.prompt.prompt.data = self.model.backbone.prompt.avg_patches[self.current_task].to(self.model.backbone.prompt.prompt.device)
            logging.info("===> Load avg prompt for task {}".format(self.current_task))

        if self.cfg.pet_cls == 'CODA' or self.cfg.pet_cls.startswith('ADAM'):
            self.model.backbone.train_mode = True

        if self.cfg.adapt_tasks != 'all':
            if isinstance(self.cfg.adapt_tasks, int):
                self.cfg.adapt_tasks = list(range(self.cfg.adapt_tasks+1))
            if self.current_task not in self.cfg.adapt_tasks:
                if hasattr(self.model.backbone, 'prompt'):
                    freeze(self.model.backbone.prompt)
                elif hasattr(self, 'pets'):
                    freeze(self.pets)
                logging.info("===> Freeze pets for task {}".format(self.current_task))
                return

        if self.cfg.num_freeze_epochs == "alternating":
            assert self.cfg.pet_cls != 'LLS', "LLS does not support alternating freezing"
            if self.current_epoch % 2 == 0:
                if hasattr(self.model.backbone, 'prompt'):
                    freeze(self.model.backbone.prompt)
                elif hasattr(self, 'pets'):
                    freeze(self.pets)
                unfreeze(self.model.head)
                logging.info("===> Freeze pets, unfreeze head for epoch {}".format(self.current_epoch))
                return
            else:
                if hasattr(self.model.backbone, 'prompt'):
                    unfreeze(self.model.backbone.prompt)
                elif hasattr(self, 'pets'):
                    unfreeze(self.pets)
                freeze(self.model.head)
                logging.info("===> Unfreeze pets, freeze head for epoch {}".format(self.current_epoch))
                return

        if self.cfg.pet_cls == 'LLS':
            if self.current_epoch == 0 and self.cfg.num_freeze_epochs > 0:
                for prompt in self.model.backbone.prompt.prompt_pool:
                    prompt.requires_grad_(False)
                logging.info("===> Freeze pets")
            elif self.current_epoch == self.cfg.num_freeze_epochs:
                if self.cfg.model.backbone_kwargs.lls_freeze:
                    self.model.backbone.prompt.prompt_pool[-1].requires_grad_(True)
                else:
                    for prompt in self.model.backbone.prompt.prompt_pool:
                        prompt.requires_grad_(True)
                logging.info("===> Unfreeze pets")
            return

        if (not self.cfg.freeze_first_task and self.current_task == 0) or self.cfg.num_freeze_epochs < 1:
            return

        if self.current_epoch == 0:
            if self.cfg.pet_cls.startswith('L2P') or self.cfg.pet_cls.startswith('DP'):
                freeze(self.model.backbone.prompt)
            else:
                freeze(self.pets)
            logging.info("===> Freeze pets")

        if self.current_epoch == self.cfg.num_freeze_epochs:
            if self.cfg.pet_cls.startswith('Pool') and self.cfg.pet_kwargs.embed_key == 'oracle':
                unfreeze(self.pets[self.current_task])
            if self.cfg.pet_cls.startswith('L2P') or self.cfg.pet_cls.startswith('DP'):
                unfreeze(self.model.backbone.prompt)
            elif hasattr(self, 'pets'):
                unfreeze(self.pets)
            logging.info("===> Unfreeze pets")

    def train_query_step(self, batch, batch_idx):
        """
        Train query function
        :param batch: Input batch
        :param batch_idx: Batch index
        :return:
        """
        output = self.query_function(batch[0])
        if self.query_transforms is not None:
            pos_samples = self.query_transforms(batch[0])
            pos_output = self.query_function(pos_samples)
        if self.cfg.model.query.endswith('l2p'):
            features = output['x']
            pos = pos_output['x']
        elif self.cfg.model.query.endswith('hide'):
            features = output['logits']
        else:
            features = output
            pos = pos_output
        if self.cfg.model.query_train_kwargs.loss != 'cross_entropy':
            if self.cfg.model.query_train_kwargs.embedding_key == 'cls':
                features = features[:, 0, :]
                pos = pos[:, 0, :]
            elif self.cfg.model.query_train_kwargs.embedding_key == 'mean':
                features = features.mean(dim=1)
                pos = pos.mean(dim=1)
            elif self.cfg.model.query_train_kwargs.embedding_key == 'max':
                features = features.max(dim=1)[0]
                pos = pos.max(dim=1)[0]
            else:
                raise NotImplementedError

        if self.cfg.model.query_train_kwargs.loss == 'info_nce':
            neg_idx = {}
            for c in torch.unique(batch[1]):
                neg_idx[c.item()] = (batch[1] != c).nonzero().squeeze(dim=-1)
            query = features
            neg = torch.cat([features[neg_idx[l.item()]][:5].unsqueeze(0) for l in batch[1]], dim=0)

            loss = self.query_loss_fn(query, pos, neg)
        elif self.cfg.model.query_train_kwargs.loss == 'cross_entropy':
            logits = features[:, self.datamodule.num_old_classes:]
            target = batch[1] - self.datamodule.num_old_classes
            loss = self.query_loss_fn(logits, target)
        else:
            loss = self.query_loss_fn(features, pos)
        return loss

    def eval_query_epoch(self, loader):
        """
        Evaluate query function
        :param loader: DataLoader
        :return:
        """
        features, labels = [], []
        for batch in loader:
            output = self.query_function(batch[0])
            if self.cfg.model.query.endswith('l2p'):
                features.append(output['x'].cpu())
            elif self.cfg.model.query.endswith('hide'):
                features.append(output['logits'].cpu())
            else:
                features.append(output.cpu())
            labels.append(batch[1].cpu())

        labels = torch.cat(labels, dim=0)

        if self.cfg.model.query.endswith('hide'):
            features = torch.cat(features, dim=0)
            preds = features.argmax(dim=1)
            acc = (preds == labels).float().mean().item() * 100
            return acc
        else:
            cls_feature = torch.cat(features, dim=0)[:, 0, :]
            fine_acc, coarse_acc = get_knn_acc(cls_feature, labels, k=25)
            return fine_acc, coarse_acc

    def train_step(
            self, batch, batch_idx, mode='train'
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """
        Train step
        :param batch: Input batch
        :param batch_idx: Batch index
        :param mode: 'train' or 'memtrain'
        :return: loss and accuracy
        """
        if self.datamodule.memory is not None and (self.datamodule.cfg.memory.algo_name.endswith('dumb') and mode == 'memtrain') or not self.cfg.mask_logits or mode == 'bictrain':
            num_old_classes = 0
        else:
            num_old_classes = self.datamodule.num_old_classes
        if self.cfg.pet_cls.startswith("L2P") or self.cfg.pet_cls.startswith("DP"):
            if self.cfg.model.backbone_kwargs.embedding_key == 'oracle':
                self.model.backbone.prompt.current_labels = batch[1]
            with torch.no_grad():
                if self.query_function is not None:
                    if self.query_function == 'model':
                        output = self((batch[0], 'query'))
                    else:
                        output = self.query_function(batch[0])
                    cls_features = output['pre_logits']
                else:
                    cls_features = None
            output = self((batch[0], cls_features, None))
            logits = output['logits']
        elif self.cfg.pet_cls == 'HIDE':
            with torch.no_grad():
                if self.query_function is not None:
                    if self.query_function == 'model':
                        output = self((batch[0], 'query'))
                    else:
                        output = self.query_function(batch[0])
                    logits = output['logits']

                    if self.cfg.pet_kwargs.train_mask:
                        prompt_id = torch.max(logits, dim=1)[1]
                        # translate cls to task_id
                        prompt_id = torch.tensor([v.item()//self.datamodule.cfg.num_increment_classes for v in prompt_id],
                                                 device=self.device).unsqueeze(-1)
                    else:
                        prompt_id = None
                else:
                    raise NotImplementedError("original model is None")
            output = self.model((batch[0], self.current_task, prompt_id, None, True, False,
                           self.cfg.pet_kwargs.prompt_momentum))
            logits = output['logits']
        elif self.cfg.pet_cls == 'LLS' or self.cfg.pet_cls == 'CODA':
            if self.cfg.model.backbone_kwargs.embedding_key.endswith('oracle') or self.cfg.model.backbone_kwargs.embedding_key.startswith('masked'):
                self.model.backbone.prompt.current_labels = batch[1]
            if self.query_function is not None and self.query_function != 'model':
                cls_features = self.query_function(batch[0])['pre_logits']
                output = self((batch[0], cls_features))
            else:
                output = self(batch[0])
            logits = output['logits']
        else:
            if self.cfg.pet_kwargs.embed_key == "cls":
                cls_features = self.query_function(batch[0])
                for pet in self.pets:
                    pet.cls_features = cls_features
            output = None
            if self.cfg.pet_cls.startswith('Pool'):
                for pet in self.pets:
                    pet.current_labels = batch[1]
            if self.cfg.model.head != "DynamicSimpleHead":
                logits = self(batch[0])
                if isinstance(logits, dict):
                    logits = logits['logits']
            else:
                logits = self(batch[0])
                logits = logits["logits"] if isinstance(logits, dict) else logits
        self.train_acc.update(logits, batch[1])
        logits = logits[:, num_old_classes:]
        target = batch[1] - num_old_classes
        assert (target >= 0).all(), f"target: {target}"
        if num_old_classes == 0 and False:
            batch_size, num_classes = logits.size()
            task_indices = batch[1] // self.datamodule.cfg.num_increment_classes
            start_class_indices = task_indices * self.datamodule.cfg.num_increment_classes
            start_class_indices_list = start_class_indices.tolist()
            for i in range(batch_size):
                logits[i, :start_class_indices_list[i]] = float("-inf")

        loss = self.loss_fn(logits, target)
        if not (self.cfg.pet_cls.startswith("LAE") or self.cfg.pet_cls.startswith("ADAM") or self.cfg.pet_cls.startswith(
                "Prompt") or self.cfg.pet_cls == 'CODA') and self.cfg.pet_kwargs.pull_constraint:
            if isinstance(output, dict):
                loss = loss - self.cfg.pet_kwargs.pull_constraint_coeff * output['reduce_sim']
            else:
                p_sim, n_sim = 0, 0
                for pet in self.pets:
                    p_sim += pet.reduce_sim
                    n_sim += pet.reduce_nkey_sim
                loss = loss - self.cfg.pet_kwargs.pull_constraint_coeff * p_sim.to(loss.device)
                if self.cfg.pet_kwargs.push_constraint:
                    loss = loss + self.cfg.pet_kwargs.push_constraint_coeff * n_sim.to(loss.device)
        if self.cfg.pet_cls == 'CODA':
            loss = loss + output['prompt_loss'].sum()
        if self.cfg.pet_cls.startswith("L2P") and self.cfg.pet_kwargs.embed_key == 'generate':
            prompts = output['prompted_embedding'][:, 1:output['total_prompt_len'] + 1, :].mean(dim=1)
            neg_sim = []
            pos_sim = []
            for c in torch.unique(target):
                pos = prompts[target == c]
                neg = prompts[target != c]
                pos = pos / pos.norm(dim=-1, keepdim=True)
                neg = neg / neg.norm(dim=-1, keepdim=True)
                neg_sim.append(torch.mm(pos, neg.t()).mean().unsqueeze(0))
                pos_sim.append(torch.mm(pos, pos.t()).mean().unsqueeze(0))
            neg_sim = torch.cat(neg_sim, dim=0)
            pos_sim = torch.cat(pos_sim, dim=0)
            loss = loss + self.cfg.pet_kwargs.push_constraint_coeff * neg_sim.mean() \
                   - self.cfg.pet_kwargs.pull_constraint_coeff * pos_sim.mean()
        if self.cfg.regularization == 'ewc' and self.current_task > 0:
            loss = loss + self.regularizer.penalty(self.model) * self.cfg.reg_lambda
        if self.cfg.regularization == 'si' and self.current_task > 0:
                loss = loss + self.regularizer.penalty(self.model) * self.cfg.reg_lambda
        if hasattr(self.cfg.model.backbone_kwargs, "prompt_ema") and self.cfg.model.backbone_kwargs.prompt_ema:
            if batch_idx % self.cfg.model.backbone_kwargs.prompt_ema_freq == 0:
                self.model.backbone.prompt.interpolate()
        loss_dict = {"loss": loss, "acc": self.train_acc.compute() * 100}
        return loss, loss_dict

    def post_train_step(self):
        """
        Post-train step
        :return:
        """
        if not (self.cfg.pet_cls.startswith('L2P') or self.cfg.pet_cls.startswith('DP') or self.cfg.pet_cls.startswith('ADAM') or
                self.cfg.pet_cls.startswith('LLS') or self.cfg.pet_cls == "CODA" or self.cfg.pet_cls == 'None'
                or self.cfg.pet_cls == 'HIDE' or self.cfg.pet_cls == 'OA'):
            for idx, ema in enumerate(reversed(self.pets_emas)):
                if idx == 0:  # the last one
                    ema.update(self.pets)
                else:
                    ema.update(self.pets_emas[idx - 1])

        if self.cfg.pet_cls == "CODA":
            self.model.backbone.train_mode = False

    def eval_epoch(self, loader):
        """
        Evaluate epoch
        :param loader: DataLoader
        :return: tuple (pred_on, pred_off, pred_ens, gts, sim, prompts, labels)
        """
        task_ranges = []
        n_tasks = min(self.current_task + 1, self.datamodule.cfg.num_classes//self.datamodule.cfg.num_increment_classes)
        for t in range(n_tasks):
            s = task_ranges[-1][-1] + 1 if task_ranges else 0
            e = s + self.datamodule.num_classes_of(t)
            task_ranges.append(list(range(s, e)))

        pred_on, pred_off, pred_ens, gts = [], [], [], []
        prompts, labels, sim, logits = [], [], [], []
        for _, batch in enumerate(loader):
            if self.cfg.debug and _ > 1:
                break
            input, target = batch[:2]
            if self.cfg.pet_cls.startswith('Pool'):
                for pet in self.pets:
                    pet.current_task = batch[2]
                if self.cfg.pet_cls.startswith('Pool'):
                    for pet in self.pets:
                        pet.current_labels = batch[1]
                    if self.pets_emas is not None:
                        for pet_ema in self.pets_emas:
                            for pet in pet_ema.module:
                                pet.current_labels = batch[1]
            if self.cfg.pet_cls.startswith("L2P") or self.cfg.pet_cls.startswith("DP") or self.cfg.pet_cls == 'CODA' or self.cfg.pet_cls == 'HIDE':
                if self.cfg.model.backbone_kwargs.embedding_key == 'oracle':
                    self.model.backbone.prompt.current_labels = batch[1]
                if self.cfg.pet_cls == 'HIDE':
                    with torch.no_grad():
                        if self.query_function is not None:
                            if self.query_function == 'model':
                                output = self((batch[0], 'query'))
                            else:
                                output = self.query_function(batch[0])
                            q_logits = output['logits']

                            if self.cfg.pet_kwargs.train_mask:
                                prompt_id = torch.max(q_logits, dim=1)[1]
                                # translate cls to task_id
                                prompt_id = torch.tensor(
                                    [v.item() // self.datamodule.cfg.num_increment_classes for v in prompt_id],
                                    device=self.device).unsqueeze(-1)
                            else:
                                prompt_id = None
                        else:
                            raise NotImplementedError("original model is None")
                    output = self.model((batch[0], -1, prompt_id, None, False, False, 0))
                else:
                    if self.query_function is not None and self.cfg.pet_cls != 'CODA':
                        if self.query_function == 'model':
                            output = self((batch[0], 'query'))
                        else:
                            output = self.query_function(batch[0])
                        cls_features = output['pre_logits']
                    else:
                        cls_features = None
                    output = self((input, cls_features, None))
                if self.cfg.pet_cls.startswith("L2P"):
                    prompts.append(output['prompted_embedding'][:, :output['total_prompt_len'], :].to('cpu'))
                elif self.cfg.pet_cls == 'CODA':
                    prompts.append(output['prompt_list'].mean(dim=0).to('cpu'))
                elif self.cfg.pet_cls == 'DP' or self.cfg.pet_cls == 'HIDE':
                    p = output['batched_prompt'].mean(dim=0).mean(dim=1).view(batch[0].size(0), -1, 768).to('cpu')
                    prompts.append(p)
                labels.append(batch[1].to('cpu'))
                output = output['logits']
                pred_on.append(output.argmax(dim=1))
                pred_off.append(output.argmax(dim=1))
            elif self.cfg.pet_cls == "LLS":
                if self.cfg.model.backbone_kwargs.embedding_key.endswith('oracle') or self.cfg.model.backbone_kwargs.embedding_key.startswith('masked'):
                    self.model.backbone.prompt.current_labels = batch[1]
                output = self(input)
                labels.append(batch[1].to('cpu'))
                sim.append(output['sim'].to('cpu'))
                output = output['logits']
                pred_on.append(output.argmax(dim=1))
                pred_off.append(output.argmax(dim=1))
            else:
                if self.cfg.pet_kwargs.embed_key == "cls":
                    cls_features = self.query_function(batch[0])
                    for pet in self.pets:
                        pet.cls_features = cls_features
                output, bs = self(input), input.shape[0]
                if isinstance(output, dict):
                    output = output['logits']
                logits.append(output.to('cpu'))
                labels.append(batch[1].to('cpu'))
                pred_on.append(output.argmax(dim=1))
                output_emas = [output.softmax(dim=1)]
                if hasattr(self, 'pets_emas'):
                    for ema in self.pets_emas:
                        self.attach_pets(ema.module)
                        output_emas.append(self(input).softmax(dim=1))

                    for oe in output_emas[1:]:
                        pred_off.append(oe.argmax(dim=1))
                        break

                    if self.cfg.eval_only_emas and len(output_emas) > 1:
                        output_emas = output_emas[1:]
                    self.attach_pets(self.pets)
                output = torch.stack(output_emas, dim=-1).max(dim=-1)[0]

            self.val_acc.update(output, target)
            for t in batch[2].long().unique().tolist():
                sel = batch[2] == t
                self.val_task_accs[t].update(output[sel], target[sel])
                t_range = task_ranges[t]
                output_local = output[sel][:, t_range]
                target_local = target[sel] - t_range[0]
                self.val_task_local_accs[t].update(output_local, target_local)

            pred_ens.append(output.argmax(dim=1))
            gts.append(target)
        pred_on = torch.cat(pred_on)
        labels = torch.cat(labels) if len(labels) else None
        prompts = torch.cat(prompts) if len(prompts) else None
        logits = torch.cat(logits) if len(logits) else None
        if self.cfg.pet_cls == 'LLS':
            prompt_pool = self.model.backbone.prompt.prompt_pool
            prompts = torch.cat([param for param in prompt_pool], dim=1)
        sim = torch.cat(sim) if len(sim) else None
        return pred_on, pred_off, pred_ens, gts, sim, prompts, labels

    def post_eval_epoch(self, result):
        """
        Post-evaluate epoch
        :param result: Evaluation result
        :return:
        """
        super().post_eval_epoch()
        if self.cfg.pet_cls.startswith('Pool'):
            for pet in self.pets:
                pet.log_key_stats = False
                pet.mode = 'train'
        if self.cfg.pet_cls == "L2P" and self.cfg.model.backbone_kwargs.prompt_pool == 'flex':
            self.model.backbone.prompt.mode_train = True
        if self.current_task == 0 or not self.cfg.debug:
            return

    def adam_train(self, train_loader, test_loader, train_loader_for_protonet):
        """
        Train the model with ADAM.
        :param train_loader: Training data loader.
        :param test_loader: Testing data loader.
        :param train_loader_for_protonet: Training data loader for protonet.
        :return:
        """

        if self.current_task == 0 and not self.cfg.pet_cls.endswith('SimpleCIL'):

            # Freeze the parameters for ViT.
            total_params = sum(p.numel() for p in self.model.backbone.parameters())
            print(f'{total_params:,} total parameters.')
            total_trainable_params = sum(
                p.numel() for p in self.model.backbone.parameters() if p.requires_grad)
            print(f'{total_trainable_params:,} training parameters.')

            # if some parameters are trainable, print the key name and corresponding parameter number
            if total_params != total_trainable_params:
                for name, param in self.model.backbone.named_parameters():
                    if param.requires_grad:
                        print(name, param.numel())

            if self.cfg.optimizer.algo == 'SGD':
                optimizer = optim.SGD(self.model.backbone.parameters(), momentum=0.9, lr=self.cfg.optimizer.lr,
                                      weight_decay=self.cfg.optimizer.weight_decay)
            elif self.cfg.optimizer.algo == 'Adam':
                optimizer = optim.AdamW(self.model.backbone.parameters(), lr=self.cfg.optimizer.lr, weight_decay=self.weight_decay)
            if self.cfg.scheduler.mode == 'cosine':
                scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.cfg.num_epochs if isinstance(self.cfg.num_epochs, int) else self.cfg.num_epochs[0],
                                                                 eta_min=0)
            else:
                scheduler = None

            self.adam_init_train(train_loader, test_loader, optimizer, scheduler)
            if self.cfg.pet_cls.startswith('ADAM'):
                self.construct_dual_branch_network()
            elif self.cfg.pet_cls.startswith('RP'):
                self.setup_RP()
            else:
                raise NotImplementedError
        else:
            pass

        if self.cfg.pet_cls.startswith('ADAM'):
            self.replace_fc(train_loader_for_protonet, self.model.backbone, None)
        elif self.cfg.pet_cls.startswith('RP'):
            self.replace_fc_ranpac(train_loader_for_protonet, self.model.backbone)
        else:
            raise NotImplementedError
        self.adam_compute_accuracy(test_loader)

    def construct_dual_branch_network(self):
        """
        Construct dual branch network for ADAM.
        :return:
        """
        network = MultiBranchCosineIncrementalNet(**self.cfg.model.backbone_kwargs)
        network.construct_dual_branch_network(self.model.backbone)
        self.model.backbone = network.to(self.device)

    def adam_init_train(self, train_loader, test_loader, optimizer, scheduler):
        """
        Initialize training for ADAM.
        :param train_loader: Training data loader.
        :param test_loader: Testing data loader.
        :param optimizer: Optimizer.
        :param scheduler: Learning rate scheduler.
        :return:
        """
        prog_bar = tqdm(range(self.cfg.num_epochs if isinstance(self.cfg.num_epochs, int) else self.cfg.num_epochs[0]))
        info = ''
        for _, epoch in enumerate(prog_bar):
            self.model.train()
            losses = 0.0
            correct, total = 0, 0
            for i, (inputs, targets, _) in enumerate(train_loader):
                inputs, targets = inputs.to(self.device), targets.to(self.device)
                logits = self(inputs)

                loss = torch.nn.functional.cross_entropy(logits, targets)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                losses += loss.item()

                _, preds = torch.max(logits, dim=1)
                correct += preds.eq(targets.expand_as(preds)).cpu().sum()
                total += len(targets)

                if self.cfg.debug and i > 0:
                    break

            scheduler.step()
            train_acc = np.around(correct.cpu().data.numpy() * 100 / total, decimals=2)

            if epoch % 5 == 0:
                info = "Task {}, Epoch {}/{} => Loss {:.3f}, Train_accy {:.2f}".format(
                    self.current_task,
                    epoch + 1,
                    self.cfg.num_epochs if isinstance(self.cfg.num_epochs, int) else self.cfg.num_epochs[0],
                    losses / len(train_loader),
                    train_acc,
                )
            else:
                test_acc = self.adam_compute_accuracy(test_loader)
                info = "Task {}, Epoch {}/{} => Loss {:.3f}, Train_accy {:.2f}, Test_accy {:.2f}".format(
                    self.current_task,
                    epoch + 1,
                    self.cfg.num_epochs if isinstance(self.cfg.num_epochs, int) else self.cfg.num_epochs[0],
                    losses / len(train_loader),
                    train_acc,
                    test_acc['Global Acc']
                )
            prog_bar.set_description(info)

        logging.info(info)

    def adam_compute_accuracy(self, loader):
        """
        Compute accuracy for ADAM.
        :param loader: DataLoader.
        :return: Accuracy.
        """
        self.model.eval()
        correct, total = 0, 0
        corect_tasks = [0] * (self.current_task + 1)
        for i, (inputs, targets, task_id) in enumerate(loader):
            inputs = inputs.to(self.device)
            with torch.no_grad():
                outputs = self(inputs)
                outputs = outputs['logits'] if isinstance(outputs, dict) else outputs
            predicts = torch.max(outputs, dim=1)[1]
            correct += (predicts.cpu() == targets).sum()
            for j in range(self.current_task + 1):
                corect_tasks[j] += (predicts.cpu()[task_id == j] == targets[task_id == j]).sum()
            total += len(targets)

            if self.cfg.debug and i > 0:
                break

        out = {
            'Global Acc': correct.cpu().data.numpy() * 100 / total,
        }
        for j in range(self.current_task + 1):
            out[f'Global Acc T{j+1}'] = corect_tasks[j].cpu().data.numpy() * 100 / (total/(self.current_task+1))
        if not self.cfg.debug:
            wandb.log(out)
        else:
            print(out)

        return out

    def replace_fc(self, trainloader, model, args):
        """
        Replace the fully connected layer for ADAM.
        :param trainloader: Training data loader.
        :param model: Model.
        :param args: Arguments.
        :return: Model.
        """

        model = model.eval()
        embedding_list = []
        label_list = []
        # data_list=[]
        with torch.no_grad():
            for i, batch in enumerate(trainloader):
                (data, label, _) = batch
                data = data.to(self.device)
                label = label.to(self.device)
                embedding = model(data)
                embedding = embedding['features'] if isinstance(embedding, dict) else embedding
                embedding_list.append(embedding.cpu())
                label_list.append(label.cpu())
                if self.cfg.debug and i > 10:
                    break

        embedding_list = torch.cat(embedding_list, dim=0)
        label_list = torch.cat(label_list, dim=0)

        class_list = np.unique(label_list)
        proto_list = []
        for class_index in class_list:
            # print('Replacing...',class_index)
            data_index = (label_list == class_index).nonzero().squeeze(-1)
            embedding = embedding_list[data_index]
            proto = embedding.mean(0)
            self.model.backbone.fc.weight.data[class_index] = proto
        return model
    def replace_fc_ranpac(self, trainloader, model):
        """
        Replace the fully connected layer for RP.
        :param trainloader: Training data loader.
        :param model: Model.
        :return: Model.
        """
        model = model.eval()

        if self.cfg.pet_kwargs.use_rp:
            # these lines are needed because the CosineLinear head gets deleted between streams and replaced by one with more classes (for CIL)
            model.fc.use_RP = True
            if self.cfg.pet_kwargs.M > 0:
                model.fc.W_rand = self.W_rand
            else:
                model.fc.W_rand = None

        Features_f = []
        label_list = []
        with torch.no_grad():
            for i, batch in enumerate(trainloader):
                (data, label, _) = batch
                data = data.to(self.device)
                label = label.to(self.device)
                embedding = model.convnet(data)
                Features_f.append(embedding.cpu())
                label_list.append(label.cpu())
        Features_f = torch.cat(Features_f, dim=0)
        label_list = torch.cat(label_list, dim=0)

        Y = target2onehot(label_list, self.datamodule.cfg.num_classes)
        if self.cfg.pet_kwargs.use_rp:
            # print('Number of pre-trained feature dimensions = ',Features_f.shape[-1])
            if self.cfg.pet_kwargs.M > 0:
                Features_h = torch.nn.functional.relu(Features_f @ model.fc.W_rand.cpu())
            else:
                Features_h = Features_f
            self.Q = self.Q + Features_h.T @ Y
            self.G = self.G + Features_h.T @ Features_h
            ridge = optimise_ridge_parameter(Features_h, Y)
            Wo = torch.linalg.solve(self.G + ridge * torch.eye(self.G.size(dim=0)),
                                    self.Q).T  # better numerical stability than .inv
            model.fc.weight.data = Wo[0:model.fc.weight.shape[0], :].to(device=self.device)
        else:
            for class_index in np.unique(self.train_dataset.labels):
                data_index = (label_list == class_index).nonzero().squeeze(-1)
                if self.is_dil:
                    class_prototype = Features_f[data_index].sum(0)
                    model.fc.weight.data[class_index] += class_prototype.to(
                        device=self.device)  # for dil, we update all classes in all tasks
                else:
                    # original cosine similarity approach of Zhou et al (2023)
                    class_prototype = Features_f[data_index].mean(0)
                    model.fc.weight.data[class_index] = class_prototype  # for cil, only new classes get updated
        return model

    def setup_RP(self):
        """
        Setup RanPAC.
        :return:
        """
        self.initiated_G = False
        self.model.backbone.fc.use_RP = True
        if self.cfg.pet_kwargs.M > 0:
            # RP with M > 0
            M = self.cfg.pet_kwargs.M
            self.model.backbone.fc.weight = nn.Parameter(
                torch.Tensor(self.model.backbone.fc.out_features, M).to(device=self.device))  # num classes in task x M
            self.model.backbone.fc.reset_parameters()
            self.model.backbone.fc.W_rand = torch.randn(self.model.backbone.fc.in_features, M).to(device=self.device)
            self.W_rand = copy.deepcopy(
                self.model.backbone.fc.W_rand)  # make a copy that gets passed each time the head is replaced
        else:
            # no RP, only decorrelation
            M = self.model.backbone.fc.in_features  # this M is L in the paper
        self.Q = torch.zeros(M, self.datamodule.cfg.num_classes)
        self.G = torch.zeros(M, M)

    def hide_train_task_adaptive_prediction(self, model, cls_mean, cls_cov, args):
        """
        Train task adaptive prediction for HIDE.
        :param model: Model.
        :param cls_mean: Class mean.
        :param cls_cov: Class covariance.
        :param args: Arguments.
        :return:
        """
        model.train()
        run_epochs = args.crct_epochs
        crct_num = 0
        param_list = [p for n, p in model.named_parameters() if p.requires_grad and 'prompt' not in n]
        network_params = [{'params': param_list, 'lr': args.ca_lr, 'weight_decay': args.weight_decay}]
        optimizer = optim.SGD(network_params, lr=args.ca_lr, momentum=0.9, weight_decay=5e-4)

        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=run_epochs)
        criterion = torch.nn.CrossEntropyLoss().to(self.device)

        for i in range(self.current_task):
            crct_num += self.datamodule.cfg.num_increment_classes

        # TODO: efficiency may be improved by encapsulating sampled data into Datasets class and using distributed sampler.
        for epoch in range(run_epochs):

            sampled_data = []
            sampled_label = []
            num_sampled_pcls = args.batch_size * 5

            if args.ca_storage_efficient_method in ['covariance', 'variance']:
                for c_id in cls_mean.keys():
                    mean = torch.tensor(cls_mean[c_id], dtype=torch.float64).to(self.device)
                    cov = cls_cov[c_id].to(self.device)
                    if args.ca_storage_efficient_method == 'variance':
                        cov = torch.diag(cov)
                    m = MultivariateNormal(mean.float(), cov.float())
                    sampled_data_single = m.sample(sample_shape=(num_sampled_pcls,))
                    sampled_data.append(sampled_data_single)

                    sampled_label.extend([c_id] * num_sampled_pcls)

            elif args.ca_storage_efficient_method == 'multi-centroid':
                for c_id in cls_mean.keys():
                    for cluster in range(len(cls_mean[c_id])):
                        mean = cls_mean[c_id][cluster]
                        var = cls_cov[c_id][cluster]
                        if var.mean() == 0:
                            continue
                        m = MultivariateNormal(mean.float(), (
                                    torch.diag(var) + 1e-4 * torch.eye(mean.shape[0]).to(mean.device)).float())
                        sampled_data_single = m.sample(sample_shape=(num_sampled_pcls,))
                        sampled_data.append(sampled_data_single)
                        sampled_label.extend([c_id] * num_sampled_pcls)
            else:
                raise NotImplementedError

            sampled_data = torch.cat(sampled_data, dim=0).float().to(self.device)
            sampled_label = torch.tensor(sampled_label).long().to(self.device)
            #print(sampled_data.shape)

            inputs = sampled_data
            targets = sampled_label

            sf_indexes = torch.randperm(inputs.size(0))
            inputs = inputs[sf_indexes]
            targets = targets[sf_indexes]
            #print(targets.unique())

            for _iter in range(crct_num):
                inp = inputs[_iter * num_sampled_pcls:(_iter + 1) * num_sampled_pcls]
                tgt = targets[_iter * num_sampled_pcls:(_iter + 1) * num_sampled_pcls]
                logits = model.head(inp, None)

                preds = torch.max(logits, dim=1)[1]
                acc = (preds == tgt).float().sum() / len(tgt)

                loss = criterion(logits, tgt)  # base criterion (CrossEntropyLoss)

                if _iter % 100 == 0:
                    logging.info(f'T{self.current_task}|E{epoch}|B{_iter}: loss {round(loss.item(),4)}, acc {round(acc.item(), 4)}')

                if not math.isfinite(loss.item()):
                    print("Loss is {}, stopping training".format(loss.item()))
                    sys.exit(1)

                optimizer.zero_grad()
                loss.backward()
                # for name, p in model.named_parameters():
                #    if p.requires_grad and p.grad is None:
                #        print(name)
                optimizer.step()
                #if self.device != 'cpu':
                #    torch.cuda.synchronize()

            scheduler.step()

    def configure_optimizer(self, *_, query=False, bic=False):
        """
        Configure optimizer
        :param _:
        :param query: Query training.
        :param bic: Bic training.
        :return: Optimizer.
        """
        if query:
            if self.cfg.model.query.endswith('l2p'):
                return super().configure_optimizer(self.query_function)
            if self.cfg.model.query.endswith('hide'):
                q_cfg = copy.deepcopy(self.cfg.optimizer)
                q_cfg.lr = self.cfg.model.query_train_kwargs.lr
                return super().configure_optimizer(self.query_function.head, cfg=q_cfg)
            else:
                return super().configure_optimizer(self.query_pets)
        if bic:
            bic_opt_cfg = self.cfg.optimizer
            if isinstance(bic_opt_cfg.lr, list):
                bic_opt_cfg.lr = bic_opt_cfg.lr[0]
            return super().configure_optimizer(self.model.head.bias_alphas, self.model.head.bias_betas, cfg=bic_opt_cfg)
        if self.cfg.pet_cls in ["L2P", "DP", "LLS", 'CODA', 'HIDE', 'OA']:
            if self.cfg.model.head == "DynamicNMCHead":
                return super().configure_optimizer(self.model.backbone)
            return super().configure_optimizer(self.model.head, self.model.backbone)
        elif self.cfg.pet_cls.startswith('ADAM') or self.cfg.pet_cls.startswith('RP'):
            return super().configure_optimizer(self.model.head, self.model.backbone)
        elif self.cfg.model.head == "KMeansHead" and self.cfg.pet_cls == "GDumb":
            return super().configure_optimizer(self.model.backbone)
        elif self.cfg.pet_cls == 'None':
            return super().configure_optimizer(self.model.head)
        else:
            if self.cfg.model.head != "DynamicSimpleHead":
                return super().configure_optimizer(self.pets)
            return super().configure_optimizer(self.model.head, self.pets)


@hydra.main(version_base=None, config_path='config/experiments/')
def main(cfg: DictConfig):
    kwargs = dict(
        config=cfg, datamodule_cls=MyDataModule, module_cls=MyModule, config_cls=MyConf
    )
    App(**kwargs).run()


if __name__ == "__main__":
    main()