core/model/MInfLoRA3.py

"""
Code Reference:
https://github.com/liangyanshuo/InfLoRA/blob/main/methods/inflora.py
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
import numpy as np

from tqdm import tqdm
from .backbone.transformer import MultiHeadAttention_MultiMaskedLoRA3

Epsilon = 0.5

class TopK:

    '''
    A class to maintain a collection of the top K items based on a specified attribute.

    This class allows for the dynamic addition of items, each represented as a dictionary, 
    where each dictionary must have a key 'proj_norm' that represents the value used 
    to determine the ranking. The class keeps track of the top K items with the highest 
    'proj_norm' values.
    '''

    def __init__(self, k):
        self.k = k
        self.top_k_list = []

    def add(self, dict):
        if len(self.top_k_list) < self.k:
            self.top_k_list.append(dict)
        elif dict['proj_norm'] > min(self.top_k_list, key=lambda x: x['proj_norm'])['proj_norm']:
            self.top_k_list.remove(min(self.top_k_list, key=lambda x: x['proj_norm']))
            self.top_k_list.append(dict)
        elif dict['proj_norm'] == min(self.top_k_list, key=lambda x: x['proj_norm'])['proj_norm'] and \
            dict['proj_norm'] == max(self.top_k_list, key=lambda x: x['proj_norm'])['proj_norm']:
            self.top_k_list.remove(min(self.top_k_list, key=lambda x: x['task_id']))
            self.top_k_list.append(dict)

    def get_top_k(self):
        return self.top_k_list

class SiNet(nn.Module):
    def __init__(self, backbone, **kwargs):
        super().__init__()

        self._cur_task_id = -1
        self.backbone = backbone
        self.init_cls_num = kwargs["init_cls_num"]
        self.inc_cls_num = kwargs["inc_cls_num"]

        self.classifier_pool = nn.ModuleList([
            nn.Linear(kwargs["embd_dim"], kwargs['init_cls_num'], bias=True)] + 
            [nn.Linear(kwargs["embd_dim"], kwargs['inc_cls_num'], bias=True) for _ in range(kwargs['task_num'] - 1)])

        for name, module in self.backbone.named_modules():
            if 'transformer' in name and 'blocks' not in name:
                self.transformer_module = module

    def update_fc(self):
        self._cur_task_id += 1

    def forward(self, x, expert_id, inference = False):
        logits = []
        features = self.backbone(x, expert_id = expert_id)

        if inference:

            # Bayesian
            for i, prompts in enumerate(self.classifier_pool[:self._cur_task_id + 1]):
                # No Masking
                logits.append(prompts(features))

            logits = torch.cat(logits, dim=1)

            return logits

        else:
            logits.append(self.classifier_pool[self._cur_task_id](features))
            return torch.cat(logits, dim=1)

    def update_input_matrix(self, x):
        self.backbone(x, expert_id = -1, get_input_matrix = True)

class MInfLoRA3(nn.Module):

    def __init__(self, backbone, device, **kwargs):
        super().__init__()

        self.device = device
        self.init_cls_num = kwargs["init_cls_num"]
        self.inc_cls_num = kwargs["inc_cls_num"]
        self.task_num = kwargs["task_num"]
        self.lame = kwargs["lame"]
        self.lamb = kwargs["lamb"]
        self.embd_dim = kwargs["embd_dim"]
        self.eval_mat = kwargs['eval_mat']

        self._known_classes = 0
        self.feature_list = []
        self.project_type = []

        self.distributed = torch.distributed.is_initialized()
        self.local_rank = torch.distributed.get_rank() if self.distributed else 0
        self._network = SiNet(backbone, **kwargs)

        self.attention_modules = [module for module in self._network.modules() if isinstance(module, MultiHeadAttention_MultiMaskedLoRA3)]

        # TRGP Implementation
        self.feature_list_each_tasks = [[np.zeros((1)) for _ in range(len(self.attention_modules))] for _ in range(self.task_num)]
        self.final_decision = [[np.zeros((1)) for _ in range(len(self.attention_modules))] for _ in range(self.task_num)]
        self.before_mat = [[0 for _ in range(len(self.attention_modules))] for _ in range(self.task_num)]

        self.experts_distributions = []

        # Class Alignment Implementation
        self._use_class_alignment = kwargs['use_ca']
        self._class_means = None
        self._class_covs = None
        self._dataset = kwargs['dataset']
        if self._dataset == 'cifar':
            self.logit_norm = None
        else:
            self.logit_norm = 0.1   
    
        self.lll = []

        self._network.to(self.device)
        
    def observe(self, data):

        x, y = data['image'].to(self.device, non_blocking=True), data['label'].to(self.device, non_blocking=True) - self._known_classes

        logits = self._network(x, expert_id = self._network._cur_task_id)
        loss = F.cross_entropy(logits, y)

        preds = logits.max(1)[1]
        acc = preds.eq(y).sum().item() / y.shape[0]

        return preds, acc, loss
    
    def inference(self, data, **kwargs):

        task_id = kwargs['task_id'] if 'task_id' in kwargs else -1
        x, y = data['image'].to(self.device, non_blocking=True), data['label'].to(self.device, non_blocking=True)

        logits = self._network(x, expert_id = task_id, inference = True)
        preds = logits.max(1)[1]
        acc = preds.eq(y).sum().item() / y.shape[0]

        return preds, acc

    @torch.no_grad()
    def before_task(self, task_idx, buffer, train_loader, test_loaders):
        
        self._network.update_fc()

        [module.init_param() for module in self.attention_modules]
        
        self._update_input_matrix(train_loader, test_loaders[0].dataset.trfms)
        
        '''
        for i, module in enumerate(self.attention_modules):
            
            topk = TopK(1)

            mat = module.cur_matrix.cpu().numpy()
            mat_norm = np.linalg.norm(mat)

            for task_id in range(task_idx):
            
                proj_norm = np.linalg.norm(self.feature_list_each_tasks[task_id][i] @ self.feature_list_each_tasks[task_id][i].T @ mat)
                
                if proj_norm > Epsilon * mat_norm:
                    topk.add({'proj_norm':proj_norm, 'task_id': task_id})

            self.final_decision[task_idx][i] = [dic['task_id'] for dic in topk.get_top_k()]
            print(f'Layer {i} of {task_idx} consider {self.final_decision[task_idx][i]} as trust region')
        '''

        if self.local_rank == 0:

            if task_idx == 0:
                for i, module in enumerate(self.attention_modules):
                    
                    U, _, _ = torch.linalg.svd(module.cur_matrix)
                    U = torch.Tensor(U).to(self.device)

                    module.lora_A_k.weight.data.copy_(U[:,:module.lora_rank].T/math.sqrt(3))
                    module.lora_A_v.weight.data.copy_(U[:,:module.lora_rank].T/math.sqrt(3))
            else:
                for i, module in enumerate(self.attention_modules):
                    assert self.project_type[i] == 'remove' or self.project_type[i] == 'retain'

                    #tr = self.final_decision[task_idx][i][0]
                    #feature_mat = torch.Tensor(self.feature_list_each_tasks[tr][i] @ self.feature_list_each_tasks[tr][i].T).to(self.device)

                    #target_shape = max(70, self.feature_list[i].shape[1]) # constant 50 and whole feature_list and no QQ^T, get best result for now
                    target_shape = 768

                    # either /math.sqrt(3) or no /math.sqrt(3) is bad

                    cur_matrix = module.cur_matrix.to(self.device)
                    feature_mat = torch.Tensor(self.feature_list[i] @ self.feature_list[i].T).to(self.device)

                    q_weight, k_weight, v_weight = module.qkv.weight.chunk(3, dim=0)
                    kk = feature_mat - k_weight.data @ feature_mat
                    vv = feature_mat - v_weight.data @ feature_mat

                    U, _, _ = np.linalg.svd(kk.cpu().numpy(), full_matrices = False)
                    U = torch.Tensor(U).to(self.device)
                    module.space_k[task_idx] = U[:, :target_shape].T/math.sqrt(3)

                    U, _, _ = np.linalg.svd(vv.cpu().numpy(), full_matrices = False)
                    U = torch.Tensor(U).to(self.device)
                    module.space_v[task_idx] = U[:, :target_shape].T/math.sqrt(3)

                    if self.project_type[i] == 'remove':
                        cur_matrix = cur_matrix - feature_mat @ cur_matrix
                    else:
                        cur_matrix = feature_mat @ cur_matrix

                    U, _, _ = np.linalg.svd(cur_matrix.cpu().numpy(), full_matrices = False)
                    U = torch.Tensor(U).to(self.device)

                    module.lora_A_k.weight.data.copy_(U[:, :module.lora_rank].T/math.sqrt(3))
                    module.lora_A_v.weight.data.copy_(U[:, :module.lora_rank].T/math.sqrt(3))

        # Initilize space_k and space_v before sync
        if self.local_rank != 0 and task_idx != 0:
            for module in self.attention_modules:
                module.space_k[task_idx] = torch.empty((50, self.embd_dim)).to(self.device)
                module.space_v[task_idx] = torch.empty((50, self.embd_dim)).to(self.device)

        if self.distributed and task_idx != 0:
            dist.barrier()
            for module in self.attention_modules:
                dist.broadcast(module.lora_A_k.weight.data, 0)
                dist.broadcast(module.lora_A_v.weight.data, 0)
                dist.broadcast(module.space_k[task_idx].contiguous(), 0)
                dist.broadcast(module.space_v[task_idx].contiguous(), 0)

        for name, param in self._network.named_parameters():
            param.requires_grad_(False)
            if f"classifier_pool.{task_idx}" in name or \
               f"lora_B_k_list.{task_idx}" in name or \
               f"lora_B_v_list.{task_idx}" in name or \
               f"scale_param.{task_idx}" in name:
                param.requires_grad_(True)
        
        if self.local_rank == 0:
            for name, param in self._network.named_parameters():
                if param.requires_grad:
                    print(name)

    def after_task(self, task_idx, buffer, train_loader, test_loaders):
        '''
        Called after each task before final testing, it is used to perform preliminary operations on the mapping matrix to facilitate the update of lora_a layer in the next round of before_task
        '''

        self._known_classes += self.init_cls_num if task_idx == 0 else self.inc_cls_num

        [module.merge_weight() for module in self.attention_modules]

        self._update_feature(task_idx, train_loader, test_loaders)

    @torch.no_grad()
    def _update_feature(self, task_idx, train_loader, test_loaders):
        '''
        Update feature lists and the corresponding type
        '''

        self._update_input_matrix(train_loader, test_loaders[0].dataset.trfms)

        if self.local_rank == 0:

            threshold = (self.lame - self.lamb)*task_idx/self.task_num + self.lamb

            if task_idx == 0:
                for i, module in enumerate(self.attention_modules):

                    activation = module.cur_matrix

                    U, S, _ = np.linalg.svd(activation, full_matrices=False)
                    sval_ratio = (S**2)/(S**2).sum()
                    r = max(np.sum(np.cumsum(sval_ratio) < threshold), 1)
                    assert r < activation.shape[0]/2

                    self.feature_list_each_tasks[task_idx][i] = U[:, :r]
                    self.feature_list.append(U[:, :r])
                    self.project_type.append('remove')
           
            else:
                for i, module in enumerate(self.attention_modules):

                    activation = module.cur_matrix
                    _, S, _ = np.linalg.svd(activation, full_matrices=False)
                    sval_total = (S**2).sum()

                    if self.project_type[i] == 'remove':

                        act_hat = activation - torch.Tensor(self.feature_list[i] @ self.feature_list[i].T) @ activation
                        U, S, _ = np.linalg.svd(act_hat, full_matrices = False)
                        sigma = S**2

                        delta = (torch.Tensor(self.feature_list[i]).T @ activation @ activation.T @ torch.Tensor(self.feature_list[i])).diagonal()

                        stack = np.hstack((delta, sigma))
                        stack_index = np.argsort(stack)[::-1] # the index of each element in descending sorted array
                        stack = np.sort(stack)[::-1] # descending sorted array

                        if threshold * sval_total <= 0:
                            r = 0
                        else:
                            r = min(np.sum(np.cumsum(stack) < threshold * sval_total) + 1, activation.shape[0])

                        Ui = np.hstack((self.feature_list[i], U))
                        sel_each = stack_index[:r]
                        sel_overall = sel_each[sel_each >= len(delta)] # without overlap

                        self.feature_list[i] = np.hstack((self.feature_list[i], Ui[:, sel_overall]))
                        self.feature_list_each_tasks[task_idx][i] = Ui[:, sel_each]

                        if sel_overall.shape[0] == 0:
                            print(f'Skip Updating Space for layer: {i+1}')

                    else:
                        act_hat = torch.Tensor(self.feature_list[i] @ self.feature_list[i].T) @ activation
                        U,S,_ = np.linalg.svd(act_hat, full_matrices = False)
                        sval_hat = (S**2).sum()
                        sval_ratio = (S**2)/sval_total     
                        accumulated_sval = sval_hat/sval_total          

                        if accumulated_sval < 1 - threshold:
                            print (f'Skip Updating Space for layer: {i+1}')
                        else:
                            r = np.sum(accumulated_sval - np.cumsum(sval_ratio) >= 1 - threshold) + 1
                            act_feature = self.feature_list[i] - U[:,0:r] @ U[:,0:r].T @ self.feature_list[i]
                            U, _, _ = np.linalg.svd(act_feature)
                            self.feature_list[i]=U[:,:self.feature_list[i].shape[1]-r]

            print('-'*40)
            print(f'Threshold: {threshold}')
            print('-'*40)
            for i in range(len(self.feature_list)):
                '''
                if self.project_type[i]=='remove' and (self.feature_list[i].shape[1] > (self.feature_list[i].shape[0]/2)):
                    feature = self.feature_list[i]
                    U, S, V = np.linalg.svd(feature)
                    new_feature = U[:,feature.shape[1]:]
                    self.feature_list[i] = new_feature
                    self.project_type[i] = 'retain'
                elif self.project_type[i]=='retain':
                    assert self.feature_list[i].shape[1] <= (self.feature_list[i].shape[0]/2)
                '''
                print ('Layer {} : {}/{} type {}'.format(i+1,self.feature_list[i].shape[1], self.feature_list[i].shape[0], self.project_type[i]))
            print('-'*40)

    @torch.no_grad()
    def _update_input_matrix(self, train_loader, test_trfms):

        if self.eval_mat:
            self._network.eval()
            train_trfms = train_loader.dataset.trfms
            train_loader.dataset.trfms = test_trfms

        for module in self.attention_modules:
            module.reset_input_matrix()

        for batch in tqdm(train_loader, desc="Forwarding to get input matrix", disable=(self.local_rank != 0)):
            self._network.update_input_matrix(batch['image'].to(self.device, non_blocking=True))

        if self.distributed: # Combine input matrix across all GPUs
            for module in self.attention_modules:
                n_cur_matrix = torch.tensor(module.n_cur_matrix).to(self.device)
                cur_matrix = (module.cur_matrix * module.n_cur_matrix).to(self.device)

                dist.all_reduce(cur_matrix, op=dist.ReduceOp.SUM)
                dist.all_reduce(n_cur_matrix, op=dist.ReduceOp.SUM)

                module.n_cur_matrix = n_cur_matrix.item()
                module.cur_matrix = cur_matrix.cpu() / module.n_cur_matrix

        if self.eval_mat:
            self._network.train()
            train_loader.dataset.trfms = train_trfms

    def get_parameters(self, config):
        return self._network.parameters()