reconstruct_phema.py

#!/usr/bin/env python
# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is licensed under a Creative Commons
# Attribution-NonCommercial-ShareAlike 4.0 International License.
# You should have received a copy of the license along with this
# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/

"""Perform post-hoc EMA reconstruction.

This is a direct copy of `edm2/reconstruct_phema.py`, wired to use
`training.phema` from this repo so that PHEMA reconstruction keeps working.
"""

import os
import re
import copy
import warnings
import click
import tqdm
import pickle
import numpy as np
import torch
import dnnlib
import training.phema
from torch_utils import persistence

warnings.filterwarnings('ignore', 'You are using `torch.load` with `weights_only=False`')

#----------------------------------------------------------------------------
# Persistence import hook to make training snapshots robust to relative imports.
# Mirrors the logic used in generate.py so that checkpoints that include
# training.loss_cd and torch_utils.* can be loaded without
# `attempted relative import with no known parent package` errors.

@persistence.import_hook
def _fix_relative_imports(meta: dnnlib.EasyDict) -> dnnlib.EasyDict:
    src = meta.module_src

    # training.loss_cd: relative import -> absolute
    src = src.replace(
        "from .consistency_ops import (",
        "from training.consistency_ops import (",
    )

    # torch_utils.*: relative imports -> absolute
    src = src.replace(
        "from . import distributed as dist",
        "from torch_utils import distributed as dist",
    )
    src = src.replace(
        "from . import training_stats",
        "from torch_utils import training_stats",
    )
    src = src.replace(
        "from . import misc",
        "from torch_utils import misc",
    )

    # dnnlib.__init__: relative import -> absolute
    src = src.replace(
        "from .util import EasyDict, make_cache_dir_path",
        "from dnnlib.util import EasyDict, make_cache_dir_path",
    )

    meta.module_src = src
    return meta

#----------------------------------------------------------------------------
# Construct the full path of a network pickle.

def pkl_path(dir, prefix, nimg, std):
    name = prefix + f'-{nimg//1000:07d}-{std:.3f}.pkl'
    if dir is None:
        return None
    if dnnlib.util.is_url(dir):
        return f'{dir}/{name}'
    return os.path.join(dir, name)

#----------------------------------------------------------------------------
# Recover exact nimg from kimg and batch_size.
#
# cur_nimg is always a multiple of batch_size.  The filename stores
# kimg = cur_nimg // 1000.  Given batch_size we find the unique multiple
# of batch_size in [kimg*1000, (kimg+1)*1000).

def kimg_to_nimg(kimg, batch_size=None):
    if batch_size is None:
        return kimg * 1000
    lo = kimg * 1000
    hi = lo + 1000
    # When batch_size > 1000, there is at most one multiple of batch_size in
    # [lo, hi), and it is the exact cur_nimg that the training loop had when
    # writing this file — giving identical math to EDM2's original intent.
    nimg = ((lo + batch_size - 1) // batch_size) * batch_size
    if nimg < hi:
        return nimg
    raise click.ClickException(
        f'Cannot recover exact nimg for kimg={kimg} with batch_size={batch_size}: '
        f'no multiple of {batch_size} in [{lo}, {hi}). '
        f'This file may not have been produced with this batch size.'
    )

#----------------------------------------------------------------------------
# List input pickles for post-hoc EMA reconstruction.
# Returns a list of dnnlib.EasyDict(path, nimg, std).

def list_input_pickles(
    in_dir,             # Directory containing the input pickles.
    in_prefix   = None, # Input filename prefix. None = anything goes.
    in_std      = None, # Relative standard deviations of the input pickles. None = anything goes.
    batch_size  = None, # Training batch size for exact nimg recovery (optional).
):
    if not os.path.isdir(in_dir):
        raise click.ClickException('Input directory does not exist')
    in_std = set(in_std) if in_std is not None else None

    pkls = []
    with os.scandir(in_dir) as it:
        for e in it:
            m = re.fullmatch(r'(.*)-(\d+)-(\d+\.\d+)\.pkl', e.name)
            if not m or not e.is_file():
                continue
            prefix = m.group(1)
            kimg = int(m.group(2))
            std = float(m.group(3))
            if in_prefix is not None and prefix != in_prefix:
                continue
            if in_std is not None and std not in in_std:
                continue

            # Try to read exact nimg from the pickle (written by training loop).
            nimg = None
            try:
                with open(e.path, 'rb') as f:
                    data = pickle.load(f)
                if 'nimg' in data:
                    nimg = int(data['nimg'])
                del data
            except Exception:
                pass

            if nimg is None:
                nimg = kimg_to_nimg(kimg, batch_size=batch_size)

            pkls.append(dnnlib.EasyDict(path=e.path, nimg=nimg, std=std))
    pkls = sorted(pkls, key=lambda pkl: (pkl.nimg, pkl.std))
    return pkls

#----------------------------------------------------------------------------
# Perform post-hoc EMA reconstruction.
# Returns an iterable that yields dnnlib.EasyDict(out, step_idx, num_steps),
# where 'out' is a list of dnnlib.EasyDict(net, nimg, std, pkl_data, pkl_path)

def reconstruct_phema(
    in_pkls,                    # List of input pickles, expressed as dnnlib.EasyDict(path, nimg, std).
    out_std,                    # List of relative standard deviations to reconstruct.
    out_nimg        = None,     # Training time of the snapshot to reconstruct. None = highest input time.
    out_dir         = None,     # Where to save the reconstructed network pickles. None = do not save.
    out_prefix      = 'phema',  # Output filename prefix.
    skip_existing   = False,    # Skip output files that already exist?
    max_batch_size  = 8,        # Maximum simultaneous reconstructions
    verbose         = True,     # Enable status prints?
):
    # Validate input pickles.
    if out_nimg is None:
        out_nimg = max((pkl.nimg for pkl in in_pkls), default=0)
    elif not any(out_nimg == pkl.nimg for pkl in in_pkls):
        raise click.ClickException('Reconstruction time must match one of the input pickles')
    in_pkls = [pkl for pkl in in_pkls if 0 < pkl.nimg <= out_nimg]
    if len(in_pkls) == 0:
        raise click.ClickException('No valid input pickles found')
    in_nimg = [pkl.nimg for pkl in in_pkls]
    in_std = [pkl.std for pkl in in_pkls]
    if verbose:
        print(f'Loading {len(in_pkls)} input pickles...')
        for pkl in in_pkls:
            print('    ' + pkl.path)

    # Determine output pickles.
    out_std = [out_std] if isinstance(out_std, float) else sorted(set(out_std))
    if skip_existing and out_dir is not None:
        out_std = [std for std in out_std if not os.path.isfile(pkl_path(out_dir, out_prefix, out_nimg, std))]
    num_batches = (len(out_std) - 1) // max_batch_size + 1
    out_std_batches = np.array_split(out_std, num_batches)
    if verbose:
        print(f'Reconstructing {len(out_std)} output pickles in {num_batches} batches...')
        for i, batch in enumerate(out_std_batches):
            for std in batch:
                print(f'    batch {i}: ', end='')
                print(pkl_path(out_dir, out_prefix, out_nimg, std) if out_dir is not None else pkl_path('', '<yield>', out_nimg, std))

    # Return an iterable over the reconstruction steps.
    class ReconstructionIterable:
        def __len__(self):
            return num_batches * len(in_pkls)

        def __iter__(self):
            # Loop over batches.
            r = dnnlib.EasyDict(step_idx=0, num_steps=len(self))
            for out_std_batch in out_std_batches:
                coefs = training.phema.solve_posthoc_coefficients(in_nimg, in_std, out_nimg, out_std_batch)
                out = [dnnlib.EasyDict(net=None, nimg=out_nimg, std=std) for std in out_std_batch]
                r.out = []

                # Loop over input pickles.
                for i in range(len(in_pkls)):
                    with dnnlib.util.open_url(in_pkls[i].path, verbose=False) as f:
                        in_pkl_data = pickle.load(f)
                        in_net = in_pkl_data['ema'].to(torch.float32)

                    # Accumulate weights for each output pickle.
                    for j in range(len(out)):
                        if out[j].net is None:
                            out[j].pkl_data = copy.deepcopy(in_pkl_data)
                            out[j].net = out[j].pkl_data['ema']
                            for pj in out[j].net.parameters():
                                pj.zero_()
                        for pi, pj in zip(in_net.parameters(), out[j].net.parameters()):
                            pj += pi * coefs[i, j]
                        for pi, pj in zip(in_net.buffers(), out[j].net.buffers()):
                            pj.copy_(pi)

                    # Finalize outputs.
                    if i == len(in_pkls) - 1:
                        for j in range(len(out)):
                            out[j].net.to(torch.float16)
                            out[j].pkl_path = pkl_path(out_dir, out_prefix, out_nimg, out[j].std)
                            if out[j].pkl_path is not None:
                                os.makedirs(out_dir, exist_ok=True)
                                with open(out[j].pkl_path, 'wb') as f:
                                    pickle.dump(out[j].pkl_data, f)
                        r.out = out

                    # Yield results.
                    del in_pkl_data, in_net # conserve memory
                    yield r
                    r.step_idx += 1

    return ReconstructionIterable()

#----------------------------------------------------------------------------
# Parse a comma separated list of relative standard deviations.
# The special token '...' interpreted as an evenly spaced interval.
# Example: '0.01,0.02,...,0.05' returns [0.01, 0.02, 0.03, 0.04, 0.05]

def parse_std_list(s):
    if isinstance(s, list):
        return s

    # Parse raw values.
    raw = [None if v == '...' else float(v) for v in s.split(',')]

    # Fill in '...' tokens.
    out = []
    for i, v in enumerate(raw):
        if v is not None:
            out.append(v)
            continue
        if i - 2 < 0 or raw[i - 2] is None or raw[i - 1] is None:
            raise click.ClickException("'...' must be preceded by at least two floats")
        if i + 1 >= len(raw) or raw[i + 1] is None:
            raise click.ClickException("'...' must be followed by at least one float")
        if raw[i - 2] == raw[i - 1]:
            raise click.ClickException("The floats preceding '...' must not be equal")
        approx_num = (raw[i + 1] - raw[i - 1]) / (raw[i - 1] - raw[i - 2]) - 1
        num = round(approx_num)
        if num <= 0:
            raise click.ClickException("'...' must correspond to a non-empty interval")
        if abs(num - approx_num) > 1e-4:
            raise click.ClickException("'...' must correspond to an evenly spaced interval")
        for j in range(num):
            out.append(raw[i - 1] + (raw[i - 1] - raw[i - 2]) * (j + 1))

    # Validate.
    out = sorted(set(out))
    if not all(0.000 < v < 0.289 for v in out):
        raise click.ClickException('Relative standard deviation must be positive and less than 0.289')
    return out

#----------------------------------------------------------------------------
# Command line interface.

@click.command()
@click.option('--indir', 'in_dir',          help='Directory containing the input pickles', metavar='DIR',           type=str, required=True)
@click.option('--inprefix', 'in_prefix',    help='Filter inputs based on filename prefix', metavar='STR',           type=str, default=None)
@click.option('--instd', 'in_std',          help='Filter inputs based on standard deviations', metavar='LIST',      type=parse_std_list, default=None)

@click.option('--outdir', 'out_dir',        help='Where to save the reconstructed network pickles', metavar='DIR',  type=str, required=True)
@click.option('--outprefix', 'out_prefix',  help='Output filename prefix', metavar='STR',                           type=str, default='phema', show_default=True)
@click.option('--outstd', 'out_std',        help='List of desired relative standard deviations', metavar='LIST',    type=parse_std_list, required=True)
@click.option('--outkimg', 'out_kimg',      help='Training time of the snapshot to reconstruct', metavar='KIMG',    type=click.IntRange(min=1), default=None)

@click.option('--skip', 'skip_existing',    help='Skip output files that already exist',                            is_flag=True)
@click.option('--batch', 'max_batch_size',  help='Maximum simultaneous reconstructions', metavar='INT',             type=click.IntRange(min=1), default=8, show_default=True)
@click.option('--training_batch',           help='Training batch size for exact nimg recovery from kimg filenames',  metavar='INT', type=click.IntRange(min=1), default=None)

def cmdline(in_dir, in_prefix, in_std, out_kimg, training_batch, **opts):
    """Perform post-hoc EMA reconstruction.

    Examples:

    \b
    # Download raw snapshots for the pre-trained edm2-img512-xs model
    rclone copy --progress --http-url https://nvlabs-fi-cdn.nvidia.com/edm2 \\
        :http:raw-snapshots/edm2-img512-xs/ raw-snapshots/edm2-img512-xs/

    \b
    # Reconstruct a new EMA profile with std=0.150
    python reconstruct_phema.py --indir=raw-snapshots/edm2-img512-xs \\
        --outdir=out --outstd=0.150

    \b
    # Reconstruct a set of 31 EMA profiles, streaming over the input data 4 times
    python reconstruct_phema.py --indir=raw-snapshots/edm2-img512-xs \\
        --outdir=out --outstd=0.010,0.015,...,0.250 --batch=8

    \b
    # Perform reconstruction for the latest snapshot of a given training run
    python reconstruct_phema.py --indir=training-runs/00000-edm2-img512-xs \\
        --outdir=out --outstd=0.150
    """
    if os.environ.get('WORLD_SIZE', '1') != '1':
        raise click.ClickException('Distributed execution is not supported')
    out_nimg = kimg_to_nimg(out_kimg, batch_size=training_batch) if out_kimg is not None else None
    in_pkls = list_input_pickles(in_dir=in_dir, in_prefix=in_prefix, in_std=in_std, batch_size=training_batch)
    rec_iter = reconstruct_phema(in_pkls=in_pkls, out_nimg=out_nimg, **opts)
    for _r in tqdm.tqdm(rec_iter, unit='step'):
        pass

#----------------------------------------------------------------------------

if __name__ == "__main__":
    cmdline()