-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
70 lines (53 loc) · 1.91 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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import shutil
import time
from pathlib import Path
import torch
def get_git_state():
import git
repo = git.Repo(search_parent_directories=True)
try:
branch_name = repo.active_branch.name
except Exception:
branch_name = "NONE"
return branch_name, repo.head.object.hexsha
def time_print(*args, **kwargs):
print(time.strftime("%Y-%m-%d %H:%M:%S"), *args, **kwargs)
def save_checkpoint(state, is_best=True, directory=".", filename='checkpoint.pth.tar'):
directory = Path(directory)
directory.mkdir(exist_ok=True, parents=True)
filepath = str(directory / filename)
torch.save(state, filepath)
if is_best:
shutil.copyfile(filepath, str(directory / 'model_best.pth.tar'))
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self, name, fmt=':f'):
self.name = name
self.fmt = fmt
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def __str__(self):
fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
return fmtstr.format(**self.__dict__)
class ProgressMeter(object):
def __init__(self, num_batches, meters, prefix=""):
self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
self.meters = meters
self.prefix = prefix
def display(self, batch, print_fn=print):
entries = [self.prefix + self.batch_fmtstr.format(batch)]
entries += [str(meter) for meter in self.meters]
print_fn('\t'.join(entries))
def _get_batch_fmtstr(self, num_batches):
num_digits = len(str(num_batches // 1))
fmt = '{:' + str(num_digits) + 'd}'
return '[' + fmt + '/' + fmt.format(num_batches) + ']'