nerfstudio/mattport/utils/stats_tracker.py

"""
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="")