from typing import List import torch from methods.base.model import BaseModel import tqdm from torch import nn import torch.nn.functional as F from abc import abstractmethod from methods.elasticdnn.model.base import ElasticDNNUtil from methods.elasticdnn.pipeline.offline.fm_to_md.base import FM_to_MD_Util from methods.elasticdnn.pipeline.offline.fm_lora.base import FMLoRA_Util from utils.dl.common.model import LayerActivation class ElasticDNN_OfflineFMModel(BaseModel): def get_required_model_components(self) -> List[str]: return ['main'] @abstractmethod def generate_md_by_reducing_width(self, reducing_width_ratio, samples: torch.Tensor): pass @abstractmethod def forward_to_get_task_loss(self, x, y, *args, **kwargs): raise NotImplementedError @abstractmethod def get_feature_hook(self) -> LayerActivation: pass @abstractmethod def get_elastic_dnn_util(self) -> ElasticDNNUtil: pass @abstractmethod def get_lora_util(self) -> FMLoRA_Util: pass @abstractmethod def get_task_head_params(self): pass class ElasticDNN_OfflineClsFMModel(ElasticDNN_OfflineFMModel): def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: x, y = x.to(self.device), y.to(self.device) output = self.infer(x) pred = F.softmax(output.logits, dim=1).argmax(dim=1) #correct = torch.eq(torch.argmax(output.logits,dim = 1), y).sum().item() correct = torch.eq(pred, y).sum().item() acc += correct sample_num += len(y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_correct: {correct}, ' f'cur_batch_acc: {(correct / len(y)):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'](x) import numpy as np class StreamSegMetrics: """ Stream Metrics for Semantic Segmentation Task """ def __init__(self, n_classes): self.n_classes = n_classes self.confusion_matrix = np.zeros((n_classes, n_classes)) def update(self, label_trues, label_preds): for lt, lp in zip(label_trues, label_preds): self.confusion_matrix += self._fast_hist( lt.flatten(), lp.flatten() ) @staticmethod def to_str(results): string = "\n" for k, v in results.items(): if k!="Class IoU": string += "%s: %f\n"%(k, v) return string def _fast_hist(self, label_true, label_pred): mask = (label_true >= 0) & (label_true < self.n_classes) hist = np.bincount( self.n_classes * label_true[mask].astype(int) + label_pred[mask], minlength=self.n_classes ** 2, ).reshape(self.n_classes, self.n_classes) return hist def get_results(self): """Returns accuracy score evaluation result. - overall accuracy - mean accuracy - mean IU - fwavacc """ hist = self.confusion_matrix acc = np.diag(hist).sum() / hist.sum() acc_cls = np.diag(hist) / hist.sum(axis=1) acc_cls = np.nanmean(acc_cls) iu = np.diag(hist) / (hist.sum(axis=1) + hist.sum(axis=0) - np.diag(hist)) mean_iu = np.nanmean(iu) freq = hist.sum(axis=1) / hist.sum() fwavacc = (freq[freq > 0] * iu[freq > 0]).sum() cls_iu = dict(zip(range(self.n_classes), iu)) return { "Overall Acc": acc, "Mean Acc": acc_cls, "FreqW Acc": fwavacc, "Mean IoU": mean_iu, "Class IoU": cls_iu, } def reset(self): self.confusion_matrix = np.zeros((self.n_classes, self.n_classes)) class ElasticDNN_OfflineSegFMModel(ElasticDNN_OfflineFMModel): def __init__(self, name: str, models_dict_path: str, device: str, num_classes): super().__init__(name, models_dict_path, device) self.num_classes = num_classes def get_accuracy(self, test_loader, *args, **kwargs): device = self.device self.to_eval_mode() metrics = StreamSegMetrics(self.num_classes) metrics.reset() import tqdm pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), leave=False, dynamic_ncols=True) with torch.no_grad(): for batch_index, (x, y) in pbar: x, y = x.to(device, dtype=x.dtype, non_blocking=True, copy=False), \ y.to(device, dtype=y.dtype, non_blocking=True, copy=False) output = self.infer(x) pred = output.detach().max(dim=1)[1].cpu().numpy() metrics.update((y + 0).cpu().numpy(), pred) res = metrics.get_results() pbar.set_description(f'cur batch mIoU: {res["Mean IoU"]:.4f}') res = metrics.get_results() return res['Mean IoU'] def infer(self, x, *args, **kwargs): return self.models_dict['main'](x) class ElasticDNN_OfflineDetFMModel(ElasticDNN_OfflineFMModel): def __init__(self, name: str, models_dict_path: str, device: str, num_classes): super().__init__(name, models_dict_path, device) self.num_classes = num_classes def get_accuracy(self, test_loader, *args, **kwargs): # print('DeeplabV3: start test acc') _d = test_loader.dataset from data import build_dataloader if _d.__class__.__name__ == 'MergedDataset': # print('\neval on merged datasets') datasets = _d.datasets test_loaders = [build_dataloader(d, test_loader.batch_size, test_loader.num_workers, False, None) for d in datasets] accs = [self.get_accuracy(loader) for loader in test_loaders] # print(accs) return sum(accs) / len(accs) # print('dataset len', len(test_loader.dataset)) model = self.models_dict['main'] device = self.device model.eval() # print('# classes', model.num_classes) model = model.to(device) from dnns.yolov3.coco_evaluator import COCOEvaluator from utils.common.others import HiddenPrints with torch.no_grad(): with HiddenPrints(): evaluator = COCOEvaluator( dataloader=test_loader, img_size=(224, 224), confthre=0.01, nmsthre=0.65, num_classes=self.num_classes, testdev=False ) res = evaluator.evaluate(model, False, False) map50 = res[1] # print('eval info', res[-1]) return map50 def infer(self, x, *args, **kwargs): if len(args) > 0: print(args, len(args)) return self.models_dict['main'](x, *args) # forward(x, label) return self.models_dict['main'](x) class ElasticDNN_OfflineSenClsFMModel(ElasticDNN_OfflineFMModel): def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) y = y.to(self.device) output = self.infer(x) pred = F.softmax(output, dim=1).argmax(dim=1) correct = torch.eq(pred, y).sum().item() acc += correct sample_num += len(y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_correct: {correct}, ' f'cur_batch_acc: {(correct / len(y)):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'].forward(**x) from accelerate.utils.operations import pad_across_processes class ElasticDNN_OfflineTrFMModel(ElasticDNN_OfflineFMModel): def get_accuracy(self, test_loader, *args, **kwargs): # TODO: BLEU from sacrebleu import corpus_bleu acc = 0 num_batches = 0 self.to_eval_mode() from data.datasets.sentiment_classification.global_bert_tokenizer import get_tokenizer tokenizer = get_tokenizer() def _decode(o): # https://github.com/huggingface/transformers/blob/main/examples/research_projects/seq2seq-distillation/finetune.py#L133 o = tokenizer.batch_decode(o, skip_special_tokens=True) return [oi.strip() for oi in o] with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) label = y.to(self.device) # generated_tokens = self.infer(x, generate=True) generated_tokens = self.infer(x).logits.argmax(-1) # pad tokens generated_tokens = pad_across_processes( generated_tokens, dim=1, pad_index=tokenizer.pad_token_id ) # pad label label = pad_across_processes( label, dim=1, pad_index=tokenizer.pad_token_id ) label = label.cpu().numpy() label = np.where(label != -100, label, tokenizer.pad_token_id) decoded_output = _decode(generated_tokens) decoded_y = _decode(y) decoded_y = [decoded_y] if batch_index == 0: print(decoded_y, decoded_output) bleu = corpus_bleu(decoded_output, decoded_y).score pbar.set_description(f'cur_batch_bleu: {bleu:.4f}') acc += bleu num_batches += 1 acc /= num_batches return acc def infer(self, x, *args, **kwargs): if 'token_type_ids' in x.keys(): del x['token_type_ids'] if 'generate' in kwargs: return self.models_dict['main'].generate( x['input_ids'], attention_mask=x["attention_mask"], max_length=512 ) return self.models_dict['main'](**x) from nltk.metrics import accuracy as nltk_acc class ElasticDNN_OfflineTokenClsFMModel(ElasticDNN_OfflineFMModel): def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) # print(x) y = y.to(self.device) output = self.infer(x) # torch.Size([16, 512, 43]) torch.Size([16, 512]) for oi, yi, xi in zip(output, y, x['input_ids']): # oi: 512, 43; yi: 512 seq_len = xi.nonzero().size(0) # print(output.size(), y.size()) pred = F.softmax(oi, dim=-1).argmax(dim=-1) correct = torch.eq(pred[1: seq_len], yi[1: seq_len]).sum().item() # print(output.size(), y.size()) acc += correct sample_num += seq_len pbar.set_description(f'seq_len: {seq_len}, cur_seq_acc: {(correct / seq_len):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'](**x) class ElasticDNN_OfflineMMClsFMModel(ElasticDNN_OfflineFMModel): # def __init__(self, name: str, models_dict_path: str, device: str, class_to_label_idx_map): # super().__init__(name, models_dict_path, device) # self.class_to_label_idx_map = class_to_label_idx_map def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() batch_size = 1 with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: if batch_index * batch_size > 2000: break for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) y = y.to(self.device) # print(x) raw_texts = x['texts'][:] x['texts'] = list(set(x['texts'])) # print(x['texts']) batch_size = len(y) x['for_training'] = False output = self.infer(x) output = output.logits_per_image # print(output.size()) # exit() # y = torch.arange(len(y), device=self.device) y = torch.LongTensor([x['texts'].index(rt) for rt in raw_texts]).to(self.device) # print(y) # exit() # print(output.size(), y.size()) pred = F.softmax(output, dim=1).argmax(dim=1) correct = torch.eq(pred, y).sum().item() acc += correct sample_num += len(y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_correct: {correct}, ' f'cur_batch_acc: {(correct / len(y)):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): x['for_training'] = self.models_dict['main'].training return self.models_dict['main'](**x) class VQAScore: def __init__(self): # self.add_state("score", default=torch.tensor(0.0), dist_reduce_fx="sum") # self.add_state("total", default=torch.tensor(0.0), dist_reduce_fx="sum") self.score = torch.tensor(0.0) self.total = torch.tensor(0.0) def update(self, logits, target): logits, target = ( logits.detach().float().to(self.score.device), target.detach().float().to(self.score.device), ) logits = torch.max(logits, 1)[1] one_hots = torch.zeros(*target.size()).to(target) one_hots.scatter_(1, logits.view(-1, 1), 1) scores = one_hots * target self.score += scores.sum() self.total += len(logits) def compute(self): return self.score / self.total class ElasticDNN_OfflineVQAFMModel(ElasticDNN_OfflineFMModel): # def __init__(self, name: str, models_dict_path: str, device: str, class_to_label_idx_map): # super().__init__(name, models_dict_path, device) # self.class_to_label_idx_map = class_to_label_idx_map def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 vqa_score = VQAScore() self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) y = y.to(self.device) output = self.infer(x).logits vqa_score.update(output, y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_acc: {vqa_score.compute():.4f}') return vqa_score.compute() def infer(self, x, *args, **kwargs): return self.models_dict['main'](**x) class ElasticDNN_OfflineMDModel(BaseModel): def get_required_model_components(self) -> List[str]: return ['main'] @abstractmethod def forward_to_get_task_loss(self, x, y, *args, **kwargs): raise NotImplementedError @abstractmethod def get_feature_hook(self) -> LayerActivation: pass @abstractmethod def get_distill_loss(self, student_output, teacher_output): pass @abstractmethod def get_matched_param_of_fm(self, self_param_name, fm: nn.Module): pass class ElasticDNN_OfflineClsMDModel(ElasticDNN_OfflineMDModel): def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: x, y = x.to(self.device), y.to(self.device) output = self.infer(x) pred = F.softmax(output, dim=1).argmax(dim=1) correct = torch.eq(pred, y).sum().item() acc += correct sample_num += len(y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_correct: {correct}, ' f'cur_batch_acc: {(correct / len(y)):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'](x) class ElasticDNN_OfflineSegMDModel(ElasticDNN_OfflineMDModel): def __init__(self, name: str, models_dict_path: str, device: str, num_classes): super().__init__(name, models_dict_path, device) self.num_classes = num_classes def get_accuracy(self, test_loader, *args, **kwargs): device = self.device self.to_eval_mode() metrics = StreamSegMetrics(self.num_classes) metrics.reset() import tqdm pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), leave=False, dynamic_ncols=True) with torch.no_grad(): for batch_index, (x, y) in pbar: x, y = x.to(device, dtype=x.dtype, non_blocking=True, copy=False), \ y.to(device, dtype=y.dtype, non_blocking=True, copy=False) output = self.infer(x) pred = output.detach().max(dim=1)[1].cpu().numpy() metrics.update((y + 0).cpu().numpy(), pred) res = metrics.get_results() pbar.set_description(f'cur batch mIoU: {res["Mean IoU"]:.4f}') res = metrics.get_results() return res['Mean IoU'] def infer(self, x, *args, **kwargs): return self.models_dict['main'](x) class ElasticDNN_OfflineDetMDModel(ElasticDNN_OfflineMDModel): def __init__(self, name: str, models_dict_path: str, device: str, num_classes): super().__init__(name, models_dict_path, device) self.num_classes = num_classes def get_accuracy(self, test_loader, *args, **kwargs): # print('DeeplabV3: start test acc') _d = test_loader.dataset from data import build_dataloader if _d.__class__.__name__ == 'MergedDataset': # print('\neval on merged datasets') datasets = _d.datasets test_loaders = [build_dataloader(d, test_loader.batch_size, test_loader.num_workers, False, None) for d in datasets] accs = [self.get_accuracy(loader) for loader in test_loaders] # print(accs) return sum(accs) / len(accs) # print('dataset len', len(test_loader.dataset)) model = self.models_dict['main'] device = self.device model.eval() # print('# classes', model.num_classes) model = model.to(device) from dnns.yolov3.coco_evaluator import COCOEvaluator from utils.common.others import HiddenPrints with torch.no_grad(): with HiddenPrints(): evaluator = COCOEvaluator( dataloader=test_loader, img_size=(224, 224), confthre=0.01, nmsthre=0.65, num_classes=self.num_classes, testdev=False ) res = evaluator.evaluate(model, False, False) map50 = res[1] # print('eval info', res[-1]) return map50 def infer(self, x, *args, **kwargs): if len(args) > 0: return self.models_dict['main'](x, *args) # forward(x, label) return self.models_dict['main'](x) class ElasticDNN_OfflineSenClsMDModel(ElasticDNN_OfflineMDModel): def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) y = y.to(self.device) output = self.infer(x) pred = F.softmax(output, dim=1).argmax(dim=1) # print(pred, y) correct = torch.eq(pred, y).sum().item() acc += correct sample_num += len(y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_correct: {correct}, ' f'cur_batch_acc: {(correct / len(y)):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'](**x) class ElasticDNN_OfflineTrMDModel(ElasticDNN_OfflineMDModel): def get_accuracy(self, test_loader, *args, **kwargs): # TODO: BLEU from sacrebleu import corpus_bleu acc = 0 num_batches = 0 self.to_eval_mode() from data.datasets.sentiment_classification.global_bert_tokenizer import get_tokenizer tokenizer = get_tokenizer() def _decode(o): # https://github.com/huggingface/transformers/blob/main/examples/research_projects/seq2seq-distillation/finetune.py#L133 o = tokenizer.batch_decode(o, skip_special_tokens=True, clean_up_tokenization_spaces=True) return [oi.strip().replace(' ', '') for oi in o] with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) y = y.to(self.device) output = self.infer(x) decoded_output = _decode(output.argmax(-1)) decoded_y = _decode(y) decoded_y = [decoded_y] # print(x, decoded_y, decoded_output, output.argmax(-1)) bleu = corpus_bleu(decoded_output, decoded_y).score pbar.set_description(f'cur_batch_bleu: {bleu:.4f}') acc += bleu num_batches += 1 acc /= num_batches return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'](**x) class ElasticDNN_OfflineTokenClsMDModel(ElasticDNN_OfflineMDModel): def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) # print(x) y = y.to(self.device) output = self.infer(x) # torch.Size([16, 512, 43]) torch.Size([16, 512]) for oi, yi, xi in zip(output, y, x['input_ids']): # oi: 512, 43; yi: 512 seq_len = xi.nonzero().size(0) # print(output.size(), y.size()) pred = F.softmax(oi, dim=-1).argmax(dim=-1) correct = torch.eq(pred[1: seq_len], yi[1: seq_len]).sum().item() # print(output.size(), y.size()) acc += correct sample_num += seq_len pbar.set_description(f'seq_len: {seq_len}, cur_seq_acc: {(correct / seq_len):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'](**x) class ElasticDNN_OfflineMMClsMDModel(ElasticDNN_OfflineMDModel): def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 self.to_eval_mode() batch_size = 1 with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: if batch_index * batch_size > 2000: break for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) y = y.to(self.device) # print(x) raw_texts = x['texts'][:] x['texts'] = list(set(x['texts'])) # print(x['texts']) batch_size = len(y) x['for_training'] = False output = self.infer(x) output = output.logits_per_image # print(output.size()) # exit() # y = torch.arange(len(y), device=self.device) y = torch.LongTensor([x['texts'].index(rt) for rt in raw_texts]).to(self.device) # print(y) # exit() # print(output.size(), y.size()) pred = F.softmax(output, dim=1).argmax(dim=1) correct = torch.eq(pred, y).sum().item() acc += correct sample_num += len(y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_correct: {correct}, ' f'cur_batch_acc: {(correct / len(y)):.4f}') acc /= sample_num return acc def infer(self, x, *args, **kwargs): return self.models_dict['main'](**x) class ElasticDNN_OfflineVQAMDModel(ElasticDNN_OfflineMDModel): # def __init__(self, name: str, models_dict_path: str, device: str, class_to_label_idx_map): # super().__init__(name, models_dict_path, device) # self.class_to_label_idx_map = class_to_label_idx_map def get_accuracy(self, test_loader, *args, **kwargs): acc = 0 sample_num = 0 vqa_score = VQAScore() self.to_eval_mode() with torch.no_grad(): pbar = tqdm.tqdm(enumerate(test_loader), total=len(test_loader), dynamic_ncols=True, leave=False) for batch_index, (x, y) in pbar: for k, v in x.items(): if isinstance(v, torch.Tensor): x[k] = v.to(self.device) y = y.to(self.device) output = self.infer(x).logits vqa_score.update(output, y) pbar.set_description(f'cur_batch_total: {len(y)}, cur_batch_acc: {vqa_score.compute():.4f}') return vqa_score.compute() def infer(self, x, *args, **kwargs): return self.models_dict['main'](**x)