Add colorstr() (#1887)
* Add colorful() * update * newline fix * add git description * --always * update loss scaling * update loss scaling 2 * rename to colorstr()
This commit is contained in:
1 parent
3e25f1e9e5
commit
6ab589583c
5 files changed
+60
-22
No files matched your search
+13
-3
@@ -3,9 +3,11 @@
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
@@ -41,9 +43,17 @@ def init_torch_seeds(seed=0):
|
||||
cudnn.benchmark, cudnn.deterministic = True, False
|
||||
|
||||
|
||||
def git_describe():
|
||||
# return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
|
||||
if Path('.git').exists():
|
||||
return subprocess.check_output('git describe --tags --long --always', shell=True).decode('utf-8')[:-1]
|
||||
else:
|
||||
return ''
|
||||
|
||||
|
||||
def select_device(device='', batch_size=None):
|
||||
# device = 'cpu' or '0' or '0,1,2,3'
|
||||
s = f'Using torch {torch.__version__} ' # string
|
||||
s = f'YOLOv5 {git_describe()} torch {torch.__version__} ' # string
|
||||
cpu = device.lower() == 'cpu'
|
||||
if cpu:
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
|
||||
@@ -61,9 +71,9 @@ def select_device(device='', batch_size=None):
|
||||
p = torch.cuda.get_device_properties(i)
|
||||
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB
|
||||
else:
|
||||
s += 'CPU'
|
||||
s += 'CPU\n'
|
||||
|
||||
logger.info(f'{s}\n') # skip a line
|
||||
logger.info(s) # skip a line
|
||||
return torch.device('cuda:0' if cuda else 'cpu')
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user