You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

577 line
27KB

  1. """Utilities and tools for tracking runs with Weights & Biases."""
  2. import logging
  3. import os
  4. import sys
  5. from contextlib import contextmanager
  6. from pathlib import Path
  7. from typing import Dict
  8. import yaml
  9. from tqdm import tqdm
  10. FILE = Path(__file__).resolve()
  11. ROOT = FILE.parents[3] # YOLOv5 root directory
  12. if str(ROOT) not in sys.path:
  13. sys.path.append(str(ROOT)) # add ROOT to PATH
  14. from utils.datasets import LoadImagesAndLabels, img2label_paths
  15. from utils.general import LOGGER, check_dataset, check_file
  16. try:
  17. import wandb
  18. assert hasattr(wandb, '__version__') # verify package import not local dir
  19. except (ImportError, AssertionError):
  20. wandb = None
  21. RANK = int(os.getenv('RANK', -1))
  22. WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
  23. def remove_prefix(from_string, prefix=WANDB_ARTIFACT_PREFIX):
  24. return from_string[len(prefix):]
  25. def check_wandb_config_file(data_config_file):
  26. wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path
  27. if Path(wandb_config).is_file():
  28. return wandb_config
  29. return data_config_file
  30. def check_wandb_dataset(data_file):
  31. is_trainset_wandb_artifact = False
  32. is_valset_wandb_artifact = False
  33. if check_file(data_file) and data_file.endswith('.yaml'):
  34. with open(data_file, errors='ignore') as f:
  35. data_dict = yaml.safe_load(f)
  36. is_trainset_wandb_artifact = isinstance(data_dict['train'],
  37. str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX)
  38. is_valset_wandb_artifact = isinstance(data_dict['val'],
  39. str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX)
  40. if is_trainset_wandb_artifact or is_valset_wandb_artifact:
  41. return data_dict
  42. else:
  43. return check_dataset(data_file)
  44. def get_run_info(run_path):
  45. run_path = Path(remove_prefix(run_path, WANDB_ARTIFACT_PREFIX))
  46. run_id = run_path.stem
  47. project = run_path.parent.stem
  48. entity = run_path.parent.parent.stem
  49. model_artifact_name = 'run_' + run_id + '_model'
  50. return entity, project, run_id, model_artifact_name
  51. def check_wandb_resume(opt):
  52. process_wandb_config_ddp_mode(opt) if RANK not in [-1, 0] else None
  53. if isinstance(opt.resume, str):
  54. if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
  55. if RANK not in [-1, 0]: # For resuming DDP runs
  56. entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
  57. api = wandb.Api()
  58. artifact = api.artifact(entity + '/' + project + '/' + model_artifact_name + ':latest')
  59. modeldir = artifact.download()
  60. opt.weights = str(Path(modeldir) / "last.pt")
  61. return True
  62. return None
  63. def process_wandb_config_ddp_mode(opt):
  64. with open(check_file(opt.data), errors='ignore') as f:
  65. data_dict = yaml.safe_load(f) # data dict
  66. train_dir, val_dir = None, None
  67. if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX):
  68. api = wandb.Api()
  69. train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias)
  70. train_dir = train_artifact.download()
  71. train_path = Path(train_dir) / 'data/images/'
  72. data_dict['train'] = str(train_path)
  73. if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX):
  74. api = wandb.Api()
  75. val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias)
  76. val_dir = val_artifact.download()
  77. val_path = Path(val_dir) / 'data/images/'
  78. data_dict['val'] = str(val_path)
  79. if train_dir or val_dir:
  80. ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml')
  81. with open(ddp_data_path, 'w') as f:
  82. yaml.safe_dump(data_dict, f)
  83. opt.data = ddp_data_path
  84. class WandbLogger():
  85. """Log training runs, datasets, models, and predictions to Weights & Biases.
  86. This logger sends information to W&B at wandb.ai. By default, this information
  87. includes hyperparameters, system configuration and metrics, model metrics,
  88. and basic data metrics and analyses.
  89. By providing additional command line arguments to train.py, datasets,
  90. models and predictions can also be logged.
  91. For more on how this logger is used, see the Weights & Biases documentation:
  92. https://docs.wandb.com/guides/integrations/yolov5
  93. """
  94. def __init__(self, opt, run_id=None, job_type='Training'):
  95. """
  96. - Initialize WandbLogger instance
  97. - Upload dataset if opt.upload_dataset is True
  98. - Setup trainig processes if job_type is 'Training'
  99. arguments:
  100. opt (namespace) -- Commandline arguments for this run
  101. run_id (str) -- Run ID of W&B run to be resumed
  102. job_type (str) -- To set the job_type for this run
  103. """
  104. # Pre-training routine --
  105. self.job_type = job_type
  106. self.wandb, self.wandb_run = wandb, None if not wandb else wandb.run
  107. self.val_artifact, self.train_artifact = None, None
  108. self.train_artifact_path, self.val_artifact_path = None, None
  109. self.result_artifact = None
  110. self.val_table, self.result_table = None, None
  111. self.bbox_media_panel_images = []
  112. self.val_table_path_map = None
  113. self.max_imgs_to_log = 16
  114. self.wandb_artifact_data_dict = None
  115. self.data_dict = None
  116. # It's more elegant to stick to 1 wandb.init call,
  117. # but useful config data is overwritten in the WandbLogger's wandb.init call
  118. if isinstance(opt.resume, str): # checks resume from artifact
  119. if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
  120. entity, project, run_id, model_artifact_name = get_run_info(opt.resume)
  121. model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name
  122. assert wandb, 'install wandb to resume wandb runs'
  123. # Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config
  124. self.wandb_run = wandb.init(id=run_id,
  125. project=project,
  126. entity=entity,
  127. resume='allow',
  128. allow_val_change=True)
  129. opt.resume = model_artifact_name
  130. elif self.wandb:
  131. self.wandb_run = wandb.init(config=opt,
  132. resume="allow",
  133. project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem,
  134. entity=opt.entity,
  135. name=opt.name if opt.name != 'exp' else None,
  136. job_type=job_type,
  137. id=run_id,
  138. allow_val_change=True) if not wandb.run else wandb.run
  139. if self.wandb_run:
  140. if self.job_type == 'Training':
  141. if opt.upload_dataset:
  142. if not opt.resume:
  143. self.wandb_artifact_data_dict = self.check_and_upload_dataset(opt)
  144. if opt.resume:
  145. # resume from artifact
  146. if isinstance(opt.resume, str) and opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
  147. self.data_dict = dict(self.wandb_run.config.data_dict)
  148. else: # local resume
  149. self.data_dict = check_wandb_dataset(opt.data)
  150. else:
  151. self.data_dict = check_wandb_dataset(opt.data)
  152. self.wandb_artifact_data_dict = self.wandb_artifact_data_dict or self.data_dict
  153. # write data_dict to config. useful for resuming from artifacts. Do this only when not resuming.
  154. self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict}, allow_val_change=True)
  155. self.setup_training(opt)
  156. if self.job_type == 'Dataset Creation':
  157. self.wandb_run.config.update({"upload_dataset": True})
  158. self.data_dict = self.check_and_upload_dataset(opt)
  159. def check_and_upload_dataset(self, opt):
  160. """
  161. Check if the dataset format is compatible and upload it as W&B artifact
  162. arguments:
  163. opt (namespace)-- Commandline arguments for current run
  164. returns:
  165. Updated dataset info dictionary where local dataset paths are replaced by WAND_ARFACT_PREFIX links.
  166. """
  167. assert wandb, 'Install wandb to upload dataset'
  168. config_path = self.log_dataset_artifact(opt.data, opt.single_cls,
  169. 'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem)
  170. with open(config_path, errors='ignore') as f:
  171. wandb_data_dict = yaml.safe_load(f)
  172. return wandb_data_dict
  173. def setup_training(self, opt):
  174. """
  175. Setup the necessary processes for training YOLO models:
  176. - Attempt to download model checkpoint and dataset artifacts if opt.resume stats with WANDB_ARTIFACT_PREFIX
  177. - Update data_dict, to contain info of previous run if resumed and the paths of dataset artifact if downloaded
  178. - Setup log_dict, initialize bbox_interval
  179. arguments:
  180. opt (namespace) -- commandline arguments for this run
  181. """
  182. self.log_dict, self.current_epoch = {}, 0
  183. self.bbox_interval = opt.bbox_interval
  184. if isinstance(opt.resume, str):
  185. modeldir, _ = self.download_model_artifact(opt)
  186. if modeldir:
  187. self.weights = Path(modeldir) / "last.pt"
  188. config = self.wandb_run.config
  189. opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp, opt.imgsz = str(
  190. self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs,\
  191. config.hyp, config.imgsz
  192. data_dict = self.data_dict
  193. if self.val_artifact is None: # If --upload_dataset is set, use the existing artifact, don't download
  194. self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(
  195. data_dict.get('train'), opt.artifact_alias)
  196. self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(
  197. data_dict.get('val'), opt.artifact_alias)
  198. if self.train_artifact_path is not None:
  199. train_path = Path(self.train_artifact_path) / 'data/images/'
  200. data_dict['train'] = str(train_path)
  201. if self.val_artifact_path is not None:
  202. val_path = Path(self.val_artifact_path) / 'data/images/'
  203. data_dict['val'] = str(val_path)
  204. if self.val_artifact is not None:
  205. self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
  206. columns = ["epoch", "id", "ground truth", "prediction"]
  207. columns.extend(self.data_dict['names'])
  208. self.result_table = wandb.Table(columns)
  209. self.val_table = self.val_artifact.get("val")
  210. if self.val_table_path_map is None:
  211. self.map_val_table_path()
  212. if opt.bbox_interval == -1:
  213. self.bbox_interval = opt.bbox_interval = (opt.epochs // 10) if opt.epochs > 10 else 1
  214. if opt.evolve:
  215. self.bbox_interval = opt.bbox_interval = opt.epochs + 1
  216. train_from_artifact = self.train_artifact_path is not None and self.val_artifact_path is not None
  217. # Update the the data_dict to point to local artifacts dir
  218. if train_from_artifact:
  219. self.data_dict = data_dict
  220. def download_dataset_artifact(self, path, alias):
  221. """
  222. download the model checkpoint artifact if the path starts with WANDB_ARTIFACT_PREFIX
  223. arguments:
  224. path -- path of the dataset to be used for training
  225. alias (str)-- alias of the artifact to be download/used for training
  226. returns:
  227. (str, wandb.Artifact) -- path of the downladed dataset and it's corresponding artifact object if dataset
  228. is found otherwise returns (None, None)
  229. """
  230. if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX):
  231. artifact_path = Path(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias)
  232. dataset_artifact = wandb.use_artifact(artifact_path.as_posix().replace("\\", "/"))
  233. assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'"
  234. datadir = dataset_artifact.download()
  235. return datadir, dataset_artifact
  236. return None, None
  237. def download_model_artifact(self, opt):
  238. """
  239. download the model checkpoint artifact if the resume path starts with WANDB_ARTIFACT_PREFIX
  240. arguments:
  241. opt (namespace) -- Commandline arguments for this run
  242. """
  243. if opt.resume.startswith(WANDB_ARTIFACT_PREFIX):
  244. model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest")
  245. assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist'
  246. modeldir = model_artifact.download()
  247. # epochs_trained = model_artifact.metadata.get('epochs_trained')
  248. total_epochs = model_artifact.metadata.get('total_epochs')
  249. is_finished = total_epochs is None
  250. assert not is_finished, 'training is finished, can only resume incomplete runs.'
  251. return modeldir, model_artifact
  252. return None, None
  253. def log_model(self, path, opt, epoch, fitness_score, best_model=False):
  254. """
  255. Log the model checkpoint as W&B artifact
  256. arguments:
  257. path (Path) -- Path of directory containing the checkpoints
  258. opt (namespace) -- Command line arguments for this run
  259. epoch (int) -- Current epoch number
  260. fitness_score (float) -- fitness score for current epoch
  261. best_model (boolean) -- Boolean representing if the current checkpoint is the best yet.
  262. """
  263. model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model',
  264. type='model',
  265. metadata={
  266. 'original_url': str(path),
  267. 'epochs_trained': epoch + 1,
  268. 'save period': opt.save_period,
  269. 'project': opt.project,
  270. 'total_epochs': opt.epochs,
  271. 'fitness_score': fitness_score})
  272. model_artifact.add_file(str(path / 'last.pt'), name='last.pt')
  273. wandb.log_artifact(model_artifact,
  274. aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else ''])
  275. LOGGER.info(f"Saving model artifact on epoch {epoch + 1}")
  276. def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False):
  277. """
  278. Log the dataset as W&B artifact and return the new data file with W&B links
  279. arguments:
  280. data_file (str) -- the .yaml file with information about the dataset like - path, classes etc.
  281. single_class (boolean) -- train multi-class data as single-class
  282. project (str) -- project name. Used to construct the artifact path
  283. overwrite_config (boolean) -- overwrites the data.yaml file if set to true otherwise creates a new
  284. file with _wandb postfix. Eg -> data_wandb.yaml
  285. returns:
  286. the new .yaml file with artifact links. it can be used to start training directly from artifacts
  287. """
  288. upload_dataset = self.wandb_run.config.upload_dataset
  289. log_val_only = isinstance(upload_dataset, str) and upload_dataset == 'val'
  290. self.data_dict = check_dataset(data_file) # parse and check
  291. data = dict(self.data_dict)
  292. nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names'])
  293. names = {k: v for k, v in enumerate(names)} # to index dictionary
  294. # log train set
  295. if not log_val_only:
  296. self.train_artifact = self.create_dataset_table(LoadImagesAndLabels(data['train'], rect=True, batch_size=1),
  297. names,
  298. name='train') if data.get('train') else None
  299. if data.get('train'):
  300. data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train')
  301. self.val_artifact = self.create_dataset_table(
  302. LoadImagesAndLabels(data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None
  303. if data.get('val'):
  304. data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val')
  305. path = Path(data_file)
  306. # create a _wandb.yaml file with artifacts links if both train and test set are logged
  307. if not log_val_only:
  308. path = (path.stem if overwrite_config else path.stem + '_wandb') + '.yaml' # updated data.yaml path
  309. path = ROOT / 'data' / path
  310. data.pop('download', None)
  311. data.pop('path', None)
  312. with open(path, 'w') as f:
  313. yaml.safe_dump(data, f)
  314. LOGGER.info(f"Created dataset config file {path}")
  315. if self.job_type == 'Training': # builds correct artifact pipeline graph
  316. if not log_val_only:
  317. self.wandb_run.log_artifact(
  318. self.train_artifact) # calling use_artifact downloads the dataset. NOT NEEDED!
  319. self.wandb_run.use_artifact(self.val_artifact)
  320. self.val_artifact.wait()
  321. self.val_table = self.val_artifact.get('val')
  322. self.map_val_table_path()
  323. else:
  324. self.wandb_run.log_artifact(self.train_artifact)
  325. self.wandb_run.log_artifact(self.val_artifact)
  326. return path
  327. def map_val_table_path(self):
  328. """
  329. Map the validation dataset Table like name of file -> it's id in the W&B Table.
  330. Useful for - referencing artifacts for evaluation.
  331. """
  332. self.val_table_path_map = {}
  333. LOGGER.info("Mapping dataset")
  334. for i, data in enumerate(tqdm(self.val_table.data)):
  335. self.val_table_path_map[data[3]] = data[0]
  336. def create_dataset_table(self, dataset: LoadImagesAndLabels, class_to_id: Dict[int, str], name: str = 'dataset'):
  337. """
  338. Create and return W&B artifact containing W&B Table of the dataset.
  339. arguments:
  340. dataset -- instance of LoadImagesAndLabels class used to iterate over the data to build Table
  341. class_to_id -- hash map that maps class ids to labels
  342. name -- name of the artifact
  343. returns:
  344. dataset artifact to be logged or used
  345. """
  346. # TODO: Explore multiprocessing to slpit this loop parallely| This is essential for speeding up the the logging
  347. artifact = wandb.Artifact(name=name, type="dataset")
  348. img_files = tqdm([dataset.path]) if isinstance(dataset.path, str) and Path(dataset.path).is_dir() else None
  349. img_files = tqdm(dataset.im_files) if not img_files else img_files
  350. for img_file in img_files:
  351. if Path(img_file).is_dir():
  352. artifact.add_dir(img_file, name='data/images')
  353. labels_path = 'labels'.join(dataset.path.rsplit('images', 1))
  354. artifact.add_dir(labels_path, name='data/labels')
  355. else:
  356. artifact.add_file(img_file, name='data/images/' + Path(img_file).name)
  357. label_file = Path(img2label_paths([img_file])[0])
  358. artifact.add_file(str(label_file), name='data/labels/' +
  359. label_file.name) if label_file.exists() else None
  360. table = wandb.Table(columns=["id", "train_image", "Classes", "name"])
  361. class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()])
  362. for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)):
  363. box_data, img_classes = [], {}
  364. for cls, *xywh in labels[:, 1:].tolist():
  365. cls = int(cls)
  366. box_data.append({
  367. "position": {
  368. "middle": [xywh[0], xywh[1]],
  369. "width": xywh[2],
  370. "height": xywh[3]},
  371. "class_id": cls,
  372. "box_caption": "%s" % (class_to_id[cls])})
  373. img_classes[cls] = class_to_id[cls]
  374. boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
  375. table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()),
  376. Path(paths).name)
  377. artifact.add(table, name)
  378. return artifact
  379. def log_training_progress(self, predn, path, names):
  380. """
  381. Build evaluation Table. Uses reference from validation dataset table.
  382. arguments:
  383. predn (list): list of predictions in the native space in the format - [xmin, ymin, xmax, ymax, confidence, class]
  384. path (str): local path of the current evaluation image
  385. names (dict(int, str)): hash map that maps class ids to labels
  386. """
  387. class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()])
  388. box_data = []
  389. avg_conf_per_class = [0] * len(self.data_dict['names'])
  390. pred_class_count = {}
  391. for *xyxy, conf, cls in predn.tolist():
  392. if conf >= 0.25:
  393. cls = int(cls)
  394. box_data.append({
  395. "position": {
  396. "minX": xyxy[0],
  397. "minY": xyxy[1],
  398. "maxX": xyxy[2],
  399. "maxY": xyxy[3]},
  400. "class_id": cls,
  401. "box_caption": f"{names[cls]} {conf:.3f}",
  402. "scores": {
  403. "class_score": conf},
  404. "domain": "pixel"})
  405. avg_conf_per_class[cls] += conf
  406. if cls in pred_class_count:
  407. pred_class_count[cls] += 1
  408. else:
  409. pred_class_count[cls] = 1
  410. for pred_class in pred_class_count.keys():
  411. avg_conf_per_class[pred_class] = avg_conf_per_class[pred_class] / pred_class_count[pred_class]
  412. boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
  413. id = self.val_table_path_map[Path(path).name]
  414. self.result_table.add_data(self.current_epoch, id, self.val_table.data[id][1],
  415. wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set),
  416. *avg_conf_per_class)
  417. def val_one_image(self, pred, predn, path, names, im):
  418. """
  419. Log validation data for one image. updates the result Table if validation dataset is uploaded and log bbox media panel
  420. arguments:
  421. pred (list): list of scaled predictions in the format - [xmin, ymin, xmax, ymax, confidence, class]
  422. predn (list): list of predictions in the native space - [xmin, ymin, xmax, ymax, confidence, class]
  423. path (str): local path of the current evaluation image
  424. """
  425. if self.val_table and self.result_table: # Log Table if Val dataset is uploaded as artifact
  426. self.log_training_progress(predn, path, names)
  427. if len(self.bbox_media_panel_images) < self.max_imgs_to_log and self.current_epoch > 0:
  428. if self.current_epoch % self.bbox_interval == 0:
  429. box_data = [{
  430. "position": {
  431. "minX": xyxy[0],
  432. "minY": xyxy[1],
  433. "maxX": xyxy[2],
  434. "maxY": xyxy[3]},
  435. "class_id": int(cls),
  436. "box_caption": f"{names[int(cls)]} {conf:.3f}",
  437. "scores": {
  438. "class_score": conf},
  439. "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
  440. boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
  441. self.bbox_media_panel_images.append(wandb.Image(im, boxes=boxes, caption=path.name))
  442. def log(self, log_dict):
  443. """
  444. save the metrics to the logging dictionary
  445. arguments:
  446. log_dict (Dict) -- metrics/media to be logged in current step
  447. """
  448. if self.wandb_run:
  449. for key, value in log_dict.items():
  450. self.log_dict[key] = value
  451. def end_epoch(self, best_result=False):
  452. """
  453. commit the log_dict, model artifacts and Tables to W&B and flush the log_dict.
  454. arguments:
  455. best_result (boolean): Boolean representing if the result of this evaluation is best or not
  456. """
  457. if self.wandb_run:
  458. with all_logging_disabled():
  459. if self.bbox_media_panel_images:
  460. self.log_dict["BoundingBoxDebugger"] = self.bbox_media_panel_images
  461. try:
  462. wandb.log(self.log_dict)
  463. except BaseException as e:
  464. LOGGER.info(
  465. f"An error occurred in wandb logger. The training will proceed without interruption. More info\n{e}"
  466. )
  467. self.wandb_run.finish()
  468. self.wandb_run = None
  469. self.log_dict = {}
  470. self.bbox_media_panel_images = []
  471. if self.result_artifact:
  472. self.result_artifact.add(self.result_table, 'result')
  473. wandb.log_artifact(self.result_artifact,
  474. aliases=[
  475. 'latest', 'last', 'epoch ' + str(self.current_epoch),
  476. ('best' if best_result else '')])
  477. wandb.log({"evaluation": self.result_table})
  478. columns = ["epoch", "id", "ground truth", "prediction"]
  479. columns.extend(self.data_dict['names'])
  480. self.result_table = wandb.Table(columns)
  481. self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation")
  482. def finish_run(self):
  483. """
  484. Log metrics if any and finish the current W&B run
  485. """
  486. if self.wandb_run:
  487. if self.log_dict:
  488. with all_logging_disabled():
  489. wandb.log(self.log_dict)
  490. wandb.run.finish()
  491. @contextmanager
  492. def all_logging_disabled(highest_level=logging.CRITICAL):
  493. """ source - https://gist.github.com/simon-weber/7853144
  494. A context manager that will prevent any logging messages triggered during the body from being processed.
  495. :param highest_level: the maximum logging level in use.
  496. This would only need to be changed if a custom level greater than CRITICAL is defined.
  497. """
  498. previous_level = logging.root.manager.disable
  499. logging.disable(highest_level)
  500. try:
  501. yield
  502. finally:
  503. logging.disable(previous_level)