utils.py

import os
import time
import random
import numpy as np
import torch


class BlackHole(object):
    def __setattr__(self, name, value):
        pass

    def __call__(self, *args, **kwargs):
        return self

    def __getattr__(self, name):
        return self


def get_new_log_dir(root="./logs", prefix="", postfix=""):
    log_dir = os.path.join(
        root, prefix + time.strftime("%Y_%m_%d__%H_%M_%S", time.localtime()) + postfix
    )
    os.makedirs(log_dir)
    return log_dir


class CheckpointManager(object):
    def __init__(self, save_dir, logger=BlackHole()):
        super().__init__()
        os.makedirs(save_dir, exist_ok=True)
        self.save_dir = save_dir
        self.ckpts = []
        self.logger = logger

        for f in os.listdir(self.save_dir):
            if f[:4] != "ckpt":
                continue
            _, score, it = f.split("_")
            it = it.split(".")[0]
            self.ckpts.append(
                {
                    "score": float(score),
                    "file": f,
                    "iteration": int(it),
                }
            )

    def get_worst_ckpt_idx(self):
        idx = -1
        worst = float("-inf")
        for i, ckpt in enumerate(self.ckpts):
            if ckpt["score"] >= worst:
                idx = i
                worst = ckpt["score"]
        return idx if idx >= 0 else None

    def get_best_ckpt_idx(self):
        idx = -1
        best = float("inf")
        for i, ckpt in enumerate(self.ckpts):
            if ckpt["score"] <= best:
                idx = i
                best = ckpt["score"]
        return idx if idx >= 0 else None

    def get_latest_ckpt_idx(self):
        idx = -1
        latest_it = -1
        for i, ckpt in enumerate(self.ckpts):
            if ckpt["iteration"] > latest_it:
                idx = i
                latest_it = ckpt["iteration"]
        return idx if idx >= 0 else None

    def save(self, model, args, score, others=None, step=None):
        if step is None:
            fname = "ckpt_%.6f_.pt" % float(score)
        else:
            fname = "ckpt_%.6f_%d.pt" % (float(score), int(step))
        path = os.path.join(self.save_dir, fname)

        torch.save(
            {
                "args": args,
                "state_dict": model.state_dict(),
                "others": others,
            },
            path,
        )

        self.ckpts.append({"score": score, "file": fname})

        return True

    def load_best(self):
        idx = self.get_best_ckpt_idx()
        if idx is None:
            raise IOError("No checkpoints found.")
        ckpt = torch.load(os.path.join(self.save_dir, self.ckpts[idx]["file"]))
        return ckpt

    def load_latest(self):
        idx = self.get_latest_ckpt_idx()
        if idx is None:
            raise IOError("No checkpoints found.")
        ckpt = torch.load(os.path.join(self.save_dir, self.ckpts[idx]["file"]))
        return ckpt

    def load_selected(self, file):
        ckpt = torch.load(os.path.join(self.save_dir, file))
        return ckpt


def seed_all(seed=42):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.enabled = True


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


# assumes  : batch, particles, feats
def normalize_tensor(tensor, mean, std, sigma=1):
    for i in range(len(mean)):
        tensor[..., i] = (tensor[..., i] - mean[i]) / (std[i] / sigma)
    return tensor


def inverse_normalize_tensor(tensor, mean, std, sigma=1):
    for i in range(len(mean)):
        tensor[..., i] = (tensor[..., i] * (std[i] / sigma)) + mean[i]
    return tensor