| 30 | |
| 31 | |
| 32 | def get_batch_data(batch_data): |
| 33 | inputs = batch_data['img'] |
| 34 | labels = batch_data['label'] |
| 35 | targets = batch_data['target'] |
| 36 | heatmaps = batch_data['heatmap'] |
| 37 | cstency_heatmaps = None |
| 38 | offsets = None |
| 39 | |
| 40 | if 'cstency_heatmap' in batch_data: |
| 41 | cstency_heatmaps = batch_data['cstency_heatmap'] |
| 42 | |
| 43 | if 'offset' in batch_data: |
| 44 | offsets = batch_data['offset'] |
| 45 | |
| 46 | return inputs, labels, targets, heatmaps, cstency_heatmaps, offsets |
| 47 | |
| 48 | |
| 49 | def train(cfg, model, critetion, optimizer, epoch, data_loader, logger, writer, devices, trainIters, metrics_base='combine'): |