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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
|
- # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
- """
- Logging utils
- """
- import warnings
- from threading import Thread
- import torch
- from torch.utils.tensorboard import SummaryWriter
- from utils.general import colorstr, emojis
- from utils.loggers.wandb.wandb_utils import WandbLogger
- from utils.plots import plot_images, plot_results
- from utils.torch_utils import de_parallel
- LOGGERS = ('csv', 'tb', 'wandb') # text-file, TensorBoard, Weights & Biases
- try:
- import wandb
- assert hasattr(wandb, '__version__') # verify package import not local dir
- except (ImportError, AssertionError):
- wandb = None
- class Loggers():
- # YOLOv5 Loggers class
- def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS):
- self.save_dir = save_dir
- self.weights = weights
- self.opt = opt
- self.hyp = hyp
- self.logger = logger # for printing results to console
- self.include = include
- self.keys = ['train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss
- 'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', # metrics
- 'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss
- 'x/lr0', 'x/lr1', 'x/lr2'] # params
- for k in LOGGERS:
- setattr(self, k, None) # init empty logger dictionary
- self.csv = True # always log to csv
- # Message
- if not wandb:
- prefix = colorstr('Weights & Biases: ')
- s = f"{prefix}run 'pip install wandb' to automatically track and visualize YOLOv5 🚀 runs (RECOMMENDED)"
- print(emojis(s))
- # TensorBoard
- s = self.save_dir
- if 'tb' in self.include and not self.opt.evolve:
- prefix = colorstr('TensorBoard: ')
- self.logger.info(f"{prefix}Start with 'tensorboard --logdir {s.parent}', view at http://localhost:6006/")
- self.tb = SummaryWriter(str(s))
- # W&B
- if wandb and 'wandb' in self.include:
- wandb_artifact_resume = isinstance(self.opt.resume, str) and self.opt.resume.startswith('wandb-artifact://')
- run_id = torch.load(self.weights).get('wandb_id') if self.opt.resume and not wandb_artifact_resume else None
- self.opt.hyp = self.hyp # add hyperparameters
- self.wandb = WandbLogger(self.opt, run_id)
- else:
- self.wandb = None
- def on_pretrain_routine_end(self):
- # Callback runs on pre-train routine end
- paths = self.save_dir.glob('*labels*.jpg') # training labels
- if self.wandb:
- self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]})
- def on_train_batch_end(self, ni, model, imgs, targets, paths, plots, sync_bn):
- # Callback runs on train batch end
- if plots:
- if ni == 0:
- if not sync_bn: # tb.add_graph() --sync known issue https://github.com/ultralytics/yolov5/issues/3754
- with warnings.catch_warnings():
- warnings.simplefilter('ignore') # suppress jit trace warning
- self.tb.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), [])
- if ni < 3:
- f = self.save_dir / f'train_batch{ni}.jpg' # filename
- Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
- if self.wandb and ni == 10:
- files = sorted(self.save_dir.glob('train*.jpg'))
- self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]})
- def on_train_epoch_end(self, epoch):
- # Callback runs on train epoch end
- if self.wandb:
- self.wandb.current_epoch = epoch + 1
- def on_val_image_end(self, pred, predn, path, names, im):
- # Callback runs on val image end
- if self.wandb:
- self.wandb.val_one_image(pred, predn, path, names, im)
- def on_val_end(self):
- # Callback runs on val end
- if self.wandb:
- files = sorted(self.save_dir.glob('val*.jpg'))
- self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]})
- def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
- # Callback runs at the end of each fit (train+val) epoch
- x = {k: v for k, v in zip(self.keys, vals)} # dict
- if self.csv:
- file = self.save_dir / 'results.csv'
- n = len(x) + 1 # number of cols
- s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n') # add header
- with open(file, 'a') as f:
- f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n')
- if self.tb:
- for k, v in x.items():
- self.tb.add_scalar(k, v, epoch)
- if self.wandb:
- self.wandb.log(x)
- self.wandb.end_epoch(best_result=best_fitness == fi)
- def on_model_save(self, last, epoch, final_epoch, best_fitness, fi):
- # Callback runs on model save event
- if self.wandb:
- if ((epoch + 1) % self.opt.save_period == 0 and not final_epoch) and self.opt.save_period != -1:
- self.wandb.log_model(last.parent, self.opt, epoch, fi, best_model=best_fitness == fi)
- def on_train_end(self, last, best, plots, epoch):
- # Callback runs on training end
- if plots:
- plot_results(file=self.save_dir / 'results.csv') # save results.png
- files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
- files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter
- if self.tb:
- import cv2
- for f in files:
- self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC')
- if self.wandb:
- self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
- # Calling wandb.log. TODO: Refactor this into WandbLogger.log_model
- if not self.opt.evolve:
- wandb.log_artifact(str(best if best.exists() else last), type='model',
- name='run_' + self.wandb.wandb_run.id + '_model',
- aliases=['latest', 'best', 'stripped'])
- self.wandb.finish_run()
- else:
- self.wandb.finish_run()
- self.wandb = WandbLogger(self.opt)
|