methods/elasticdnn/api/algs/md_pretraining_bk.py

from typing import Any, Dict
from schema import Schema, Or
import schema
from data import Scenario, MergedDataset
from methods.base.alg import BaseAlg
from data import build_dataloader
from ..model import ElasticDNN_OfflineFMModel, ElasticDNN_OfflineMDModel
from ...model.base import ElasticDNNUtil
import torch.optim
import tqdm
import torch.nn.functional as F
from torch import nn
from utils.dl.common.env import create_tbwriter
import os
import random
import numpy as np
from copy import deepcopy
from utils.dl.common.model import LayerActivation, get_module
from utils.common.log import logger


class ElasticDNN_MDPretrainingAlg(BaseAlg):
    """
    TODO: fine-tuned FM -> init MD -> trained MD -> construct indexes (only between similar weights) and fine-tune
    """
    def get_required_models_schema(self) -> Schema:
        return Schema({
            'fm': ElasticDNN_OfflineFMModel,
            'md': ElasticDNN_OfflineMDModel
        })
        
    def get_required_hyp_schema(self) -> Schema:
        return Schema({
            'launch_tbboard': bool,
            
            'samples_size': (int, int, int, int),
            'generate_md_width_ratio': int,
            
            'FBS_r': int,
            'FBS_ignore_layers': [str],
            
            'train_batch_size': int,
            'val_batch_size': int,
            'num_workers': int,            
            'optimizer': str,
            'md_optimizer_args': dict,
            'indexes_optimizer_args': dict,
            'scheduler': str,
            'scheduler_args': dict,
            'num_iters': int,
            'val_freq': int,
            'max_sparsity': float,
            'min_sparsity': float,
            'distill_loss_weight': float,
            'index_loss_weight': float,
            'val_num_sparsities': int,
            
            'bn_cal_num_iters': int,
            
            'index_guided_linear_comb_split_size': Or(int, None)
        })
        
    def upsample_2d_tensor(self, p: torch.Tensor, target_len: int):
        assert p.dim() == 2 # regard 2d weight as (batch_size, 1d_vector_dim)
        return F.upsample(p.unsqueeze(1).unsqueeze(3),
                          size=(target_len, 1),
                          mode='bilinear').squeeze(3).squeeze(1)
        
    def two_params_diff_fast(self, trained_p: torch.Tensor, ref_p: torch.Tensor, 
                             index: torch.Tensor, 
                             split_size: int):

        assert trained_p.dim() == ref_p.dim()
        assert index.size(0) == trained_p.size(0) and index.size(1) == ref_p.size(0)
        
        # print(trained_p.size(), ref_p.size(), index.size())

        ref_p = ref_p.detach()
        if trained_p.dim() > 1:
            trained_p = trained_p.flatten(1)
            ref_p = ref_p.flatten(1)
            
            # the weight size of master DNN and foundation model may be totally different
            
            # MD -> FM: upsample first
            # FM -> MD: downsample first
            if trained_p.size(1) < ref_p.size(1):
                trained_p = self.upsample_2d_tensor(trained_p, ref_p.size(1))
            
            index = index.unsqueeze(-1)
        #     linear_combed_ref_p = (ref_p.unsqueeze(0) * index).sum(1)
        # else:
        
        # print(trained_p.size(), ref_p.size(), index.size())
        
        if split_size is None:
            # old version: huge memory consumption, not recommended (although this is fastest)
            # print('old version')
            linear_combed_ref_p = (ref_p.unsqueeze(0) * index).sum(1)
        
        else:
            # new version
            linear_combed_ref_p = 0
            
            cur_split_size = split_size
            while index.size(1) % cur_split_size != 0:
                cur_split_size -= 1
            # print(cur_split_size) 
            
            for i in range(0, index.size(1), cur_split_size):
                # if not isinstance(linear_combed_ref_p, int):
                    # print(linear_combed_ref_p.size(), ref_p.unsqueeze(0)[:, i: i + cur_split_size].size(), index[:, i: i + cur_split_size].size())
                linear_combed_ref_p += ref_p.unsqueeze(0)[:, i: i + cur_split_size] * index[:, i: i + cur_split_size]
            linear_combed_ref_p = linear_combed_ref_p.sum(1)
            
        diff = (linear_combed_ref_p - trained_p).norm(2) ** 2
        return diff
        
    def get_index_loss(self, fm, md, indexes, match_fn, split_size):
        res = 0.

        for name, p in md.named_parameters():
            if p.dim() == 0:
                continue
            
            raw_p = match_fn(name, fm)
            if raw_p is None:
                continue

            index = indexes[name]
            
            # print(name)
            res += self.two_params_diff_fast(p, raw_p, index, split_size)
        return res
    
    def bn_cal(self, model: nn.Module, train_loader, num_iters, device):
        has_bn = False
        for n, m in model.named_modules():
            if isinstance(m, nn.BatchNorm2d):
                has_bn = True
                break
        
        if not has_bn:
            return {}
        
        def bn_calibration_init(m):
            """ calculating post-statistics of batch normalization """
            if getattr(m, 'track_running_stats', False):
                # reset all values for post-statistics
                m.reset_running_stats()
                # set bn in training mode to update post-statistics
                m.training = True
                
        with torch.no_grad():
            model.eval()
            model.apply(bn_calibration_init)
            for _ in range(num_iters):
                x, _ = next(train_loader)
                model(x.to(device))
            model.eval()
            
        bn_stats = {}
        for n, m in model.named_modules():
            if isinstance(m, nn.BatchNorm2d):
                bn_stats[n] = m
        return bn_stats
        
    def run(self, scenario: Scenario, hyps: Dict) -> Dict[str, Any]:
        super().run(scenario, hyps)
        
        # sanity check
        # a=  torch.tensor([[1, 2, 3], [1, 2, 4]])
        # index = torch.tensor([[1, 2, 3],
        # [1, 2, 4]])
        # b = torch.tensor([[1, 2, 3], [1, 2, 4], [2, 3, 4]])
        # print(self.two_params_diff_fast(a, b, index, hyps['index_guided_linear_comb_split_size']))
        
        assert isinstance(self.models['md'], ElasticDNN_OfflineMDModel) # for auto completion
        assert isinstance(self.models['fm'], ElasticDNN_OfflineFMModel) # for auto completion
        
        # 1. add FBS
        device = self.models['md'].device
        
        if self.models['md'].models_dict['main'] == -1:
            logger.info(f'init master DNN by reducing width of an adapted foundation model (already tuned by LoRA)...')
            
            before_fm_model = deepcopy(self.models['fm'].models_dict['main'])
            lora_util = self.models['fm'].get_lora_util()
            lora_absorbed_fm_model = lora_util.absorb_lora_and_recover_net_structure(self.models['fm'].models_dict['main'], 
                                                                                    torch.rand(hyps['samples_size']).to(device))
            self.models['fm'].models_dict['main'] = lora_absorbed_fm_model
            master_dnn = self.models['fm'].generate_md_by_reducing_width(hyps['generate_md_width_ratio'], 
                                                                        torch.rand(hyps['samples_size']).to(device))
            self.models['fm'].models_dict['main'] = before_fm_model
            
            elastic_dnn_util = self.models['fm'].get_elastic_dnn_util()
            master_dnn = elastic_dnn_util.convert_raw_dnn_to_master_dnn_with_perf_test(master_dnn,
                                                                                    hyps['FBS_r'], hyps['FBS_ignore_layers'])
            self.models['md'].models_dict['main'] = master_dnn
            self.models['md'].to(device)
        
        # 2. train (knowledge distillation, index relationship)
        offline_datasets = scenario.get_offline_datasets()
        train_dataset = MergedDataset([d['train'] for d in offline_datasets.values()])
        val_dataset = MergedDataset([d['val'] for d in offline_datasets.values()])
        train_loader = iter(build_dataloader(train_dataset, hyps['train_batch_size'], hyps['num_workers'],
                                        True, None))
        val_loader = build_dataloader(val_dataset, hyps['val_batch_size'], hyps['num_workers'],
                                      False, False)
        
        # 2.1 train whole master DNN (knowledge distillation)
        for p in master_dnn.parameters():
            p.requires_grad = True
        self.models['md'].to_train_mode()
        
        optimizer = torch.optim.__dict__[hyps['optimizer']]([
            {'params': self.models['md'].models_dict['main'].parameters(), **hyps['md_optimizer_args']}
        ])
        scheduler = torch.optim.lr_scheduler.__dict__[hyps['scheduler']](optimizer, **hyps['scheduler_args'])
        tb_writer = create_tbwriter(os.path.join(self.res_save_dir, 'tb_log'), launch_tbboard=hyps['launch_tbboard'])
        pbar = tqdm.tqdm(range(hyps['num_iters']), dynamic_ncols=True)
        best_avg_val_acc = 0.
        
        md_output_hook = None
        
        for iter_index in pbar:
            self.models['md'].to_train_mode()
            self.models['fm'].to_eval_mode()
            
            rand_sparsity = random.random() * (hyps['max_sparsity'] - hyps['min_sparsity']) + hyps['min_sparsity']
            elastic_dnn_util.set_master_dnn_sparsity(self.models['md'].models_dict['main'], rand_sparsity)
            
            x, y = next(train_loader)
            x, y = x.to(device), y.to(device)
            
            with torch.no_grad():
                fm_output = self.models['fm'].infer(x)

            if md_output_hook is None:
                md_output_hook = LayerActivation(self.models['md'].models_dict['main'], False, device)
            task_loss = self.models['md'].forward_to_get_task_loss(x, y)
            md_output = md_output_hook.output
            distill_loss = hyps['distill_loss_weight'] * self.models['md'].get_distill_loss(md_output, fm_output)
            total_loss = task_loss + distill_loss
            
            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()
            scheduler.step()
            
            if (iter_index + 1) % hyps['val_freq'] == 0:
                
                elastic_dnn_util.clear_cached_channel_attention_in_master_dnn(self.models['md'].models_dict['main'])
                md_output_hook.remove()
                md_output_hook = None
                
                cur_md = self.models['md'].models_dict['main']
                md_for_test = deepcopy(self.models['md'].models_dict['main'])
                val_accs = {}
                avg_val_acc = 0.
                bn_stats = {}
                
                for val_sparsity in np.linspace(hyps['min_sparsity'], hyps['max_sparsity'], num=hyps['val_num_sparsities']):
                    elastic_dnn_util.set_master_dnn_sparsity(md_for_test, val_sparsity)
                    bn_stats[f'{val_sparsity:.4f}'] = self.bn_cal(md_for_test, train_loader, hyps['bn_cal_num_iters'], device)
                    self.models['md'].models_dict['main'] = md_for_test
                    self.models['md'].to_eval_mode()
                    val_acc = self.models['md'].get_accuracy(val_loader)
                    
                    val_accs[f'{val_sparsity:.4f}'] = val_acc
                    avg_val_acc += val_acc
                    
                avg_val_acc /= hyps['val_num_sparsities']
                
                self.models['md'].models_dict['main'] = cur_md
                self.models['md'].models_dict['bn_stats'] = bn_stats
                
                self.models['md'].save_model(os.path.join(self.res_save_dir, 'models/md_last.pt'))
                self.models['fm'].save_model(os.path.join(self.res_save_dir, 'models/fm_last.pt'))
                
                if avg_val_acc > best_avg_val_acc:
                    best_avg_val_acc = avg_val_acc
                    self.models['md'].save_model(os.path.join(self.res_save_dir, 'models/md_best.pt'))
                    self.models['fm'].save_model(os.path.join(self.res_save_dir, 'models/fm_best.pt'))
                
            tb_writer.add_scalars(f'losses', dict(task=task_loss, distill=distill_loss, total=total_loss), iter_index)
            pbar.set_description(f'loss: {total_loss:.6f}')
            if (iter_index + 1) >= hyps['val_freq']:
                tb_writer.add_scalars(f'accs/val_accs', val_accs, iter_index)
                tb_writer.add_scalar(f'accs/avg_val_acc', avg_val_acc, iter_index)
                pbar.set_description(f'loss: {total_loss:.6f}, avg_val_acc: {avg_val_acc:.4f}')