def evaluate_and_save_model(self):
evalmodel = self.ema_model.ema if self.use_model_ema else self.model
ap50_95, ap50, summary = self.exp.eval(evalmodel, self.evaluator,
self.is_distributed)
self.model.train()
if self.rank == 0:
self.tblogger.add_scalar("val/COCOAP50", ap50, self.epoch + 1)
self.tblogger.add_scalar("val/COCOAP50_95", ap50_95,
self.epoch + 1)
logger.info("\n" + summary)
synchronize()
self.save_ckpt("last_epoch", ap50_95 > self.best_ap)
self.best_ap = max(self.best_ap, ap50_95)
def random_resize(data_loader, exp, epoch, rank, is_distributed):
tensor = torch.LongTensor(1).cuda()
if is_distributed:
synchronize()
if rank == 0:
if epoch > exp.max_epoch - 10:
size = exp.input_size
else:
size = random.randint(*exp.random_size)
size = int(32 * size)
tensor.fill_(size)
if is_distributed:
synchronize()
dist.broadcast(tensor, 0)
input_size = data_loader.change_input_dim(multiple=tensor.item(),
random_range=None)
return input_size
def evaluate(self,
model,
distributed=False,
half=False,
trt_file=None,
decoder=None,
test_size=None):
"""
VOC average precision (AP) Evaluation. Iterate inference on the test dataset
and the results are evaluated by COCO API.
NOTE: This function will change training mode to False, please save states if needed.
Args:
model : model to evaluate.
Returns:
ap50_95 (float) : COCO style AP of IoU=50:95
ap50 (float) : VOC 2007 metric AP of IoU=50
summary (sr): summary info of evaluation.
"""
# TODO half to amp_test
tensor_type = torch.cuda.HalfTensor if half else torch.cuda.FloatTensor
model = model.eval()
if half:
model = model.half()
ids = []
data_list = {}
progress_bar = tqdm if is_main_process() else iter
inference_time = 0
nms_time = 0
n_samples = len(self.dataloader) - 1
if trt_file is not None:
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load(trt_file))
x = torch.ones(1, 3, test_size[0], test_size[1]).cuda()
model(x)
model = model_trt
for cur_iter, (imgs, _, info_imgs,
ids) in enumerate(progress_bar(self.dataloader)):
with torch.no_grad():
imgs = imgs.type(tensor_type)
# skip the the last iters since batchsize might be not enough for batch inference
is_time_record = cur_iter < len(self.dataloader) - 1
if is_time_record:
start = time.time()
outputs = model(imgs)
if decoder is not None:
outputs = decoder(outputs, dtype=outputs.type())
if is_time_record:
infer_end = time_synchronized()
inference_time += infer_end - start
outputs = postprocess(outputs, self.num_classes, self.confthre,
self.nmsthre)
if is_time_record:
nms_end = time_synchronized()
nms_time += nms_end - infer_end
data_list.update(
self.convert_to_voc_format(outputs, info_imgs, ids))
statistics = torch.cuda.FloatTensor(
[inference_time, nms_time, n_samples])
if distributed:
data_list = gather(data_list, dst=0)
data_list = ChainMap(*data_list)
torch.distributed.reduce(statistics, dst=0)
eval_results = self.evaluate_prediction(data_list, statistics)
synchronize()
return eval_results