def main(logger, args):
if not torch.cuda.is_available():
raise Exception("need gpu to train network!")
torch.cuda.empty_cache()
if args.seed is not None:
random.seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
cudnn.deterministic = True
gpus = torch.cuda.device_count()
logger.info(f'use {gpus} gpus')
logger.info(f"args: {args}")
cudnn.benchmark = True
cudnn.enabled = True
start_time = time.time()
# dataset and dataloader
logger.info('start loading data')
collater = Collater()
train_loader = DataLoader(Config.train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=collater.next)
logger.info('finish loading data')
model = retinanet.__dict__[args.network](**{
"pretrained": args.pretrained,
"num_classes": args.num_classes,
})
for name, param in model.named_parameters():
logger.info(f"{name},{param.requires_grad}")
flops_input = torch.randn(1, 3, args.input_image_size,
args.input_image_size)
flops, params = profile(model, inputs=(flops_input, ))
flops, params = clever_format([flops, params], "%.3f")
logger.info(f"model: '{args.network}', flops: {flops}, params: {params}")
criterion = RetinaLoss(image_w=args.input_image_size,
image_h=args.input_image_size).cuda()
decoder = RetinaDecoder(image_w=args.input_image_size,
image_h=args.input_image_size).cuda()
model = model.cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
if args.apex:
amp.register_float_function(torch, 'sigmoid')
amp.register_float_function(torch, 'softmax')
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
model = nn.DataParallel(model)
if args.evaluate:
if not os.path.isfile(args.evaluate):
raise Exception(
f"{args.resume} is not a file, please check it again")
logger.info('start only evaluating')
logger.info(f"start resuming model from {args.evaluate}")
checkpoint = torch.load(args.evaluate,
map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state_dict'])
logger.info(f"start eval.")
all_eval_result = validate(Config.val_dataset, model, decoder, args)
logger.info(f"eval done.")
if all_eval_result is not None:
logger.info(
f"val: epoch: {checkpoint['epoch']:0>5d}, IoU=0.5:0.95,area=all,maxDets=100,mAP:{all_eval_result[0]:.3f}, IoU=0.5,area=all,maxDets=100,mAP:{all_eval_result[1]:.3f}, IoU=0.75,area=all,maxDets=100,mAP:{all_eval_result[2]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAP:{all_eval_result[3]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAP:{all_eval_result[4]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAP:{all_eval_result[5]:.3f}, IoU=0.5:0.95,area=all,maxDets=1,mAR:{all_eval_result[6]:.3f}, IoU=0.5:0.95,area=all,maxDets=10,mAR:{all_eval_result[7]:.3f}, IoU=0.5:0.95,area=all,maxDets=100,mAR:{all_eval_result[8]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAR:{all_eval_result[9]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAR:{all_eval_result[10]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAR:{all_eval_result[11]:.3f}"
)
return
best_map = 0.0
start_epoch = 1
# resume training
if os.path.exists(args.resume):
logger.info(f"start resuming model from {args.resume}")
checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
start_epoch += checkpoint['epoch']
best_map = checkpoint['best_map']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
logger.info(
f"finish resuming model from {args.resume}, epoch {checkpoint['epoch']}, best_map: {checkpoint['best_map']}, "
f"loss: {checkpoint['loss']:3f}, cls_loss: {checkpoint['cls_loss']:2f}, reg_loss: {checkpoint['reg_loss']:2f}"
)
if not os.path.exists(args.checkpoints):
os.makedirs(args.checkpoints)
logger.info('start training')
for epoch in range(start_epoch, args.epochs + 1):
cls_losses, reg_losses, losses = train(train_loader, model, criterion,
optimizer, scheduler, epoch,
logger, args)
logger.info(
f"train: epoch {epoch:0>3d}, cls_loss: {cls_losses:.2f}, reg_loss: {reg_losses:.2f}, loss: {losses:.2f}"
)
if epoch % 1 == 0 or epoch == args.epochs:
logger.info(f"start eval.")
all_eval_result = validate(Config.val_dataset, model, decoder,
args)
logger.info(f"eval done.")
if all_eval_result is not None:
logger.info(
f"val: epoch: {epoch:0>5d}, IoU=0.5:0.95,area=all,maxDets=100,mAP:{all_eval_result[0]:.3f}, IoU=0.5,area=all,maxDets=100,mAP:{all_eval_result[1]:.3f}, IoU=0.75,area=all,maxDets=100,mAP:{all_eval_result[2]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAP:{all_eval_result[3]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAP:{all_eval_result[4]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAP:{all_eval_result[5]:.3f}, IoU=0.5:0.95,area=all,maxDets=1,mAR:{all_eval_result[6]:.3f}, IoU=0.5:0.95,area=all,maxDets=10,mAR:{all_eval_result[7]:.3f}, IoU=0.5:0.95,area=all,maxDets=100,mAR:{all_eval_result[8]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAR:{all_eval_result[9]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAR:{all_eval_result[10]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAR:{all_eval_result[11]:.3f}"
)
if all_eval_result[0] > best_map:
torch.save(model.module.state_dict(),
os.path.join(args.checkpoints, "best.pth"))
best_map = all_eval_result[0]
torch.save(
{
'epoch': epoch,
'best_map': best_map,
'cls_loss': cls_losses,
'reg_loss': reg_losses,
'loss': losses,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
}, os.path.join(args.checkpoints, 'latest.pth'))
logger.info(f"finish training, best_map: {best_map:.3f}")
training_time = (time.time() - start_time) / 3600
logger.info(
f"finish training, total training time: {training_time:.2f} hours")
def main():
args = parse_args()
global local_rank
local_rank = args.local_rank
if local_rank == 0:
global logger
logger = get_logger(__name__, args.log)
torch.cuda.empty_cache()
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
cudnn.deterministic = True
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl', init_method='env://')
global gpus_num
gpus_num = torch.cuda.device_count()
if local_rank == 0:
logger.info(f'use {gpus_num} gpus')
logger.info(f"args: {args}")
cudnn.benchmark = True
cudnn.enabled = True
start_time = time.time()
# dataset and dataloader
if local_rank == 0:
logger.info('start loading data')
train_sampler = torch.utils.data.distributed.DistributedSampler(
Config.train_dataset, shuffle=True)
train_loader = DataLoader(Config.train_dataset,
batch_size=args.per_node_batch_size,
shuffle=False,
num_workers=args.num_workers,
collate_fn=collater,
sampler=train_sampler)
if local_rank == 0:
logger.info('finish loading data')
model = retinanet.__dict__[args.network](**{
"pretrained": args.pretrained,
"num_classes": args.num_classes,
})
for name, param in model.named_parameters():
if local_rank == 0:
logger.info(f"{name},{param.requires_grad}")
flops_input = torch.randn(1, 3, args.input_image_size,
args.input_image_size)
flops, params = profile(model, inputs=(flops_input, ))
flops, params = clever_format([flops, params], "%.3f")
if local_rank == 0:
logger.info(
f"model: '{args.network}', flops: {flops}, params: {params}")
criterion = RetinaLoss(image_w=args.input_image_size,
image_h=args.input_image_size).cuda()
decoder = RetinaDecoder(image_w=args.input_image_size,
image_h=args.input_image_size).cuda()
model = model.cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=3,
verbose=True)
if args.sync_bn:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.apex:
amp.register_float_function(torch, 'sigmoid')
amp.register_float_function(torch, 'softmax')
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
model = apex.parallel.DistributedDataParallel(model,
delay_allreduce=True)
if args.sync_bn:
model = apex.parallel.convert_syncbn_model(model)
else:
model = nn.parallel.DistributedDataParallel(model,
device_ids=[local_rank],
output_device=local_rank)
if args.evaluate:
if not os.path.isfile(args.evaluate):
if local_rank == 0:
logger.exception(
'{} is not a file, please check it again'.format(
args.resume))
sys.exit(-1)
if local_rank == 0:
logger.info('start only evaluating')
logger.info(f"start resuming model from {args.evaluate}")
checkpoint = torch.load(args.evaluate,
map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state_dict'])
if local_rank == 0:
logger.info(f"start eval.")
all_eval_result = validate(Config.val_dataset, model, decoder)
logger.info(f"eval done.")
if all_eval_result is not None:
logger.info(
f"val: epoch: {checkpoint['epoch']:0>5d}, IoU=0.5:0.95,area=all,maxDets=100,mAP:{all_eval_result[0]:.3f}, IoU=0.5,area=all,maxDets=100,mAP:{all_eval_result[1]:.3f}, IoU=0.75,area=all,maxDets=100,mAP:{all_eval_result[2]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAP:{all_eval_result[3]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAP:{all_eval_result[4]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAP:{all_eval_result[5]:.3f}, IoU=0.5:0.95,area=all,maxDets=1,mAR:{all_eval_result[6]:.3f}, IoU=0.5:0.95,area=all,maxDets=10,mAR:{all_eval_result[7]:.3f}, IoU=0.5:0.95,area=all,maxDets=100,mAR:{all_eval_result[8]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAR:{all_eval_result[9]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAR:{all_eval_result[10]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAR:{all_eval_result[11]:.3f}"
)
return
best_map = 0.0
start_epoch = 1
# resume training
if os.path.exists(args.resume):
if local_rank == 0:
logger.info(f"start resuming model from {args.resume}")
checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
start_epoch += checkpoint['epoch']
best_map = checkpoint['best_map']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
if local_rank == 0:
logger.info(
f"finish resuming model from {args.resume}, epoch {checkpoint['epoch']}, best_map: {checkpoint['best_map']}, "
f"loss: {checkpoint['loss']:3f}, cls_loss: {checkpoint['cls_loss']:2f}, reg_loss: {checkpoint['reg_loss']:2f}"
)
if local_rank == 0:
if not os.path.exists(args.checkpoints):
os.makedirs(args.checkpoints)
if local_rank == 0:
logger.info('start training')
for epoch in range(start_epoch, args.epochs + 1):
train_sampler.set_epoch(epoch)
cls_losses, reg_losses, losses = train(train_loader, model, criterion,
optimizer, scheduler, epoch,
args)
if local_rank == 0:
logger.info(
f"train: epoch {epoch:0>3d}, cls_loss: {cls_losses:.2f}, reg_loss: {reg_losses:.2f}, loss: {losses:.2f}"
)
if epoch % 5 == 0 or epoch == args.epochs:
if local_rank == 0:
logger.info(f"start eval.")
all_eval_result = validate(Config.val_dataset, model, decoder)
logger.info(f"eval done.")
if all_eval_result is not None:
logger.info(
f"val: epoch: {epoch:0>5d}, IoU=0.5:0.95,area=all,maxDets=100,mAP:{all_eval_result[0]:.3f}, IoU=0.5,area=all,maxDets=100,mAP:{all_eval_result[1]:.3f}, IoU=0.75,area=all,maxDets=100,mAP:{all_eval_result[2]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAP:{all_eval_result[3]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAP:{all_eval_result[4]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAP:{all_eval_result[5]:.3f}, IoU=0.5:0.95,area=all,maxDets=1,mAR:{all_eval_result[6]:.3f}, IoU=0.5:0.95,area=all,maxDets=10,mAR:{all_eval_result[7]:.3f}, IoU=0.5:0.95,area=all,maxDets=100,mAR:{all_eval_result[8]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAR:{all_eval_result[9]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAR:{all_eval_result[10]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAR:{all_eval_result[11]:.3f}"
)
if all_eval_result[0] > best_map:
torch.save(model.module.state_dict(),
os.path.join(args.checkpoints, "best.pth"))
best_map = all_eval_result[0]
if local_rank == 0:
torch.save(
{
'epoch': epoch,
'best_map': best_map,
'cls_loss': cls_losses,
'reg_loss': reg_losses,
'loss': losses,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
}, os.path.join(args.checkpoints, 'latest.pth'))
if local_rank == 0:
logger.info(f"finish training, best_map: {best_map:.3f}")
training_time = (time.time() - start_time) / 3600
if local_rank == 0:
logger.info(
f"finish training, total training time: {training_time:.2f} hours")
def test_model(args):
print(args)
if args.use_gpu:
# use one Graphics card to test
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
if not torch.cuda.is_available():
raise Exception("need gpu to test network!")
torch.cuda.empty_cache()
if args.seed is not None:
random.seed(args.seed)
if args.use_gpu:
torch.cuda.manual_seed_all(args.seed)
cudnn.deterministic = True
if args.use_gpu:
cudnn.benchmark = True
cudnn.enabled = True
if args.detector == "retinanet":
model = _retinanet(args.backbone, args.pretrained_model_path,
args.num_classes)
decoder = RetinaDecoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "fcos":
model = _fcos(args.backbone, args.pretrained_model_path,
args.num_classes)
decoder = FCOSDecoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "centernet":
model = _centernet(args.backbone, args.pretrained_model_path,
args.num_classes)
decoder = CenterNetDecoder(
image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "yolov3":
model = _yolov3(args.backbone, args.pretrained_model_path,
args.num_classes)
decoder = YOLOV3Decoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
else:
print("unsupport detection model!")
return
flops_input = torch.randn(1, 3, args.input_image_size,
args.input_image_size)
flops, params = profile(model, inputs=(flops_input, ))
flops, params = clever_format([flops, params], "%.3f")
print(
f"backbone:{args.backbone},detector: '{args.detector}', flops: {flops}, params: {params}"
)
model.eval()
if args.use_gpu:
model = model.cuda()
decoder = decoder.cuda()
model = nn.DataParallel(model)
# load image and image preprocessing
img = cv2.imread(args.test_image_path)
origin_img = img
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
height, width, _ = img.shape
max_image_size = max(height, width)
resize_factor = args.input_image_size / max_image_size
resize_height, resize_width = int(height * resize_factor), int(
width * resize_factor)
img = cv2.resize(img, (resize_width, resize_height))
resized_img = np.zeros((args.input_image_size, args.input_image_size, 3))
resized_img[0:resize_height, 0:resize_width] = img
scale = resize_factor
resized_img = torch.tensor(resized_img)
print(resized_img.shape)
if args.use_gpu:
resized_img = resized_img.cuda()
# inference image
cls_heads, reg_heads, batch_anchors = model(
resized_img.permute(2, 0, 1).float().unsqueeze(0))
scores, classes, boxes = decoder(cls_heads, reg_heads, batch_anchors)
scores, classes, boxes = scores.cpu(), classes.cpu(), boxes.cpu()
# snap boxes to fit origin image
boxes /= scale
scores = scores.squeeze(0)
classes = classes.squeeze(0)
boxes = boxes.squeeze(0)
# draw all boxes
for per_score, per_class_index, per_box in zip(scores, classes, boxes):
per_score = per_score.numpy()
per_class_index = per_class_index.numpy().astype(np.int32)
per_box = per_box.numpy().astype(np.int32)
class_name = COCO_CLASSES[per_class_index]
color = coco_class_colors[per_class_index]
text = '{}:{:.3f}'.format(class_name, per_score)
cv2.putText(origin_img,
text, (per_box[0], per_box[1] - 10),
cv2.FONT_HERSHEY_PLAIN,
1,
color=color,
thickness=2)
cv2.rectangle(origin_img, (per_box[0], per_box[1]),
(per_box[2], per_box[3]),
color=color,
thickness=2)
if args.save_detected_image:
cv2.imwrite('detection_result.jpg', origin_img)
cv2.imshow('detection_result', origin_img)
cv2.waitKey(0)
cv2.destroyAllWindows()
return
def test_model(args):
print(args)
if args.use_gpu:
# use one Graphics card to test
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
if not torch.cuda.is_available():
raise Exception("need gpu to test network!")
torch.cuda.empty_cache()
if args.seed is not None:
random.seed(args.seed)
if args.use_gpu:
torch.cuda.manual_seed_all(args.seed)
cudnn.deterministic = True
if args.use_gpu:
cudnn.benchmark = True
cudnn.enabled = True
coco_val_dataset = CocoDetection(
image_root_dir=os.path.join(COCO2017_path, 'images/val2017'),
annotation_root_dir=os.path.join(COCO2017_path, 'annotations'),
set="val2017",
transform=transforms.Compose([
Normalize(),
Resize(resize=args.input_image_size),
]))
if args.detector == "retinanet":
model = _retinanet(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = RetinaDecoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "fcos":
model = _fcos(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = FCOSDecoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "centernet":
model = _centernet(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = CenterNetDecoder(
image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "yolov3":
model = _yolov3(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = YOLOV3Decoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
else:
print("unsupport detection model!")
return
flops_input = torch.randn(1, 3, args.input_image_size,
args.input_image_size)
flops, params = profile(model, inputs=(flops_input, ))
flops, params = clever_format([flops, params], "%.3f")
print(
f"backbone:{args.backbone},detector: '{args.detector}', flops: {flops}, params: {params}"
)
if args.use_gpu:
model = model.cuda()
decoder = decoder.cuda()
model = nn.DataParallel(model)
print(f"start eval.")
all_eval_result = validate(coco_val_dataset, model, decoder, args)
print(f"eval done.")
if all_eval_result is not None:
print(
f"val: backbone: {args.backbone}, detector: {args.detector}, IoU=0.5:0.95,area=all,maxDets=100,mAP:{all_eval_result[0]:.3f}, IoU=0.5,area=all,maxDets=100,mAP:{all_eval_result[1]:.3f}, IoU=0.75,area=all,maxDets=100,mAP:{all_eval_result[2]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAP:{all_eval_result[3]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAP:{all_eval_result[4]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAP:{all_eval_result[5]:.3f}, IoU=0.5:0.95,area=all,maxDets=1,mAR:{all_eval_result[6]:.3f}, IoU=0.5:0.95,area=all,maxDets=10,mAR:{all_eval_result[7]:.3f}, IoU=0.5:0.95,area=all,maxDets=100,mAR:{all_eval_result[8]:.3f}, IoU=0.5:0.95,area=small,maxDets=100,mAR:{all_eval_result[9]:.3f}, IoU=0.5:0.95,area=medium,maxDets=100,mAR:{all_eval_result[10]:.3f}, IoU=0.5:0.95,area=large,maxDets=100,mAR:{all_eval_result[11]:.3f}"
)
return
def test_model(args):
print(args)
if args.use_gpu:
# use one Graphics card to test
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
if not torch.cuda.is_available():
raise Exception("need gpu to test network!")
torch.cuda.empty_cache()
if args.seed is not None:
random.seed(args.seed)
if args.use_gpu:
torch.cuda.manual_seed_all(args.seed)
cudnn.deterministic = True
if args.use_gpu:
cudnn.benchmark = True
cudnn.enabled = True
voc_val_dataset = VocDetection(root_dir=VOCdataset_path,
image_sets=[('2007', 'test')],
transform=transforms.Compose([
Normalize(),
Resize(resize=args.input_image_size),
]),
keep_difficult=False)
if args.detector == "retinanet":
model = _retinanet(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = RetinaDecoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "fcos":
model = _fcos(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = FCOSDecoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "centernet":
model = _centernet(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = CenterNetDecoder(
image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
elif args.detector == "yolov3":
model = _yolov3(args.backbone, args.use_pretrained_model,
args.pretrained_model_path, args.num_classes)
decoder = YOLOV3Decoder(image_w=args.input_image_size,
image_h=args.input_image_size,
min_score_threshold=args.min_score_threshold)
else:
print("unsupport detection model!")
return
flops_input = torch.randn(1, 3, args.input_image_size,
args.input_image_size)
flops, params = profile(model, inputs=(flops_input, ))
flops, params = clever_format([flops, params], "%.3f")
print(
f"backbone:{args.backbone},detector: '{args.detector}', flops: {flops}, params: {params}"
)
if args.use_gpu:
model = model.cuda()
decoder = decoder.cuda()
model = nn.DataParallel(model)
print(f"start eval.")
all_ap, mAP = validate(voc_val_dataset, model, decoder, args)
print(f"eval done.")
for class_index, class_AP in all_ap.items():
print(f"class: {class_index},AP: {class_AP:.3f}")
print(f"mAP: {mAP:.3f}")
return
