tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Consolidate StatsTracker as LocalWriter; cleaned up functions to write (#88)

Browse Source 
* removed stats tracker; consolidate writers into local writer

* basic revamp of stats tracker/writers
This commit is contained in:
evonneng 2022-05-20 08:51:45 -07:00 committed by GitHub
parent 2287eb3455
commit 10d713f84b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 256 additions and 311 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
configs/logging/default_logging.yaml
View File

@ -1,12 +1,12 @@
steps_per_log: 10
# stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR)
enable_stats: True
stats_tracker:
max_history: 10
stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] # see mattport/utils/stats_tracker.py for options
max_history: 20
# writer logs losses and images per iteration
writer:
TensorboardWriter:
save_dir: './'
LocalWriter:
save_dir: './'
stats_to_track: ["Data Load (time)", "Train Iter (time)", "Rays Per Sec (time)", "Test PSNR"]
# profiler logs run times of functions and prints at end of training
enable_profiler: True

2
mattport/nerf/graph/nerfw.py
View File

@ -212,4 +212,4 @@ class NerfWGraph(NeRFGraph):
row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2)
combined_image = torch.cat([row0, row1, row2], dim=-3)
writer.write_image(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img")
writer.put_scalar(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img")

2
mattport/nerf/graph/semantic_nerf.py
View File

@ -155,4 +155,4 @@ class SemanticNerfGraph(NeRFGraph):
self.stuff_classes
)
semantic_labels_image = visualization.apply_colormap(semantic_labels[..., None])
writer.write_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics")
writer.put_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics")

20
mattport/nerf/graph/vanilla_nerf.py
View File

@ -23,7 +23,7 @@ from mattport.nerf.sampler import PDFSampler, UniformSampler
from mattport.structures import colors
from mattport.structures.rays import PointSamples, RayBundle
from mattport.nerf.colliders import NearFarCollider
from mattport.utils import stats_tracker, visualization, writer
from mattport.utils import visualization, writer
from mattport.utils.misc import is_not_none
@ -232,9 +232,9 @@ class NeRFGraph(Graph):
combined_acc = torch.cat([acc_coarse, acc_fine], dim=1)
combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)
writer.write_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img")
writer.write_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation")
writer.write_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth")
writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img")
writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation")
writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth")
# Switch images from [H, W, C] to [1, C, H, W] for metrics computations
image = torch.moveaxis(image, -1, 0)[None, ...]
@ -246,13 +246,11 @@ class NeRFGraph(Graph):
fine_ssim = self.ssim(image, rgb_fine)
fine_lpips = self.lpips(image, rgb_fine)
writer.write_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr")
writer.write_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr")
writer.write_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim")
writer.write_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips")
writer.put_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr")
writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr")
writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim")
writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips")
stats_tracker.update_stats(
{"name": stats_tracker.Stats.CURR_TEST_PSNR, "value": float(fine_psnr), "step": step}
)
writer.put_scalar(name="Test PSNR", scalar=float(fine_psnr), step=step)
return fine_psnr.item()

2
mattport/nerf/optimizers.py
View File

@ -83,4 +83,4 @@ class Optimizers:
scheduler.step()
# TODO(ethan): clean this up. why is there indexing into a list?
lr = scheduler.get_last_lr()[0]
writer.write_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate")
writer.put_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate")

60
mattport/nerf/trainer.py
View File

@ -18,7 +18,7 @@ from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size
from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs
from mattport.nerf.optimizers import Optimizers
from mattport.utils import profiler, stats_tracker, writer
from mattport.utils import profiler, writer
from mattport.utils.decorators import check_main_thread
logging.getLogger("PIL").setLevel(logging.WARNING)
@ -51,8 +51,7 @@ class Trainer:
self.optimizers = None
self.start_step = 0
# logging variables
stats_tracker.setup_stats_tracker(config)
writer.setup_event_writers(config.logging.writer)
writer.setup_event_writers(config)
profiler.setup_profiler(config.logging)
self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}"
@ -199,47 +198,54 @@ class Trainer:
train_start = time()
num_iterations = self.config.graph.max_num_iterations
iter_dataset = iter(self.train_dataloader)
for i, step in enumerate(range(self.start_step, self.start_step + num_iterations)):
for step in range(self.start_step, self.start_step + num_iterations):
data_start = time()
batch = next(iter_dataset)
stats_tracker.update_stats(
{"name": stats_tracker.Stats.ITER_LOAD_TIME, "start_time": data_start, "end_time": time(), "step": step}
writer.put_time(
name="Data Load (time)", start_time=data_start, end_time=time(), step=step, avg_over_iters=True
)
iter_start = time()
loss_dict = self.train_iteration(batch, step)
stats_tracker.update_stats(
{
"name": stats_tracker.Stats.RAYS_PER_SEC,
"start_time": iter_start,
"end_time": time(),
"step": step,
"batch_size": batch["indices"].shape[0],
},
writer.put_time(
name="Rays Per Sec (time)",
start_time=iter_start,
end_time=time(),
step=step,
avg_over_iters=True,
avg_over_batch=batch["indices"].shape[0],
)
stats_tracker.update_stats(
{
"name": stats_tracker.Stats.ITER_TRAIN_TIME,
"start_time": iter_start,
"end_time": time(),
"step": step,
}
writer.put_time(
name="Train Iter (time)",
start_time=iter_start,
end_time=time(),
step=step,
avg_over_iters=True,
update_eta=True,
)
if step != 0 and step % self.config.logging.steps_per_log == 0:
writer.write_scalar_dict(scalar_dict=loss_dict, step=step, group="Loss", prefix="train-")
writer.put_dict(name="loss_dict", scalar_dict=loss_dict, step=step, group="Loss", prefix="train-")
# TODO: add the learning rates to tensorboard/logging
if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0:
self.save_checkpoint(self.config.graph.model_dir, step)
if step % self.config.graph.steps_per_test == 0: # NOTE(ethan): we should still run this in dry-run mode!
for image_idx in self.config.data.val_image_indices:
_ = self.test_image(image_idx=image_idx, step=step)
stats_tracker.print_stats(i / num_iterations)
self._write_out_storage(step)
stats_tracker.update_stats(
{"name": stats_tracker.Stats.TOTAL_TRAIN_TIME, "start_time": train_start, "end_time": time()}
)
stats_tracker.print_stats(1.0)
writer.put_time(name="Train Total (time)", start_time=train_start, end_time=time(), step=num_iterations)
self._write_out_storage(num_iterations)
def _write_out_storage(self, step):
"""Perform writes only during appropriate time steps"""
if (
step % self.config.logging.steps_per_log == 0
or (self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0)
or step % self.config.graph.steps_per_test == 0
or step == self.config.graph.max_num_iterations
):
writer.write_out_storage()
@profiler.time_function
def train_iteration(self, batch: dict, step: int):

34
mattport/utils/decorators.py
View File

@ -19,24 +19,12 @@ def decorate_all(decorators: List[Callable]) -> Callable:
return decorate
def check_stats_enabled(func: Callable) -> Callable:
"""Decorator: check if stats tracker is enabled"""
def wrapper(self, *args, **kwargs):
ret = None
if self.config.logging.enable_stats:
ret = func(self, *args, **kwargs)
return ret
return wrapper
def check_profiler_enabled(func: Callable) -> Callable:
"""Decorator: check if profiler is enabled"""
def wrapper(self, *args, **kwargs):
ret = None
if self.config.enable_stats:
if self.config.enable_profiler:
ret = func(self, *args, **kwargs)
return ret
@ -46,26 +34,10 @@ def check_profiler_enabled(func: Callable) -> Callable:
def check_main_thread(func: Callable) -> Callable:
"""Decorator: check if you are on main thread"""
def wrapper(self, *args, **kwargs):
def wrapper(*args, **kwargs):
ret = None
if comms.is_main_process():
ret = func(self, *args, **kwargs)
return ret
return wrapper
def check_print_stats_step(func: Callable) -> Callable:
"""Decorator: check if it is time to print stats update"""
def wrapper(self, *args, **kwargs):
ret = None
if (
self.step % self.config.logging.steps_per_log == 0
or (self.config.graph.steps_per_save and self.step % self.config.graph.steps_per_save == 0)
or self.step % self.config.graph.steps_per_test == 0
):
ret = func(self, *args, **kwargs)
ret = func(*args, **kwargs)
return ret
return wrapper

192
mattport/utils/stats_tracker.py
View File

@ -1,192 +0,0 @@
"""
Stats tracker base class and functionality
"""
import datetime
import enum
import logging
from typing import Any, Dict
from omegaconf import DictConfig
from mattport.utils import comms
from mattport.utils.decorators import check_main_thread, check_print_stats_step, check_stats_enabled, decorate_all
STATS_TRACKER = []
def update_stats(args: Dict[str, Any]):
"""update global stats tracker based on arguments"""
assert comms.is_main_process(), "Writing out with process other than main"
if "value" in args:
STATS_TRACKER[0].update_value(**args)
elif "start_time" in args:
STATS_TRACKER[0].update_time(**args)
else:
raise NotImplementedError
def print_stats(fraction_done: float):
"""print current statistics information"""
assert comms.is_main_process(), "Writing out with process other than main"
STATS_TRACKER[0].print_stats(fraction_done)
def setup_stats_tracker(config: DictConfig):
"""Initialization of stats tracker"""
if comms.is_main_process():
STATS_TRACKER.append(StatsTracker(config))
class Stats(enum.Enum):
"""Possible Stats values for StatsTracker
Args:
enum (_type_): _description_
"""
ITER_LOAD_TIME = "Data Load (ms)"
ITER_TRAIN_TIME = "Train Iter (ms)"
TOTAL_TRAIN_TIME = "Train Total (time)"
RAYS_PER_SEC = "Rays Per Sec (1/s)"
CURR_TEST_PSNR = "Test PSNR"
ETA = "ETA (time)"
@decorate_all([check_stats_enabled, check_main_thread])
class StatsTracker:
"""Stats Tracker class"""
def __init__(self, config: DictConfig):
self.config = config
if self.config.logging.enable_stats:
self.max_history = self.config.logging.stats_tracker.max_history
self.step = 0
self.stats_dict = {}
self.past_stats = []
self.new_key = False
self.stats_to_track = set()
for name in self.config.logging.stats_tracker.stats_to_track:
self.stats_to_track.add(Stats[name])
logging.info("Successfully set up StatsTracker")
else:
logging.info("StatsTracker disabled; enable in config.")
def update_value(self, name: enum.Enum, value: float, step: int):
"""update stats dictionary with key value pair
Args:
name (enum.Enum): Enum name of statistic we are logging
value (float): value to update.
step (int): number of total iteration steps.
"""
if name in self.stats_to_track:
self.step = step
self.new_key = not name in self.stats_dict or self.new_key
self.stats_dict[name] = value
def update_time(
self, name: enum.Enum, start_time: float, end_time: float, step: int = None, batch_size: int = None
):
"""update the stats dictionary with running averages/cumulative durations
Args:
name (enum.Enum): Enum name of statistic we are logging
start_time (float): start time for the call in seconds
end_time (float): end time when the call finished executing in seconds
step (int): number of total iteration steps. Defaults to None.
if None, reports duration without averaging
batch_size (int, optional): total number of rays in a batch;
if None, reports duration instead of batch per second. Defaults to None.
"""
if name in self.stats_to_track:
self.step = step
self.new_key = not name in self.stats_dict or self.new_key
val = end_time - start_time
if batch_size:
# calculate the batch per second stat
val = batch_size / val
if step is not None:
# calculate updated average over the buffered max history
curr_history = self.stats_dict.get(name, {"buffer": [], "avg": 0})
curr_buffer = curr_history["buffer"]
curr_avg = curr_history["avg"]
if len(curr_buffer) >= self.max_history:
curr_buffer.pop(0)
curr_buffer.append(val)
curr_avg = sum(curr_buffer) / len(curr_buffer)
self.stats_dict[name] = {"buffer": curr_buffer, "avg": curr_avg}
else:
# logging total time instead of average
self.stats_dict[name] = val
if name == Stats.ITER_TRAIN_TIME and Stats.ETA in self.stats_to_track:
# update ETA if logging iteration train time
remain_iter = self.config.graph.max_num_iterations - step
self.stats_dict[Stats.ETA] = remain_iter * self.stats_dict[name]["avg"]
def handle_header(self):
"""helper to handle the printing of the header labels"""
if self.step == 0 or self.new_key:
mssg = f"{'Step (% Done)':<20}"
for k in self.stats_dict:
mssg += f"{k.value:<20} "
if self.step > 0:
cursorup(len(self.past_stats) + 2)
print(mssg)
print("-" * len(mssg))
if self.step > 0:
for mssg in self.past_stats:
print(mssg)
def handle_stats(self, fraction_done):
"""helper to print out the stats in a readable format"""
if self.step > 0:
curr_mssg = f"{self.step} ({fraction_done*100:.02f}%)"
curr_mssg = f"{curr_mssg:<20}"
for k, v in self.stats_dict.items():
if isinstance(v, dict):
v = v["avg"]
if "(time)" in k.value:
v = str(datetime.timedelta(seconds=v))
elif "(ms)" in k.value:
ms = v * 1e3
v = f"{ms:0.4f}"
else:
v = f"{v:0.4f}"
curr_mssg += f"{v:<20} "
# update the history buffer
# if len(self.past_stats) >= self.max_history:
# self.past_stats.pop(0)
# cursorup(len(self.past_stats) + 1)
# for mssg in self.past_stats:
# print(mssg)
print(curr_mssg)
# self.past_stats.append(curr_mssg)
@check_print_stats_step
def print_stats(self, fraction_done: float):
"""helper to print out the stats dictionary.
Args:
fraction_done (float): fraction of steps executed in training iterations
"""
# print a new header line if there is a new key added
# self.handle_header()
# generate a new stats reporting message
self.handle_stats(fraction_done)
def dump_stats(self):
"""Dump stats locally to a json file"""
raise NotImplementedError
def cursorup(x: int):
"""utility tool to move the cursor up on the terminal
Args:
x (int): amount of lines to move cursor upward
"""
print(f"\r\033[{x}A", end="")

245
mattport/utils/writer.py
View File

@ -3,6 +3,8 @@ Generic Writer class
"""
import datetime
import enum
import os
from abc import abstractmethod
from typing import Dict
@ -10,66 +12,147 @@ from typing import Dict
import imageio
import numpy as np
import torch
from omegaconf import DictConfig
from omegaconf import DictConfig, ListConfig
from torch.utils.tensorboard import SummaryWriter
from torchtyping import TensorType
import mattport.utils.writer
from mattport.utils import comms
from mattport.utils.decorators import check_main_thread, decorate_all
to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8)
EVENT_WRITERS = []
EVENT_STORAGE = []
GLOBAL_BUFFER = {}
def write_image(name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None) -> None:
"""Write image to all writers
Args:
name (str): data identifier
image (TensorType["H", "W", "C"]): rendered image to write
step (int): the time step to log
group (str): the group e.g., "Loss", "Accuracy", "Time"
prefix (str): the prefix e.g., "train-", "test-"
"""
assert comms.is_main_process(), "Writing out with process other than main"
for writer in EVENT_WRITERS:
writer.write_image(name, image, step, group, prefix)
def write_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
"""Write scalar to all writers
class EventType(enum.Enum):
"""Possible Event types that can be logged
Args:
name (str): data identifier
step (int): the time step to log
group (str): the group e.g., "Loss", "Accuracy", "Time"
prefix (str): the prefix e.g., "train-", "test-"
enum (_type_): _description_
"""
assert comms.is_main_process(), "Writing out with process other than main"
for writer in EVENT_WRITERS:
writer.write_scalar(name, scalar, step, group, prefix)
IMAGE = "write_image"
SCALAR = "write_scalar"
DICT = "write_scalar_dict"
def write_scalar_dict(scalar_dict: Dict[str, float], step: int, group: str = None, prefix: str = None) -> None:
"""Write dictionary of scalars to all writers
@check_main_thread
def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None):
"""Setter function to place images into the queue to be written out"""
EVENT_STORAGE.append(
{"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix}
)
Args:
scalar_dict (dict): dictionary containing all scalar values with key names and quantities
step (int): the time step to log
group (str): the group e.g., "Loss", "Accuracy", "Time"
prefix (str): the prefix e.g., "train-", "test-"
@check_main_thread
def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None):
"""Setter function to place scalars into the queue to be written out"""
GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"]
GLOBAL_BUFFER["events"][name] = scalar
EVENT_STORAGE.append(
{"name": name, "write_type": EventType.SCALAR, "event": scalar, "step": step, "group": group, "prefix": prefix}
)
@check_main_thread
def put_dict(name: str, scalar_dict: float, step: int, group: str = None, prefix: str = None):
"""Setter function to place a dictionary of scalars into the queue to be written out"""
EVENT_STORAGE.append(
{
"name": name,
"write_type": EventType.DICT,
"event": scalar_dict,
"step": step,
"group": group,
"prefix": prefix,
}
)
@check_main_thread
def put_time(
name: str,
start_time: float,
end_time: float,
step: int,
group: str = None,
prefix: str = None,
avg_over_iters: bool = False,
avg_over_batch: int = None,
update_eta: bool = False,
):
"""Setter function to place a time element into the queue to be written out
Processes the time info according to the options:
avg_over_iters (bool): if True, calculate and record a running average of the times
avg_over_batch (int): if set, the size of the batch for which we take the average over (batch/second)
update_eta (bool): if True, update the ETA. should only be set for the training iterations/s
"""
assert comms.is_main_process(), "Writing out with process other than main"
GLOBAL_BUFFER["step"] = step
GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"]
val = end_time - start_time
if avg_over_batch:
val = avg_over_batch / val
if avg_over_iters:
curr_event = GLOBAL_BUFFER["events"].get(name, {"buffer": [], "avg": 0})
curr_buffer = curr_event["buffer"]
curr_avg = curr_event["avg"]
if len(curr_buffer) >= GLOBAL_BUFFER["max_history"]:
curr_buffer.pop(0)
curr_buffer.append(val)
curr_avg = sum(curr_buffer) / len(curr_buffer)
put_scalar(name, curr_avg, step, group, prefix)
GLOBAL_BUFFER["events"][name] = {"buffer": curr_buffer, "avg": curr_avg}
else:
GLOBAL_BUFFER["events"][name] = val
put_scalar(name, val, step, group, prefix)
if update_eta:
## NOTE: eta should be called with avg train iteration time
remain_iter = GLOBAL_BUFFER["max_iter"] - step
remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"]
GLOBAL_BUFFER["events"]["ETA"] = remain_time
put_scalar("ETA", remain_time, step, group, prefix)
@check_main_thread
def write_out_storage():
"""Function that writes all the events in storage to all the writer locations"""
for writer in EVENT_WRITERS:
writer.write_scalar_dict(scalar_dict, step, group, prefix)
for event in EVENT_STORAGE:
write_func = getattr(writer, event["write_type"].value)
if event["write_type"] == EventType.DICT:
write_func(event["event"], event["step"], event["group"], event["prefix"])
if isinstance(writer, LocalWriter):
continue
else:
write_func(event["name"], event["event"], event["step"], event["group"], event["prefix"])
if isinstance(writer, LocalWriter):
break
EVENT_STORAGE.clear()
@check_main_thread
def setup_event_writers(config: DictConfig) -> None:
"""Initialization of all event writers specified in config"""
if comms.is_main_process():
for writer_type in config:
curr_writer = getattr(mattport.utils.writer, writer_type)
EVENT_WRITERS.append(curr_writer(config[writer_type].save_dir))
logging_configs = config.logging.writer
for writer_type in logging_configs:
writer_class = getattr(mattport.utils.writer, writer_type)
curr_config = logging_configs[writer_type]
if writer_type == "LocalWriter":
curr_writer = writer_class(curr_config.save_dir, curr_config.stats_to_track, config)
else:
curr_writer = writer_class(curr_config.save_dir)
EVENT_WRITERS.append(curr_writer)
## configure all the global buffer basic information
GLOBAL_BUFFER["max_iter"] = config.graph.max_num_iterations
GLOBAL_BUFFER["max_history"] = config.logging.max_history
GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log
GLOBAL_BUFFER["new_key"] = True
GLOBAL_BUFFER["events"] = {}
def get_tensorboard_name(name: str, group: str = None, prefix: str = None):
@ -165,16 +248,94 @@ class TensorboardWriter(Writer):
self.tb_writer.add_scalar(tensorboard_name, scalar, step)
def _cursorup(x: int):
"""utility tool to move the cursor up on the terminal
Args:
x (int): amount of lines to move cursor upward
"""
print(f"\r\033[{x}A", end="\x1b[1K\r")
def _format_time(v):
if v < 1.0:
ms = v * 1e3
v = f"{ms:0.4f} ms"
else:
v = str(datetime.timedelta(seconds=v))
return v
def _consolidate_events():
latest_map = {}
for event in EVENT_STORAGE:
name = event["name"]
latest_map[name] = event["event"]
return latest_map
@decorate_all([check_main_thread])
class LocalWriter(Writer):
"""Local Writer Class"""
def __init__(self, save_dir: str, stats_to_track: ListConfig, config: DictConfig):
super().__init__(save_dir)
self.stats_to_track = stats_to_track
self.past_stats = []
self.max_mssg_len = 0
self.config = config
def write_image(
self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None
) -> None:
image = to8b(image)
image_path = os.path.join(self.save_dir, f"{name}.jpg")
imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0))
if name in self.stats_to_track:
image = to8b(image)
image_path = os.path.join(self.save_dir, f"{name}.jpg")
imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0))
def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
raise NotImplementedError
latest_map = _consolidate_events()
self.handle_header(latest_map)
self.handle_stats(latest_map)
GLOBAL_BUFFER["new_key"] = False
def handle_header(self, latest_map):
"""helper to handle the printing of the header labels"""
if len(self.past_stats) == 0 or GLOBAL_BUFFER["new_key"]:
mssg = f"{'Step (% Done)':<20}"
for name, _ in latest_map.items():
if name in self.stats_to_track:
mssg += f"{name:<20} "
if len(self.past_stats) > 1:
_cursorup(len(self.past_stats) - 1)
print(mssg)
print("-" * len(mssg))
if len(self.past_stats) > 0:
for mssg in self.past_stats:
print(mssg)
self.max_mssg_len = len(mssg)
def handle_stats(self, latest_map, padding=" "):
"""helper to print out the stats in a readable format"""
step = GLOBAL_BUFFER["step"]
fraction_done = step / GLOBAL_BUFFER["max_iter"]
curr_mssg = f"{step} ({fraction_done*100:.02f}%)"
curr_mssg = f"{curr_mssg:<20}"
for name, v in latest_map.items():
if name in self.stats_to_track:
if "(time)" in name:
v = _format_time(v)
else:
v = f"{v:0.4f}"
curr_mssg += f"{v:<20} "
# update the history buffer
if len(self.past_stats) > GLOBAL_BUFFER["max_history"]:
self.past_stats.pop(0)
_cursorup(len(self.past_stats))
elif len(self.past_stats) > 1:
_cursorup(len(self.past_stats) - 1)
for mssg in self.past_stats:
print(f"{mssg:{padding}<{self.max_mssg_len}}")
self.past_stats.append(curr_mssg)