| from collections import namedtuple |
|
|
| import numpy as np |
| import torch |
|
|
| from .detectors import build_detector |
|
|
| try: |
| import kornia |
| except: |
| pass |
| |
|
|
|
|
|
|
| def build_network(model_cfg, num_class, dataset): |
| model = build_detector( |
| model_cfg=model_cfg, num_class=num_class, dataset=dataset |
| ) |
| return model |
|
|
|
|
| def load_data_to_gpu(batch_dict): |
| for key, val in batch_dict.items(): |
| if key == 'camera_imgs': |
| batch_dict[key] = val.cuda() |
| elif not isinstance(val, np.ndarray): |
| continue |
| elif key in ['frame_id', 'metadata', 'calib', 'image_paths','ori_shape','img_process_infos']: |
| continue |
| elif key in ['images']: |
| batch_dict[key] = kornia.image_to_tensor(val).float().cuda().contiguous() |
| elif key in ['image_shape']: |
| batch_dict[key] = torch.from_numpy(val).int().cuda() |
| else: |
| batch_dict[key] = torch.from_numpy(val).float().cuda() |
|
|
|
|
| def model_fn_decorator(): |
| ModelReturn = namedtuple('ModelReturn', ['loss', 'tb_dict', 'disp_dict']) |
|
|
| def model_func(model, batch_dict): |
| load_data_to_gpu(batch_dict) |
| ret_dict, tb_dict, disp_dict = model(batch_dict) |
|
|
| loss = ret_dict['loss'].mean() |
| if hasattr(model, 'update_global_step'): |
| model.update_global_step() |
| else: |
| model.module.update_global_step() |
|
|
| return ModelReturn(loss, tb_dict, disp_dict) |
|
|
| return model_func |
|
|
