File size: 1,546 Bytes
d593c77 |
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 |
import os, shutil
import torch, math
def colorstr(*input):
# Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
*args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string
colors = {'black': '\033[30m', # basic colors
'red': '\033[31m',
'green': '\033[32m',
'yellow': '\033[33m',
'blue': '\033[34m',
'magenta': '\033[35m',
'cyan': '\033[36m',
'white': '\033[37m',
'bright_black': '\033[90m', # bright colors
'bright_red': '\033[91m',
'bright_green': '\033[92m',
'bright_yellow': '\033[93m',
'bright_blue': '\033[94m',
'bright_magenta': '\033[95m',
'bright_cyan': '\033[96m',
'bright_white': '\033[97m',
'end': '\033[0m', # misc
'bold': '\033[1m',
'underline': '\033[4m'}
return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
def SaveCheckpoint(state, last, last_path, best, best_path, is_best):
if os.path.exists(last):
shutil.rmtree(last)
last_path.mkdir(parents=True, exist_ok=True)
torch.save(state, os.path.join(last_path, 'ckpt.pth'))
if is_best:
if os.path.exists(best):
shutil.rmtree(best)
best_path.mkdir(parents=True, exist_ok=True)
torch.save(state, os.path.join(best_path, 'ckpt.pth'))
|