src/diffusion/train/training_loop.py

import os
import time
from types import SimpleNamespace
from typing import Dict
from pathlib import Path
from os.path import join as pjoin
import numpy as np
from tqdm import tqdm
from hydra.utils import instantiate

import torch
from torch.optim import AdamW
from torch.nn.utils import clip_grad_norm_
from torch.nn.parallel import DistributedDataParallel as DDP
from einops import rearrange


# from .. import logger
from ..config import Config, DataLoaderConfig
from ..logger_new import mylogger
from .train_platforms import TrainPlatform
from ..gaussian_diffusion import GaussianDiffusion
from ..model.evaluator import Evaluator
from ..utils.mics import get_device
from ..resample import create_named_schedule_sampler, LossAwareSampler
from ..data_loader.get_data import get_dataloader
from ...utils import debug
from .. import dist as dist_utils
from ...constant import gesture_list
from ..model.cls_free_sampler import ClassifierFreeSampleWrapper
from ...visualize.visualize import MultiMotionVisualizer


# For ImageNet experiments, this was a good default value.
# We found that the lg_loss_scale quickly climbed to
# 20-21 within the first ~1K steps of training.
INITIAL_LOG_LOSS_SCALE = 20.0


class TrainLoop:
    def __init__(
        self,
        cfg:Config,
        train_platform:TrainPlatform,
        model:torch.nn.Module,
        diffusion:GaussianDiffusion,
    ):
        self.cfg = cfg
        # self.dataset = cfg.dataset
        self.train_platform = train_platform
        self.diffusion = diffusion
        self.repr = cfg.data.repr
        self.step = 0
        self.resume_step = None
        self.lr = cfg.train.optimizer.lr
        self.log_interval = cfg.train.log_interval
        self.save_interval = cfg.train.save_interval
        self.weight_decay = cfg.train.optimizer.weight_decay
        if cfg.data.fixed_length > 0:
            self.sample_length = cfg.data.fixed_length
        else:
            self.sample_length = cfg.data.max_length
        self.save_dir = cfg.train.save_dir
        self.num_steps = cfg.train.num_steps

        torch.set_float32_matmul_precision("high")

        self.model = model
        self.sample_model = ClassifierFreeSampleWrapper(model, scale=cfg.train.sample.guidance_param)
        self.resume_checkpoint = cfg.train.resume_checkpoint
        self._load_and_sync_parameters()

        self.device = get_device()
        self.model.to(self.device)

        self.set_ddp()

        mylogger.info("Creating dataset ...")
        # self.data = get_dataset_loader('train', cfg, cfg.batch_size, shuffle=True)
        self.train_dataset = instantiate(cfg.data, split='train', debug=False)
        mylogger.info("Creating dataloader ...")
        self.train_dataloader = get_dataloader(self.train_dataset, cfg.train.dataloader)
        self.num_epochs = self.num_steps // len(self.train_dataloader) + 1

        self._add_model_params_to_optimizer()
        if self.resume_step:
            self._load_optimizer_state()

        self.schedule_sampler_type = 'uniform'
        self.schedule_sampler = create_named_schedule_sampler(self.schedule_sampler_type, diffusion)

        # TODO: Evaluation Module
        # if args.dataset in ['kit', 'humanml'] and args.eval_during_training:
        #     mm_num_samples = 0  # mm is super slow hence we won't run it during training
        #     mm_num_repeats = 0  # mm is super slow hence we won't run it during training
        #     gen_loader = get_dataset_loader(name=args.dataset, batch_size=args.eval_batch_size, num_frames=None,
        #                                     split=args.eval_split,
        #                                     hml_mode='eval')

        #     self.eval_gt_data = get_dataset_loader(name=args.dataset, batch_size=args.eval_batch_size, num_frames=None,
        #                                            split=args.eval_split,
        #                                            hml_mode='gt')
        #     self.eval_wrapper = EvaluatorMDMWrapper(args.dataset, dist_util.dev())
        #     self.eval_data = {
        #         'test': lambda: eval_humanml.get_mdm_loader(
        #             model, diffusion, args.eval_batch_size,
        #             gen_loader, mm_num_samples, mm_num_repeats, gen_loader.dataset.opt.max_motion_length,
        #             args.eval_num_samples, scale=1.,
        #         )
        #     }

        self.val_during_training = cfg.train.val_during_training
        if self.val_during_training:
            self.val_interval = cfg.train.val_cfg.val_interval
            self.val_dataset = instantiate(cfg.data, split='val', debug=False)
            self.val_dataloader = get_dataloader(self.val_dataset, cfg.train.val_cfg.dataloader)

        self.eval_during_training = cfg.train.eval_during_training
        if self.eval_during_training:
            self.eval_interval = cfg.train.eval_cfg.eval_interval
            self.eval_dataset_on_train = self.train_dataset
            self.eval_dataset_on_val = self.val_dataset if hasattr(self, 'val_dataset') else instantiate(cfg.data, split='val', debug=False)
            self.evaluate_helper = Evaluator(
                sample_model=self.sample_model,
                train_dataset=self.eval_dataset_on_train,
                val_dataset=self.eval_dataset_on_val,
                dataloader_cfg=cfg.train.eval_cfg.dataloader,
                sample_fn=self.diffusion.p_sample_loop,
                nfeats=self.model_without_ddp.nfeats, njoints=self.model_without_ddp.njoints,
                sample_length=self.sample_length,
                num_samples_on_train=cfg.train.eval_cfg.num_samples_on_train,
                num_samples_on_val=cfg.train.eval_cfg.num_samples_on_val,
                num_samples_per_condition=cfg.train.eval_cfg.num_samples_per_condition,
            )

        self.viz_during_training = cfg.train.viz_during_training
        if self.viz_during_training:
            self.viz_dataloader = get_dataloader(self.train_dataset, DataLoaderConfig(batch_size=dist_utils.get_world_size(), num_workers=0, shuffle=True))
            self.iter_viz_data = iter(self.viz_dataloader)
            self.viz_samples_count = cfg.train.viz_cfg.samples_count

    def _add_model_params_to_optimizer(self):
        params_dict = {
            pn: p
            for pn, p in self.model.named_parameters()
            if p.requires_grad
        }

        decay_params = [
            p
            for n, p in params_dict.items()
            if p.dim() >= 2
        ]

        nodecay_params = [
            p
            for n, p in params_dict.items()
            if p.dim() < 2
        ]

        optim_groups = [
            {"params": decay_params, "weight_decay": self.weight_decay},
            {"params": nodecay_params, "weight_decay": 0.0}
        ]

        self.opt = AdamW(
            optim_groups,
            lr=self.lr,
            betas=(0.9, 0.95),
            eps=1e-8,
            fused=True
        )

    def _remove_module_prefix(self, state_dict):
        if any(key.startswith('module.') for key in state_dict.keys()):
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:]
                new_state_dict[name] = v

            return new_state_dict
        return state_dict

    def _load_and_sync_parameters(self):
        resume_checkpoint = self.resume_checkpoint

        if resume_checkpoint:
            self.resume_step = parse_resume_step_from_filename(Path(resume_checkpoint).name)
            self.step = self.resume_step
            mylogger.info(f"loading model from checkpoint: {resume_checkpoint}...")
            state_dict = torch.load(
                resume_checkpoint, map_location='cpu'
            )
            state_dict = self._remove_module_prefix(state_dict['state_dict'])
            self.model.load_state_dict(state_dict, strict=False)

    def _load_optimizer_state(self):
        main_checkpoint = self.resume_checkpoint
        opt_checkpoint = os.path.join(
            os.path.dirname(main_checkpoint), f"opt{self.resume_step:09}.pt"
        )
        if os.path.exists(opt_checkpoint):
            mylogger.info(f"loading optimizer state from checkpoint: {opt_checkpoint}")
            state_dict = torch.load(
                opt_checkpoint, map_location='cpu'
            )
            self.opt.load_state_dict(state_dict)

    def set_ddp(self):
        self.model.to(self.device)
        if dist_utils.is_dist_avail_and_initialized():
            self.model = DDP(self.model, device_ids=[self.device.index], output_device=self.device.index, static_graph=True)
        self.model_without_ddp = self.model.module if isinstance(self.model, DDP) else self.model

    def run_loop(self):

        for epoch in range(self.num_epochs):
            print(f'Starting epoch {epoch}')
            if dist_utils.is_dist_avail_and_initialized():
                self.train_dataloader.sampler.set_epoch(epoch)
            for motion, cond in tqdm(self.train_dataloader, desc=f'Epoch {epoch + 1} / {self.num_epochs} RANK {dist_utils.get_rank()}'):

                motion = motion.to(self.device)
                cond['y'] = {key: val.to(self.device) if torch.is_tensor(val) else val for key, val in cond['y'].items()}


                self.run_step(motion, cond)
                if self.step % self.log_interval == 0:
                    for key in mylogger.get_all_metrics().keys():
                        if key.startswith('val_'):
                            continue
                        if key.startswith('train_'):
                            mylogger.debug(f'step [{self.step}]: {key} [{mylogger.get_average(key):0.5f}]')

                        if '_q1' not in key and '_q2' not in key and '_q3' not in key and '_q0' not in key:
                            self.train_platform.report_scalar(
                                name=key,
                                value=mylogger.get_average(key),
                                iteration=self.step,
                                group_name='Train'
                            )
                        mylogger.clear_metric(key)

                    self.train_platform.report_scalar(
                        name='learning_rate',
                        value=self.opt.param_groups[0]['lr'],
                        iteration=self.step,
                        group_name='Train',
                    )

                if self.val_during_training and self.step % self.val_interval == 0:
                    # calculate validation loss
                    with torch.no_grad():
                        for motion, cond in tqdm(self.val_dataloader, desc=f'Validation on val set | RANK {dist_utils.get_rank()}'):
                            motion = motion.to(self.device)
                            cond['y'] = {key: val.to(self.device) if torch.is_tensor(val) else val for key, val in cond['y'].items()}
                            self.validation_forward(motion, cond, split='val')

                        for motion, cond in tqdm(self.train_dataloader, desc=f'Validation on train set | RANK {dist_utils.get_rank()}'):
                            motion = motion.to(self.device)
                            cond['y'] = {key: val.to(self.device) if torch.is_tensor(val) else val for key, val in cond['y'].items()}
                            self.validation_forward(motion, cond, split='train')

                    for key in mylogger.get_all_metrics().keys():
                        if not key.startswith('val_'):
                            continue

                        mylogger.debug(f'step [{self.step}]: {key} [{mylogger.get_average(key):0.5f}]')
                        if '_q0' not in key and '_q1' not in key and '_q2' not in key and '_q3' not in key:
                            self.train_platform.report_scalar(
                                name=key,
                                value=mylogger.get_average(key),
                                iteration=self.step,
                                group_name='Validation'
                            )
                        mylogger.clear_metric(key)

                if self.eval_during_training and self.step % self.eval_interval == 0:
                    self.model.eval()
                    self.evaluate()
                    self.model.train()

                if self.step % self.save_interval == 0:
                    self.save()

                    if self.viz_during_training:
                        for sample_id in range(self.viz_samples_count):
                            samples, gt, lengths, label, mse_losses, vel_losses = self.sample()
                            gt_single = gt[0]
                            length = lengths[0]
                            mse_loss = mse_losses[0]
                            vel_loss = vel_losses[0]

                            # Handle both dict (treble text) and list (normal text) label formats
                            if isinstance(label, dict):
                                # For treble text, create a combined label string
                                label_str = f"step: {self.step}_{dist_utils.get_rank() * self.viz_samples_count + sample_id} mse_loss: {mse_loss:.6f} vel: {vel_loss:.6f} \n\n[LEFT] {label['left'][0]} \n[RIGHT] {label['right'][0]} \n[RELATION] {label['two_hands_relation'][0]}"
                            else:
                                label_str = label[0]

                            if self.repr == 'joint_rot':
                                from ..utils.mics import rot_motion_to_dict

                            def process_motion(motion, title=None):
                                left_motion, right_motion = np.split(
                                    motion.reshape(motion.shape[0], self.model_without_ddp.njoints, self.model_without_ddp.nfeats),
                                    indices_or_sections=[self.model_without_ddp.njoints // 2],
                                    axis=1
                                )  # (T, J_single, D), (T, J_single, D)
                                cur_motion_to_visualize = dict()
                                if self.repr == 'joint_pos':
                                    cur_motion_to_visualize.update(
                                        dict(
                                            type='skeleton',
                                            left_motion=left_motion,
                                            right_motion=right_motion,
                                        )
                                    )
                                elif self.repr in ['joint_pos_w_scalar_rot', 'joint_pos_w_axisangle_rot']:
                                    cur_motion_to_visualize.update(
                                        dict(
                                            type='skeleton',
                                            left_motion=left_motion[:, :, :3],
                                            right_motion=right_motion[:, :, :3],
                                        )
                                    )

                                elif self.repr == 'joint_rot':
                                    left_motion = rot_motion_to_dict(left_motion)
                                    right_motion = rot_motion_to_dict(right_motion)

                                    cur_motion_to_visualize.update(
                                        type='mano',
                                        left_motion=left_motion,
                                        right_motion=right_motion,
                                    )
                                if title is not None:
                                    cur_motion_to_visualize['title'] = title

                                return cur_motion_to_visualize

                            # GIF 1: Original - GT + all denoising steps
                            motions_to_visualize = []
                            motions_to_visualize.append(process_motion(gt_single[:length], title='Ground Truth'))
                            for i in range(len(samples)):
                                cur_sample = samples[i, :length]
                                if self.cfg.train.viz_cfg.denoising_steps and len(self.cfg.train.viz_cfg.denoising_steps) > 0:
                                    title = f"Denoising Step: {self.cfg.train.viz_cfg.denoising_steps[i]}"
                                else:
                                    title = None

                                motions_to_visualize.append(
                                    process_motion(cur_sample, title=title)
                                )

                            MultiMotionVisualizer.create_3d_animation(
                                motions=motions_to_visualize,
                                text=label_str,
                                save_path=pjoin(self.save_dir, f'step_{self.step}_{dist_utils.get_rank() * self.viz_samples_count + sample_id}_all_steps.gif'),
                                fps=30
                            )

                            dist_utils.barrier()

                            for rank in range(dist_utils.get_world_size()):
                                self.train_platform.report_video(
                                    name=f'video_{rank}',
                                    video_path=pjoin(self.save_dir, f'step_{self.step}_{rank * self.viz_samples_count + sample_id}_all_steps.gif'),
                                    video_format='gif',
                                    group_name='Train'
                                )


                self.step += 1



        # Save the last checkpoint if it wasn't already saved.
        if (self.step - 1) % self.save_interval != 0:
            self.save()
            self.evaluate()

    def evaluate(self):
        if not self.cfg.train.eval_during_training:
            return
        eval_result = self.evaluate_helper.evaluate(split='train')
        eval_result_on_val = self.evaluate_helper.evaluate(split='val')

        if dist_utils.is_main_process():
            eval_result.update(eval_result_on_val)
            for key in eval_result.keys():
                mylogger.debug(f"EVALUATION {key}: [{eval_result[key]:0.5f}]")
                self.train_platform.report_scalar(
                    name=key,
                    value=eval_result[key],
                    iteration=self.step,
                    group_name='Evaluation'
                )
        dist_utils.barrier()

        # start_eval = time.time()
        # if self.eval_wrapper is not None:
        #     print('Running evaluation loop: [Should take about 90 min]')
        #     log_file = os.path.join(self.save_dir, f'eval_humanml_{(self.step):09d}.log')
        #     diversity_times = 300
        #     mm_num_times = 0  # mm is super slow hence we won't run it during training
        #     eval_dict = eval_humanml.evaluation(
        #         self.eval_wrapper, self.eval_gt_data, self.eval_data, log_file,
        #         replication_times=self.cfg.eval_rep_times, diversity_times=diversity_times, mm_num_times=mm_num_times, run_mm=False)
        #     print(eval_dict)
        #     for k, v in eval_dict.items():
        #         if k.startswith('R_precision'):
        #             for i in range(len(v)):
        #                 self.train_platform.report_scalar(name=f'top{i + 1}_' + k, value=v[i],
        #                                                   iteration=self.step,
        #                                                   group_name='Eval')
        #         else:
        #             self.train_platform.report_scalar(name=k, value=v, iteration=self.step,
        #                                               group_name='Eval')

        # elif self.dataset in ['humanact12', 'uestc']:
        #     eval_args = SimpleNamespace(num_seeds=self.cfg.eval_rep_times, num_samples=self.cfg.eval_num_samples,
        #                                 batch_size=self.cfg.eval_batch_size, device=self.device, guidance_param = 1,
        #                                 dataset=self.dataset, unconstrained=self.cfg.unconstrained,
        #                                 model_path=os.path.join(self.save_dir, self.ckpt_file_name()))
        #     eval_dict = eval_humanact12_uestc.evaluate(eval_args, model=self.model, diffusion=self.diffusion, data=self.train_dataloader.dataset)
        #     print(f'Evaluation results on {self.dataset}: {sorted(eval_dict["feats"].items())}')
        #     for k, v in eval_dict["feats"].items():
        #         if 'unconstrained' not in k:
        #             self.train_platform.report_scalar(name=k, value=np.array(v).astype(float).mean(), iteration=self.step, group_name='Eval')
        #         else:
        #             self.train_platform.report_scalar(name=k, value=np.array(v).astype(float).mean(), iteration=self.step, group_name='Eval Unconstrained')

        # end_eval = time.time()
        # print(f'Evaluation time: {round(end_eval-start_eval)/60}min')

    def compute_reconstruction_loss(self, samples, gt, lengths):
        """
        Compute MSE loss and velocity loss between samples and ground truth.

        Args:
            samples: numpy array of shape (B, T, D) or (N_samples, T, D)
            gt: numpy array of shape (B, T, D)
            lengths: tensor or numpy array of shape (B,)

        Returns:
            mse_losses: numpy array of shape (B,) containing per-sample MSE loss
            vel_losses: numpy array of shape (B,) containing per-sample velocity MSE loss
        """
        # If samples has multiple denoising steps, take the last one
        if samples.ndim == 3 and gt.ndim == 3:
            # samples: (N_samples, T, D), gt: (B, T, D)
            # This happens when we have multiple denoising steps but batch_size=1
            samples = samples[-1:, :, :]  # Take the final sample, shape: (1, T, D)
        elif samples.ndim == 4:
            # samples: (N_samples, B, T, D), gt: (B, T, D)
            samples = samples[-1]  # Take the final sample, shape: (B, T, D)

        if isinstance(lengths, torch.Tensor):
            lengths = lengths.cpu().numpy()

        # Ensure samples is at least 2D
        if samples.ndim == 2:
            samples = samples[np.newaxis, :, :]  # (1, T, D)

        batch_size = gt.shape[0]
        mse_losses = []
        vel_losses = []

        for i in range(batch_size):
            length = int(lengths[i])
            sample = samples[i, :length]  # (T, D)
            gt_seq = gt[i, :length]  # (T, D)

            # MSE loss
            mse = np.mean((sample - gt_seq) ** 2)
            mse_losses.append(mse)

            # Velocity MSE loss
            if length > 1:
                sample_vel = sample[1:] - sample[:-1]  # (T-1, D)
                gt_vel = gt_seq[1:] - gt_seq[:-1]  # (T-1, D)
                vel_mse = np.mean((sample_vel - gt_vel) ** 2)
            else:
                vel_mse = 0.0
            vel_losses.append(vel_mse)

        return np.array(mse_losses), np.array(vel_losses)

    def sample(self):
        sample_fn = self.diffusion.p_sample_loop
        # x, model_kwargs = next(iter(self.viz_data))
        try:
            x, model_kwargs = next(self.iter_viz_data)
        except:
            self.iter_viz_data = iter(self.viz_dataloader)
            x, model_kwargs = next(self.iter_viz_data)
        # set model to eval mode
        self.sample_model.eval()

        dump_steps = None
        if self.cfg.train.viz_cfg.denoising_steps and len(self.cfg.train.viz_cfg.denoising_steps) > 0:
            dump_steps = self.cfg.train.viz_cfg.denoising_steps

        samples = sample_fn(
            self.sample_model,
            # (args.batch_size, model.njoints, model.nfeats, n_frames),  # BUG FIX - this one caused a mismatch between training and inference
            (x.shape[0], self.model_without_ddp.njoints, self.model_without_ddp.nfeats, self.sample_length),  # BUG FIX
            clip_denoised=False,
            model_kwargs=model_kwargs,
            device=self.device,
            skip_timesteps=0,  # 0 is the default value - i.e. don't skip any step
            init_image=None,
            progress=True,
            dump_steps=dump_steps,
            noise=None,
            const_noise=False,
        )

        if isinstance(samples, list):
            samples = torch.concat(samples, dim=0)

        gt = x
        samples = rearrange(samples, 'b j f t -> b t (j f)')
        gt = rearrange(gt, 'b j f t -> b t (j f)')

        # Compute losses in NORMALIZED space (before inv_transform)
        samples_normalized = samples.detach().cpu().numpy()
        gt_normalized = gt.detach().cpu().numpy()
        mse_losses, vel_losses = self.compute_reconstruction_loss(samples_normalized, gt_normalized, model_kwargs['y']['lengths'])

        # Then apply inv_transform for visualization
        samples = self.train_dataloader.dataset.inv_transform(samples_normalized)
        gt = self.train_dataloader.dataset.inv_transform(gt_normalized)

        self.model.train()
        action_label = model_kwargs['y'].get('action', None)
        if action_label is not None:
            action_label = [gesture_list[i] for i in action_label.cpu().numpy().reshape(-1)]
        if action_label is None:
            action_label = model_kwargs['y'].get('text', None)
        return samples, gt, model_kwargs['y']['lengths'], action_label, mse_losses, vel_losses

    def run_step(self, batch, cond):
        self.forward_backward(batch, cond)

        param_norm, grad_norm = 0.0, 0.0
        for p in self.model.parameters():
            with torch.no_grad():
                param_norm += torch.norm(p).item() ** 2
                if p.grad is not None:
                    grad_norm += torch.norm(p.grad).item() ** 2

        param_norm = np.sqrt(param_norm)
        grad_norm = np.sqrt(grad_norm)
        mylogger.record_metric("param_norm", param_norm)
        mylogger.record_metric("grad_norm", grad_norm)
        self.opt.step()

        clip_grad_norm_(self.model.parameters(), 1.0)

        # self._anneal_lr()

    def forward_backward(self, batch, cond):
        self.opt.zero_grad()
        t, weights = self.schedule_sampler.sample(batch.shape[0], self.device)
        with torch.autocast('cuda', dtype=torch.bfloat16):
            losses = self.diffusion.training_losses(
                self.model,
                batch,
                t,
                model_kwargs=cond,
                dataset=self.train_dataloader.dataset
            )
        if isinstance(self.schedule_sampler, LossAwareSampler):
            self.schedule_sampler.update_with_local_losses(
                t, losses['loss'].detach()
            )

        loss = (losses['loss'] * weights).mean()

        log_dict(
            self.diffusion, t, {k: v * weights for k, v in losses.items()}, key_prefix="train_"
        )
        loss.backward()

    def validation_forward(self, batch, cond, split):
        self.model.eval()
        with torch.no_grad():
            t, weights = self.schedule_sampler.sample(batch.shape[0], self.device)
            with torch.autocast('cuda', dtype=torch.bfloat16):
                metrics = self.diffusion.evaluate_metrics(
                    self.model,
                    batch,
                    t,
                    model_kwargs=cond,
                    inv_transform=self.train_dataloader.dataset.inv_transform
                )

            log_dict(
                self.diffusion, t, {k: v * weights for k, v in metrics.items()}, key_prefix=f"val_on_{split}_"
            )
        self.model.train()


    def _anneal_lr(self):
        if not self.lr_anneal_steps:
            return
        frac_done = (self.step) / self.lr_anneal_steps
        lr = self.lr * (1 - frac_done)
        for param_group in self.opt.param_groups:
            param_group["lr"] = lr



    def ckpt_file_name(self):
        return f"model{(self.step):09d}.pt"


    def save(self):
        def save_checkpoint(model_state_dict):
            # Do not save text encoder weights
            checkpoint = dict()
            checkpoint['state_dict'] = dict()

            trainable_param_names = {
                name for name, param in self.model_without_ddp.named_parameters() if param.requires_grad
            }

            for key, value in model_state_dict.items():
                if not key.startswith("_text_model."):
                    checkpoint['state_dict'][key] = value
                elif key in trainable_param_names:
                    checkpoint['state_dict'][key] = value

            mylogger.info(f"saving model...")
            filename = self.ckpt_file_name()
            with open(os.path.join(self.save_dir, filename), "wb") as f:
                torch.save(checkpoint, f)


        dist_utils.barrier()
        if dist_utils.is_main_process():
            save_checkpoint(self.model_without_ddp.state_dict())

            with open(os.path.join(self.save_dir, f"opt{self.step:09d}.pt"), "wb") as f:
                torch.save(self.opt.state_dict(), f)


def parse_resume_step_from_filename(filename):
    """
    Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the
    checkpoint's number of steps.
    """
    split = filename.split("model")
    if len(split) < 2:
        return 0
    split1 = split[-1].split(".")[0]
    try:
        return int(split1)
    except ValueError:
        return 0



def log_dict(diffusion : GaussianDiffusion, ts : torch.Tensor, losses : Dict[str, torch.Tensor], key_prefix: str = ""):
    for key, values in losses.items():
        mylogger.record_metric(key_prefix + key, values.mean().item())

        for sub_t, sub_loss in zip(ts.detach().cpu().numpy(), values.detach().cpu().numpy()):
            quartile = int(4 * sub_t / diffusion.num_timesteps)
            mylogger.record_metric(f"{key_prefix}{key}_q{quartile}", sub_loss)