| """Training and validation for Ref-AVS (text + audio + SAM2 multimask decoding).""" |
| import numpy |
| import torch |
| from torch.utils.data import DataLoader |
| from tqdm import tqdm |
|
|
| _DECODE_MODES = frozenset({'', 'iou_select', 'iou_occ_select'}) |
|
|
|
|
| def _decode_mode_and_wandb_tag(process): |
| """Match tmp.code: `process` is decode mode for known strings; else Ref split tag + default decode.""" |
| if process in _DECODE_MODES: |
| return process, process |
| return 'iou_select', process |
|
|
|
|
| class Trainer: |
| """Train / valid / null-valid steps with composite loss, contrastive term, and metrics.""" |
|
|
| def __init__(self, hyp_param, loss, tensorboard, metrics): |
| self.param = hyp_param |
| self.loss = loss |
| self.tensorboard = tensorboard |
| self.metrics = metrics |
| from loss.training.contrastive_learning import ContrastLoss |
| self.cl = ContrastLoss(self.param) |
|
|
| @torch.no_grad() |
| def valid_null(self, epoch, dataloader, model, process='test_n'): |
| if not isinstance(dataloader, DataLoader): |
| raise TypeError("valid_null() expects a torch.utils.data.DataLoader (do not pass iter(dataloader) first).") |
| decode_mode, wandb_tag = _decode_mode_and_wandb_tag(process) |
| self.metrics['foreground_s'].reset() |
| dataloader_length = len(dataloader) |
| tbar = range(dataloader_length) |
| tbar = tqdm(tbar, ncols=135) if self.param.local_rank <= 0 else tbar |
| p_pool = [None] * self.param.gpus |
| n_pool = [None] * self.param.gpus |
|
|
| data_iter = iter(dataloader) |
| for _ in tbar: |
| items = next(data_iter) |
| frame, spect, prompt_dicts = items['frame'], items['spectrogram'], items['text'] |
| logits = [] |
| for frame_, spect_, prompt_dicts_ in zip(frame, spect, prompt_dicts): |
| frame_ = frame_.cuda(self.param.local_rank, non_blocking=True) |
| spect_ = spect_.cuda(self.param.local_rank, non_blocking=True) |
| prompt_dicts_ = [prompt_dicts_] |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| outputs, _ = model.module(frame_, spect_, prompt_dicts_, sam_process=False) |
|
|
| logits_ = torch.cat([torch.cat(i['multistep_pred_multimasks_high_res']) for i in outputs]) |
| ious_scores = torch.cat([torch.cat(i['multistep_pred_ious']) for i in outputs]) |
| occ_scores = torch.cat([torch.cat(i['multistep_object_score_logits']) for i in outputs]) |
| if decode_mode == 'iou_select': |
| ious_scores = torch.argmax(ious_scores, dim=1) |
| logits_ = logits_[torch.arange(0, frame_.shape[0]), ious_scores, ...] |
| elif decode_mode == 'iou_occ_select': |
| ious_scores = torch.argmax(ious_scores, dim=1) |
| logits_ = logits_[torch.arange(0, frame_.shape[0]), ious_scores, ...] |
| logits_[occ_scores.squeeze() < 0, ...] = 0. |
| else: |
| logits_ = logits_[:, 0, ...] |
| logits.append(logits_) |
|
|
| logits = torch.cat(logits).reshape(frame.shape[0], -1, self.param.image_size, self.param.image_size) |
| if len(logits.shape) == 3: |
| logits = logits.unsqueeze(1) |
|
|
| foreground_s = self.metrics['foreground_s'].metric_s_for_null(logits, get_entire_list=True) |
| torch.distributed.all_gather_object(p_pool, foreground_s['foreground_p']) |
| torch.distributed.all_gather_object(n_pool, foreground_s['foreground_n']) |
| foreground_s = sum([i[0].cpu() for i in p_pool]) / sum([i[0] for i in n_pool]) |
|
|
| if self.param.local_rank <= 0: |
| tbar.set_description( |
| 'epoch {} | valid.null_s {}'.format(epoch, numpy.round(foreground_s, 5)), |
| ) |
| torch.cuda.empty_cache() |
|
|
| final_s = foreground_s |
| if self.param.local_rank <= 0 and self.tensorboard is not None: |
| self.tensorboard.upload_wandb_info({"valid.f_s/{}".format(wandb_tag): final_s}) |
|
|
| return numpy.round(final_s, 5) |
|
|
| @torch.no_grad() |
| def valid(self, epoch, dataloader, model, process='iou_select'): |
| """Evaluate IoU / F-score; `process` is decode mode (tmp) or split tag (test_s / test_u). Wandb keys like tmp.""" |
| if not isinstance(dataloader, DataLoader): |
| raise TypeError("valid() expects a torch.utils.data.DataLoader (do not pass iter(dataloader) first).") |
| decode_mode, wandb_tag = _decode_mode_and_wandb_tag(process) |
| self.metrics['foreground_iou'].reset() |
| self.metrics['foreground_f-score'].reset() |
| dataloader_length = len(dataloader) |
| tbar = range(dataloader_length) |
| tbar = tqdm(tbar, ncols=135) if self.param.local_rank <= 0 else tbar |
| iou_pool = [None] * self.param.gpus |
| fscore_pool = [None] * self.param.gpus |
|
|
| data_iter = iter(dataloader) |
| for _ in tbar: |
| items = next(data_iter) |
| frame, spect, label, prompt_dicts = ( |
| items['frame'], items['spectrogram'], items['label'], items['text'] |
| ) |
| logits = [] |
| labels = [] |
| for frame_, spect_, label_, prompt_dicts_ in zip(frame, spect, label, prompt_dicts): |
| frame_ = frame_.cuda(self.param.local_rank, non_blocking=True) |
| spect_ = spect_.cuda(self.param.local_rank, non_blocking=True) |
| label_ = label_.cuda(self.param.local_rank, non_blocking=True) |
| prompt_dicts_ = [prompt_dicts_] |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| outputs, _ = model.module(frame_, spect_, prompt_dicts_, sam_process=False) |
|
|
| logits_ = torch.cat([torch.cat(i['multistep_pred_multimasks_high_res']) for i in outputs]) |
| ious_scores = torch.cat([torch.cat(i['multistep_pred_ious']) for i in outputs]) |
| occ_scores = torch.cat([torch.cat(i['multistep_object_score_logits']) for i in outputs]) |
| if decode_mode == 'iou_select': |
| ious_scores = torch.argmax(ious_scores, dim=1) |
| logits_ = logits_[torch.arange(0, frame_.shape[0]), ious_scores, ...] |
| elif decode_mode == 'iou_occ_select': |
| ious_scores = torch.argmax(ious_scores, dim=1) |
| logits_ = logits_[torch.arange(0, frame_.shape[0]), ious_scores, ...] |
| logits_[occ_scores.squeeze() < 0, ...] = 0. |
| else: |
| logits_ = logits_[:, 0, ...] |
| logits.append(logits_) |
| labels.append(label_) |
|
|
| logits = torch.cat(logits) |
| labels = torch.cat(labels) |
| foreground_iou_rank = self.metrics['foreground_iou'].calculate_iou( |
| (logits > 0.).squeeze().long(), labels.squeeze().long(), get_entire_list=True, |
| ) |
| foreground_f_score_rank = self.metrics['foreground_f-score'].calculate_f_score( |
| logits.squeeze(), labels.squeeze().long(), get_entire_list=True, |
| ) |
| torch.distributed.all_gather_object(iou_pool, foreground_iou_rank) |
| torch.distributed.all_gather_object(fscore_pool, foreground_f_score_rank) |
| foreground_iou = sum([i['foreground_iou'][0].cpu() for i in iou_pool]) / sum( |
| [i['foreground_iou'][1] for i in iou_pool]) |
| foreground_f_score = sum([i['foreground_f-score'][0] for i in fscore_pool]) / sum( |
| [i['foreground_f-score'][1] for i in fscore_pool]) |
|
|
| if self.param.local_rank <= 0: |
| tbar.set_description( |
| 'epoch {} | valid.f_iou {}, valid.f_f-score {}'.format( |
| epoch, |
| numpy.round(foreground_iou.cpu().numpy(), 5), |
| numpy.round(foreground_f_score, 5), |
| ), |
| ) |
| torch.cuda.empty_cache() |
|
|
| final_iou = foreground_iou |
| final_fscore = foreground_f_score |
| if self.param.local_rank <= 0 and self.tensorboard is not None: |
| self.tensorboard.upload_wandb_info({ |
| "valid.f_iou/{}".format(wandb_tag): final_iou, |
| "valid.f_f-score/{}".format(wandb_tag): final_fscore, |
| }) |
|
|
| def _to_float(x): |
| if isinstance(x, torch.Tensor): |
| return float(x.detach().cpu().item()) |
| return float(x) |
|
|
| return numpy.round(_to_float(final_iou), 5), numpy.round(_to_float(final_fscore), 5) |
|
|
| def train(self, epoch, dataloader, model, optimiser): |
| if not isinstance(dataloader, DataLoader): |
| raise TypeError("train() expects a torch.utils.data.DataLoader (do not pass iter(dataloader) first).") |
| dataloader_length = len(dataloader) |
| tbar = range(dataloader_length) |
| tbar = tqdm(tbar, ncols=135) if self.param.local_rank <= 0 else tbar |
|
|
| data_iter = iter(dataloader) |
| for batch_index in tbar: |
| current_index = dataloader_length * epoch + batch_index |
| items = next(data_iter) |
| frame, spect, label, prompt_dicts = ( |
| items['frame'], items['spectrogram'], items['label'], items['text'], |
| ) |
| frame = torch.flatten(frame, start_dim=0, end_dim=1).cuda(self.param.local_rank, non_blocking=True) |
| spect = torch.flatten(spect, start_dim=0, end_dim=1).cuda(self.param.local_rank, non_blocking=True) |
| label = torch.flatten(label, start_dim=0, end_dim=1).cuda(self.param.local_rank, non_blocking=True) |
|
|
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| outputs, proj_feats = model(frame, spect, prompt_dicts, sam_process=False) |
| loss_dict = self.loss(outputs, label.unsqueeze(1)) |
| cl_loss = self.cl(proj_feats, outputs, label) |
|
|
| optimiser.zero_grad() |
| (loss_dict['core_loss'] + cl_loss).backward() |
| optimiser.step() |
|
|
| current_lr = self.param.lr * (1 - current_index / (dataloader_length * self.param.epochs)) ** 0.9 |
| for params_lr in optimiser.param_groups: |
| names = params_lr.get("name", []) |
| if names and any("vgg" in n for n in names): |
| params_lr['lr'] = current_lr * 0.1 |
| else: |
| params_lr['lr'] = current_lr |
|
|
| if self.param.local_rank <= 0 and self.tensorboard is not None: |
| logits = torch.cat([i['multistep_pred_multimasks_high_res'][0] for i in outputs]) |
| foreground_iou = self.metrics['foreground_iou'].calculate_iou( |
| (logits > 0)[:, 0, ...].long(), label.long(), |
| ) |
| self.tensorboard.upload_wandb_info({ |
| "loss": loss_dict['core_loss'].item(), "f_iou": foreground_iou.item(), |
| "lr": optimiser.param_groups[0]['lr'], |
| "loss_dice": loss_dict['loss_dice'], |
| "loss_focal": loss_dict['loss_mask'], |
| "loss_contras": cl_loss.item(), |
| }) |
| tbar.set_description( |
| 'epoch {} | loss {}, f_iou {}'.format( |
| epoch, loss_dict['core_loss'].item(), foreground_iou.item(), |
| ), |
| ) |
| return |
|
|
