mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 10:34:29 +08:00
Set up profiler and moved stats tracker to its own class; set up configurables for stats tracker and profiler (#42)
* added skeleton for terminal logger * done skeleton terminal writer logic... todo: actual calls to the write classes * moved lambda to top function * added logic to support timing * prettier printing * prettier printing * lint * separated out timing and other stats variables. printing in dict format * keeping logging history and wiping rest * updated logic to calculate total times instead of average times * made stats tracker configurable * moved stats tracker and profiler to own class file; implemented basic functionality for profiler * lint removing duplicate function
This commit is contained in:
evonneng
GitHub
parent
2e350c148f
commit
609fcb04f6
@ -3,6 +3,13 @@ defaults:
|
||||
- override hydra/job_logging: colorlog
|
||||
- override hydra/hydra_logging: colorlog
|
||||
|
||||
debug:
|
||||
enable_stats: True
|
||||
stats_tracker:
|
||||
max_history: 10
|
||||
stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME]
|
||||
enable_profiler: True
|
||||
|
||||
hydra:
|
||||
run:
|
||||
# Configure output dir of each experiment programmatically from the arguments
|
||||
|
||||
@ -13,13 +13,14 @@ from omegaconf import DictConfig
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
|
||||
from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one
|
||||
from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
|
||||
from mattport.nerf.dataset.utils import get_dataset_inputs
|
||||
from mattport.nerf.metrics import get_psnr
|
||||
from mattport.nerf.optimizers import Optimizers
|
||||
from mattport.utils import profiler
|
||||
from mattport.utils.decorators import check_main_thread
|
||||
from mattport.utils.io import Stats, StatsTracker
|
||||
from mattport.utils.stats_tracker import Stats, StatsTracker
|
||||
from mattport.utils.writer import TensorboardWriter
|
||||
|
||||
logging.getLogger("PIL").setLevel(logging.WARNING)
|
||||
@ -48,7 +49,9 @@ class Trainer:
|
||||
local_rank, world_size, save_dir=os.path.join(os.getcwd(), "writer")
|
||||
)
|
||||
self.stats = StatsTracker(config, self.is_main_thread)
|
||||
profiler.PROFILER = profiler.Profiler(config, self.is_main_thread)
|
||||
|
||||
@profiler.time_function
|
||||
def setup_dataset(self):
|
||||
"""_summary_"""
|
||||
dataset_inputs = get_dataset_inputs(**self.config.dataset)
|
||||
@ -72,6 +75,7 @@ class Trainer:
|
||||
)
|
||||
# TODO(ethan): implement the test dataset
|
||||
|
||||
@profiler.time_function
|
||||
def setup_graph(self):
|
||||
"""_summary_"""
|
||||
dataset_inputs = get_dataset_inputs(**self.config.dataset)
|
||||
@ -174,8 +178,10 @@ class Trainer:
|
||||
self.stats.print_stats(i / num_iterations)
|
||||
|
||||
self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1)
|
||||
self.stats.dump_stats()
|
||||
self.stats.print_stats(-1)
|
||||
profiler.PROFILER.print_profile()
|
||||
|
||||
@profiler.time_function
|
||||
def train_iteration(self, batch: dict, step: int):
|
||||
"""Run one iteration with a batch of inputs."""
|
||||
# move batch to correct device
|
||||
@ -190,6 +196,7 @@ class Trainer:
|
||||
self.optimizers.optimizer_step_all()
|
||||
return loss_dict
|
||||
|
||||
@profiler.time_function
|
||||
def test_image(self, image_idx, step):
|
||||
"""Test a specific image."""
|
||||
image = self.train_image_dataset[image_idx]["image"] # ground truth
|
||||
|
||||
@ -1,28 +1,69 @@
|
||||
"""
|
||||
Decorator definitions
|
||||
"""
|
||||
from typing import Callable
|
||||
from typing import Callable, List
|
||||
|
||||
|
||||
def decorate_all(decorators: List[Callable]) -> Callable:
|
||||
"""A decorator to decorate all member functions of a class"""
|
||||
|
||||
def decorate(cls):
|
||||
for attr in cls.__dict__:
|
||||
if callable(getattr(cls, attr)) and attr != "__init__":
|
||||
for decorator in decorators:
|
||||
setattr(cls, attr, decorator(getattr(cls, attr)))
|
||||
return cls
|
||||
|
||||
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.debug.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.debug.enable_stats:
|
||||
ret = func(self, *args, **kwargs)
|
||||
return ret
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def check_main_thread(func: Callable) -> Callable:
|
||||
"""Declarator: check if you are on main thread"""
|
||||
"""Decorator: check if you are on main thread"""
|
||||
|
||||
def wrapper(self, *args, **kwargs):
|
||||
ret = None
|
||||
if self.is_main_thread:
|
||||
func(self, *args, **kwargs)
|
||||
ret = func(self, *args, **kwargs)
|
||||
return ret
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def check_print_stats_step(func: Callable) -> Callable:
|
||||
"""Declarator: check if it is time to print stats update"""
|
||||
"""Decorator: check if it is time to print stats update"""
|
||||
|
||||
def wrapper(self, *args, **kwargs):
|
||||
ret = None
|
||||
if (
|
||||
self.step % self.config.steps_per_log == 0
|
||||
or (self.config.steps_per_save and self.step % self.config.steps_per_save == 0)
|
||||
or self.step % self.config.steps_per_test == 0
|
||||
):
|
||||
func(self, *args, **kwargs)
|
||||
ret = func(self, *args, **kwargs)
|
||||
return ret
|
||||
|
||||
return wrapper
|
||||
|
||||
@ -1,139 +1,8 @@
|
||||
"""
|
||||
Input/output utils.
|
||||
"""
|
||||
import datetime
|
||||
import enum
|
||||
import json
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from mattport.utils.decorators import check_main_thread, check_print_stats_step
|
||||
|
||||
|
||||
class Stats(enum.Enum):
|
||||
"""Possible Stats values for StatsTracker
|
||||
|
||||
Args:
|
||||
enum (_type_): _description_
|
||||
"""
|
||||
|
||||
ITER_LOAD_TIME = "Data Load (time)"
|
||||
ITER_TRAIN_TIME = "Train Iter (time)"
|
||||
TOTAL_TRAIN_TIME = "Train Total (time)"
|
||||
RAYS_PER_SEC = "Rays Per Sec (1/s)"
|
||||
CURR_TEST_PSNR = "Test PSNR"
|
||||
ETA = "ETA (time)"
|
||||
|
||||
|
||||
class StatsTracker:
|
||||
"""Stats Tracker class"""
|
||||
|
||||
def __init__(self, config: DictConfig, is_main_thread: bool, max_history: Optional[int] = 20):
|
||||
self.config = config
|
||||
self.is_main_thread = is_main_thread
|
||||
self.max_history = max_history
|
||||
self.step = 0
|
||||
self.stats_dict = {}
|
||||
self.past_stats = []
|
||||
self.new_key = False
|
||||
|
||||
def update_value(self, name: str, value: float, step: int):
|
||||
"""update stats dictionary with key value pair
|
||||
|
||||
Args:
|
||||
name (str): name of statistic we are logging
|
||||
value (float): value to update.
|
||||
step (int): number of total iteration steps.
|
||||
"""
|
||||
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: str, start_time: float, end_time: float, step: int, batch_size: int = None):
|
||||
"""update the stats dictionary with running averages/cumulative durations
|
||||
|
||||
Args:
|
||||
name (str): 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.
|
||||
batch_size (int, optional): total number of rays in a batch;
|
||||
if None, reports duration instead of batch per second. Defaults to None.
|
||||
"""
|
||||
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 == -1:
|
||||
# logging total time instead of average
|
||||
self.stats_dict[name] = val
|
||||
else:
|
||||
# calculate updated average
|
||||
self.stats_dict[name] = (self.stats_dict.get(name, 0) * step + val) / (step + 1)
|
||||
|
||||
if name == Stats.ITER_TRAIN_TIME:
|
||||
# update ETA if logging iteration train time
|
||||
remain_iter = self.config.max_num_iterations - step
|
||||
self.stats_dict[Stats.ETA] = remain_iter * self.stats_dict[name]
|
||||
|
||||
@check_print_stats_step
|
||||
@check_main_thread
|
||||
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
|
||||
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)
|
||||
|
||||
# generate a new stats reporting message
|
||||
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 "(time)" in k.value:
|
||||
v = str(datetime.timedelta(seconds=v))
|
||||
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_main_thread
|
||||
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="")
|
||||
|
||||
|
||||
def load_from_json(filename: str):
|
||||
|
||||
63
mattport/utils/profiler.py
Normal file
63
mattport/utils/profiler.py
Normal file
@ -0,0 +1,63 @@
|
||||
"""
|
||||
Profiler base class and functionality
|
||||
"""
|
||||
import sys
|
||||
import logging
|
||||
import time
|
||||
from typing import Callable
|
||||
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from mattport.utils.decorators import check_main_thread, check_profiler_enabled, decorate_all
|
||||
|
||||
PROFILER = None
|
||||
|
||||
|
||||
def time_function(func: Callable) -> Callable:
|
||||
"""Decorator: time a function call"""
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
start = time.time()
|
||||
ret = func(*args, **kwargs)
|
||||
vals = vars(sys.modules[func.__module__])
|
||||
class_str = ""
|
||||
for attr in func.__qualname__.split(".")[:-1]:
|
||||
class_str += f"{vals[attr].__qualname__}_"
|
||||
class_str += func.__name__
|
||||
PROFILER.update_time(class_str, start, time.time())
|
||||
return ret
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@decorate_all([check_main_thread, check_profiler_enabled])
|
||||
class Profiler:
|
||||
"""Profiler class"""
|
||||
|
||||
def __init__(self, config: DictConfig, is_main_thread: bool):
|
||||
self.config = config
|
||||
if self.config.debug.enable_profiler:
|
||||
self.is_main_thread = is_main_thread
|
||||
self.profiler_dict = {}
|
||||
|
||||
def update_time(self, func_name: str, start_time: float, end_time: float):
|
||||
"""update the profiler dictionary with running averages of durations
|
||||
|
||||
Args:
|
||||
func_name (str): the function name that is being profiled
|
||||
start_time (float): the start time when function is called
|
||||
end_time (float): the end time when function terminated
|
||||
"""
|
||||
val = end_time - start_time
|
||||
func_dict = self.profiler_dict.get(func_name, {"val": 0, "step": 0})
|
||||
prev_val = func_dict["val"]
|
||||
prev_step = func_dict["step"]
|
||||
self.profiler_dict[func_name] = {"val": (prev_val * prev_step + val) / (prev_step + 1), "step": prev_step + 1}
|
||||
|
||||
def print_profile(self):
|
||||
"""helper to print out the profiler stats"""
|
||||
logging.info("Printing profiling stats, from longest to shortest duration in seconds")
|
||||
sorted_keys = [k for k, _ in sorted(self.profiler_dict.items(), key=lambda item: item[1]["val"], reverse=True)]
|
||||
for k in sorted_keys:
|
||||
val = f"{self.profiler_dict[k]['val']:0.4f}"
|
||||
print(f"{k:<20}: {val:<20}")
|
||||
141
mattport/utils/stats_tracker.py
Normal file
141
mattport/utils/stats_tracker.py
Normal file
@ -0,0 +1,141 @@
|
||||
"""
|
||||
Stats tracker base class and functionality
|
||||
"""
|
||||
import datetime
|
||||
import enum
|
||||
import logging
|
||||
|
||||
from omegaconf import DictConfig
|
||||
from mattport.utils.decorators import check_main_thread, check_print_stats_step, check_stats_enabled, decorate_all
|
||||
|
||||
|
||||
class Stats(enum.Enum):
|
||||
"""Possible Stats values for StatsTracker
|
||||
|
||||
Args:
|
||||
enum (_type_): _description_
|
||||
"""
|
||||
|
||||
ITER_LOAD_TIME = "Data Load (time)"
|
||||
ITER_TRAIN_TIME = "Train Iter (time)"
|
||||
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_main_thread, check_stats_enabled])
|
||||
class StatsTracker:
|
||||
"""Stats Tracker class"""
|
||||
|
||||
def __init__(self, config: DictConfig, is_main_thread: bool):
|
||||
self.config = config
|
||||
if self.config.debug.enable_stats:
|
||||
self.is_main_thread = is_main_thread
|
||||
self.max_history = self.config.debug.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.debug.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, 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.
|
||||
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 == -1:
|
||||
# logging total time instead of average
|
||||
self.stats_dict[name] = val
|
||||
else:
|
||||
# calculate updated average
|
||||
self.stats_dict[name] = (self.stats_dict.get(name, 0) * step + val) / (step + 1)
|
||||
|
||||
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.max_num_iterations - step
|
||||
self.stats_dict[Stats.ETA] = remain_iter * self.stats_dict[name]
|
||||
|
||||
@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
|
||||
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)
|
||||
|
||||
# generate a new stats reporting message
|
||||
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 "(time)" in k.value:
|
||||
v = str(datetime.timedelta(seconds=v))
|
||||
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)
|
||||
|
||||
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="")
|
||||
@ -13,7 +13,7 @@ import torch
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
from torchtyping import TensorType
|
||||
|
||||
from mattport.utils.decorators import check_main_thread
|
||||
from mattport.utils.decorators import check_main_thread, decorate_all
|
||||
|
||||
to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8)
|
||||
|
||||
@ -69,6 +69,7 @@ class Writer:
|
||||
self.write_scalar(name, scalar, step, group=group, prefix=prefix)
|
||||
|
||||
|
||||
@decorate_all([check_main_thread])
|
||||
class TensorboardWriter(Writer):
|
||||
"""Tensorboard Writer Class"""
|
||||
|
||||
@ -77,7 +78,6 @@ class TensorboardWriter(Writer):
|
||||
if self.is_main_thread:
|
||||
self.tb_writer = SummaryWriter(log_dir=self.save_dir)
|
||||
|
||||
@check_main_thread
|
||||
def write_image(
|
||||
self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None
|
||||
) -> None:
|
||||
@ -90,7 +90,6 @@ class TensorboardWriter(Writer):
|
||||
x = to8b(x)
|
||||
self.tb_writer.add_image(name, x, step, dataformats="HWC")
|
||||
|
||||
@check_main_thread
|
||||
def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
|
||||
"""Tensorboard method to write a single scalar value to the logger
|
||||
|
||||
@ -105,10 +104,10 @@ class TensorboardWriter(Writer):
|
||||
self.tb_writer.add_scalar(f"{group_s}{prefix_s}{name}", scalar, step)
|
||||
|
||||
|
||||
@decorate_all([check_main_thread])
|
||||
class LocalWriter(Writer):
|
||||
"""Local Writer Class"""
|
||||
|
||||
@check_main_thread
|
||||
def write_image(
|
||||
self, name: str, x: TensorType["H", "W", 3], step: int, group: str = None, prefix: str = None
|
||||
) -> None:
|
||||
@ -116,6 +115,5 @@ class LocalWriter(Writer):
|
||||
image_path = os.path.join(self.save_dir, f"{name}.jpg")
|
||||
imageio.imwrite(image_path, np.uint8(x.cpu().numpy() * 255.0))
|
||||
|
||||
@check_main_thread
|
||||
def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
Loading…
Reference in New Issue
Block a user