zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Suppress wandb images size mismatch warning (#3611) 

* supress wandb images size mismatch warning

* supress wandb images size mismatch warning

* PEP8 reformat and optimize imports

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia 2021-06-14 22:24:58 +05:30 committed by GitHub
parent 239a11c197
commit daab682b06
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 25 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
utils/wandb_logging/wandb_utils.py
View File

@ -1,16 +1,16 @@
"""Utilities and tools for tracking runs with Weights & Biases.""" """Utilities and tools for tracking runs with Weights & Biases."""
import json import logging
import sys import sys
from contextlib import contextmanager
from pathlib import Path from pathlib import Path
import torch
import yaml import yaml
from tqdm import tqdm from tqdm import tqdm
sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path sys.path.append(str(Path(__file__).parent.parent.parent)) # add utils/ to path
from utils.datasets import LoadImagesAndLabels from utils.datasets import LoadImagesAndLabels
from utils.datasets import img2label_paths from utils.datasets import img2label_paths
from utils.general import colorstr, xywh2xyxy, check_dataset, check_file from utils.general import colorstr, check_dataset, check_file
try: try:
import wandb import wandb
@ -92,6 +92,7 @@ class WandbLogger():
For more on how this logger is used, see the Weights & Biases documentation: For more on how this logger is used, see the Weights & Biases documentation:
https://docs.wandb.com/guides/integrations/yolov5 https://docs.wandb.com/guides/integrations/yolov5
""" """
def __init__(self, opt, name, run_id, data_dict, job_type='Training'): def __init__(self, opt, name, run_id, data_dict, job_type='Training'):
# Pre-training routine -- # Pre-training routine --
self.job_type = job_type self.job_type = job_type
@ -272,7 +273,7 @@ class WandbLogger():
"box_caption": "%s" % (class_to_id[cls])}) "box_caption": "%s" % (class_to_id[cls])})
img_classes[cls] = class_to_id[cls] img_classes[cls] = class_to_id[cls]
boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space
table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), json.dumps(img_classes), table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()),
Path(paths).name) Path(paths).name)
artifact.add(table, name) artifact.add(table, name)
return artifact return artifact
@ -306,6 +307,7 @@ class WandbLogger():
def end_epoch(self, best_result=False): def end_epoch(self, best_result=False):
if self.wandb_run: if self.wandb_run:
with all_logging_disabled():
wandb.log(self.log_dict) wandb.log(self.log_dict)
self.log_dict = {} self.log_dict = {}
if self.result_artifact: if self.result_artifact:
@ -319,5 +321,21 @@ class WandbLogger():
def finish_run(self): def finish_run(self):
if self.wandb_run: if self.wandb_run:
if self.log_dict: if self.log_dict:
with all_logging_disabled():
wandb.log(self.log_dict) wandb.log(self.log_dict)
wandb.run.finish() wandb.run.finish()
@contextmanager
def all_logging_disabled(highest_level=logging.CRITICAL):
""" source - https://gist.github.com/simon-weber/7853144
A context manager that will prevent any logging messages triggered during the body from being processed.
:param highest_level: the maximum logging level in use.
This would only need to be changed if a custom level greater than CRITICAL is defined.
"""
previous_level = logging.root.manager.disable
logging.disable(highest_level)
try:
yield
finally:
logging.disable(previous_level)