(cfg, model, critetion, optimizer, epoch, data_loader, logger, writer, devices, trainIters, metrics_base='combine')
| 47 | |
| 48 | |
| 49 | def train(cfg, model, critetion, optimizer, epoch, data_loader, logger, writer, devices, trainIters, metrics_base='combine'): |
| 50 | calculate_acc = get_acc_mesure_func(metrics_base) |
| 51 | batch_time = AverageMeter() |
| 52 | data_time = AverageMeter() |
| 53 | losses = AverageMeter() |
| 54 | acc = AverageMeter() |
| 55 | |
| 56 | #Switch to train mode |
| 57 | model.train() |
| 58 | data_loader = tqdm(data_loader, dynamic_ncols=True) |
| 59 | start = time.time() |
| 60 | for i, batch_data in enumerate(data_loader): |
| 61 | inputs, labels, targets, heatmaps, cstency_heatmaps, offsets = get_batch_data(batch_data) |
| 62 | inputs = inputs.cuda().to(non_blocking=True, dtype=torch.float64) |
| 63 | #Measuring data loading time |
| 64 | data_time.update(time.time() - start) |
| 65 | |
| 66 | loop = arange(1) if cfg.TRAIN.optimizer != 'SAM' else arange(2) |
| 67 | for idx in loop: |
| 68 | outputs = model(inputs) |
| 69 | if isinstance(outputs, list): |
| 70 | outputs = outputs[0] |
| 71 | #In case outputs contain a dict key |
| 72 | if isinstance(outputs, dict): |
| 73 | outputs_hm = outputs['hm'] |
| 74 | outputs_cls = outputs['cls'] |
| 75 | outputs_offset = outputs['offset'] if 'offset' in outputs.keys() else None |
| 76 | outputs_cstency = outputs['cstency'] if 'cstency' in outputs.keys() else None |
| 77 | |
| 78 | if idx == 0: |
| 79 | first_outputs_hm = outputs_hm |
| 80 | first_outputs_cls = outputs_cls |
| 81 | |
| 82 | if 'Combined' in cfg.TRAIN.loss.type: |
| 83 | labels = labels.cuda().to(non_blocking=True, dtype=torch.float64) |
| 84 | # labels = labels.cuda().to(non_blocking=True).long() |
| 85 | |
| 86 | if offsets is not None: |
| 87 | offsets = offsets.cuda().to(non_blocking=True, dtype=torch.float64) |
| 88 | |
| 89 | if cstency_heatmaps is not None: |
| 90 | cstency_heatmaps = cstency_heatmaps.cuda().to(non_blocking=True, dtype=torch.float64) |
| 91 | |
| 92 | if cfg.TRAIN.loss.type != 'CombinedHeatmapBinaryLoss': |
| 93 | heatmaps = heatmaps.cuda().to(non_blocking=True, dtype=torch.float64) |
| 94 | else: |
| 95 | heatmaps = targets.cuda().to(non_blocking=True, dtype=torch.float64) |
| 96 | |
| 97 | loss_ = critetion(outputs_hm, heatmaps, outputs_cls.sigmoid(), labels, |
| 98 | offset_preds=outputs_offset, |
| 99 | offset_gts=offsets, |
| 100 | cstency_preds=outputs_cstency, |
| 101 | cstency_gts=cstency_heatmaps) |
| 102 | loss = loss_['hm'] |
| 103 | if 'cls' in loss_.keys(): |
| 104 | loss += loss_['cls'] |
| 105 | if 'dst_hm_cls' in loss_.keys(): |
| 106 | loss += loss_['dst_hm_cls'] |
no test coverage detected