main.py

import os
import sys
import fire
import random
from retry.api import retry_call
from tqdm import tqdm
from datetime import datetime
from functools import wraps
from model import Trainer, NanException
from diff_augment_test import DiffAugmentTest

import torch
import torch.multiprocessing as mp
import torch.distributed as dist

import numpy as np


def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def cast_list(el):
    return el if isinstance(el, list) else [el]

def timestamped_filename(prefix='generated-'):
    now = datetime.now()
    timestamp = now.strftime('%m-%d-%Y_%H-%M-%S')
    return f'{prefix}{timestamp}'

def set_seed(seed):
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)

def run_training(rank, world_size, model_args, data, load_from, new, num_train_steps, name, seed, use_aim, aim_repo, aim_run_hash):
    is_main = rank == 0
    is_ddp = world_size > 1

    if is_ddp:
        set_seed(seed)
        os.environ['MASTER_ADDR'] = 'localhost'
        os.environ['MASTER_PORT'] = '12355'
        dist.init_process_group('nccl', rank=rank, world_size=world_size)

        print(f'{rank + 1}/{world_size} process initialized.')

    model_args.update(
        is_ddp = is_ddp,
        rank = rank,
        world_size = world_size
    )

    model = Trainer(**model_args, hparams=model_args, use_aim=use_aim, aim_repo=aim_repo, aim_run_hash=aim_run_hash)

    if not new:
        model.load(load_from)
    else:
        model.clear()

    model.set_data_src(data)

    progress_bar = tqdm(initial = model.steps, total=num_train_steps, mininterval=10., desc=f'{name}<{data}>')
    while model.steps < num_train_steps:
        retry_call(model.train, tries=3, exceptions=NanException)
        progress_bar.n = model.steps
        progress_bar.refresh()
        if is_main and model.steps % 50 == 0:
            model.print_log()

    model.save(model.checkpoint_num)

    if is_ddp:
        dist.destroy_process_group()

def train_from_folder(
    data = './data',
    results_dir = './results',
    models_dir = './models',
    name = 'default',
    new = False,
    load_from = -1,
    image_size = 256,
    render_size = 256, # 1024
    render = False,
    renderer = 'default', # 'pulsar'
    nodisplace = False,
    num_points = 10**5,
    gamma = 1e-3,
    radius = None,
    smoothing = 0,
    transparent = False,
    greyscale = False,
    mesh_obj_path = None,
    smplx_model_path = None,
    labelpath = None,
    optimizer = 'adam',
    learning_rate = 2e-4,
    latent_dim = 256,
    fmap_max = 512,
    ttur_mult = 1,
    gp_weight = 10,
    batch_size = 8,
    gradient_accumulate_every = 4,
    num_train_steps = 150000,
    save_every = 1000,
    evaluate_every = 1000,
    generate = False,
    generate_types = ['default', 'ema'],
    generate_interpolation = False,
    aug_test = False,
    aug_prob = None,
    aug_types = ['cutout', 'translation'],
    dataset_aug_prob = 0.,
    attn_res_layers = [32],
    freq_chan_attn = False,
    disc_output_size = 1,
    dual_contrast_loss = False,
    antialias = False,
    interpolation_num_steps = 100,
    save_frames = False,
    num_image_tiles = None,
    num_workers = None,
    multi_gpus = False,
    calculate_fid_every = None,
    calculate_fid_num_images = 12800,
    clear_fid_cache = False,
    seed = 42,
    amp = False,
    show_progress = False,
    use_aim = False,
    aim_repo = None,
    aim_run_hash = None,
    load_strict = True
):
    resolution = render_size if render else image_size
    num_image_tiles = default(num_image_tiles, 4 if resolution > 512 else 8)

    model_args = dict(
        name = name,
        results_dir = results_dir,
        models_dir = models_dir,
        image_size = image_size,
        render_size = render_size,
        render = render,
        renderer = renderer,
        nodisplace = nodisplace,
        num_points = num_points,
        gamma = gamma,
        radius = radius,
        smoothing = smoothing,
        transparent = transparent,
        greyscale = greyscale,
        mesh_obj_path = mesh_obj_path,
        smplx_model_path = smplx_model_path,
        labelpath = labelpath,
        optimizer = optimizer,
        lr = learning_rate,
        latent_dim = latent_dim,
        fmap_max = fmap_max,
        ttur_mult = ttur_mult,
        gp_weight = gp_weight,
        batch_size = batch_size,
        gradient_accumulate_every = gradient_accumulate_every,
        save_every = save_every,
        evaluate_every = evaluate_every,
        aug_prob = aug_prob,
        aug_types = cast_list(aug_types),
        dataset_aug_prob = dataset_aug_prob,
        attn_res_layers = cast_list(attn_res_layers),
        freq_chan_attn = freq_chan_attn,
        disc_output_size = disc_output_size,
        dual_contrast_loss = dual_contrast_loss,
        antialias = antialias,
        num_image_tiles = num_image_tiles,
        num_workers = num_workers,
        calculate_fid_every = calculate_fid_every,
        calculate_fid_num_images = calculate_fid_num_images,
        clear_fid_cache = clear_fid_cache,
        amp = amp,
        load_strict = load_strict
    )

    if generate:
        model = Trainer(**model_args, use_aim=use_aim)
        model.load(load_from)
        samples_name = timestamped_filename()
        checkpoint = model.checkpoint_num
        dir_result = model.generate(samples_name, num_image_tiles, checkpoint, generate_types)
        print(f'sample images generated at {dir_result}')
        return

    if generate_interpolation:
        model = Trainer(**model_args, use_aim=use_aim)
        model.load(load_from)
        samples_name = timestamped_filename()
        model.generate_interpolation(samples_name, num_image_tiles, num_steps=interpolation_num_steps, save_frames=save_frames)
        print(f'interpolation generated at {results_dir}/{name}/{samples_name}')
        return

    if show_progress:
        model = Trainer(**model_args, use_aim=use_aim)
        model.show_progress(num_images=num_image_tiles, types=generate_types)
        return

    if aug_test:
        DiffAugmentTest(data=data, image_size=resolution, batch_size=batch_size, types=aug_types, nrow=num_image_tiles)
        return

    world_size = torch.cuda.device_count()

    if world_size == 1 or not multi_gpus:
        run_training(0, 1, model_args, data, load_from, new, num_train_steps, name, seed, use_aim, aim_repo, aim_run_hash)
        return

    mp.spawn(run_training,
        args = (world_size, model_args, data, load_from, new, num_train_steps, name, seed, use_aim, aim_repo, aim_run_hash),
        nprocs = world_size,
        join = True
    )

def main():
     fire.Fire(train_from_folder)

if __name__ == '__main__':
    sys.exit(main())