forked from NVlabs/prismer
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
40 lines (31 loc) · 1.37 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://github.com/NVlabs/prismer/blob/main/LICENSE
import math
from pycocotools.coco import COCO
from pycocoevalcap.eval import COCOEvalCap
def cosine_lr_schedule(optimizer, epoch, max_epoch, init_lr, min_lr):
"""Decay the learning rate"""
lr = (init_lr - min_lr) * 0.5 * (1. + math.cos(math.pi * epoch / max_epoch)) + min_lr
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def warmup_lr_schedule(optimizer, step, max_step, init_lr, max_lr):
"""Warmup the learning rate"""
lr = min(max_lr, init_lr + (max_lr - init_lr) * step / max_step)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def step_lr_schedule(optimizer, epoch, init_lr, min_lr, decay_rate):
"""Decay the learning rate"""
lr = max(min_lr, init_lr * (decay_rate ** epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def coco_caption_eval(coco_gt_root, results_file):
coco = COCO(coco_gt_root)
coco_result = coco.loadRes(results_file)
coco_eval = COCOEvalCap(coco, coco_result)
coco_eval.evaluate()
for metric, score in coco_eval.eval.items():
print(f'{metric}: {score:.3f}')
return coco_eval